From d1f0ad8d125c818ddcb6c0e754f1c96e90232d60 Mon Sep 17 00:00:00 2001
From: tunabenson <121179065+tunabenson@users.noreply.github.com>
Date: Thu, 13 Jun 2024 15:38:37 -0700
Subject: [PATCH] Update inference.ts

textClassification often returns 2D array, firebase cannot handle this input.
"3 INVALID_ARGUMENT: Cannot convert an array value in an array value."
---
 .../functions/src/inference.ts                               | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/extensions/firestore-huggingface-inference-api/functions/src/inference.ts b/extensions/firestore-huggingface-inference-api/functions/src/inference.ts
index eec0571..6a5c010 100644
--- a/extensions/firestore-huggingface-inference-api/functions/src/inference.ts
+++ b/extensions/firestore-huggingface-inference-api/functions/src/inference.ts
@@ -65,8 +65,11 @@ export async function runInference(
       return tasks.sentenceSimilarity(snapshot, inference);
     }
 
+    //Some text classification Models return a 2-dimensional array of JSON objects in textClassification task 
+    //firestore cannot handle 2D arrays (3 INVALID_ARGUMENT: Cannot convert an array value in an array value)
     case TaskId.textClassification: {
-      return tasks.textClassification(snapshot, inference);
+      const data= await tasks.textClassification(snapshot, inference);
+      return data.flat(Infinity);
     }
 
     // Text generation and text2text are currently the same.