Update tasks/text.py

tasks/text.py

@@ -45,7 +45,6 @@ class_labels = list(class_descriptions.keys())
 class_sentences = list(class_descriptions.values())
 class_embeddings = embedding_model.encode(class_sentences, batch_size=8, convert_to_numpy=True, normalize_embeddings=True)
 
-
 @router.post(ROUTE, tags=["Text Task"], description=DESCRIPTION)
 async def evaluate_text(request: TextEvaluationRequest):
     """
@@ -71,16 +70,23 @@ async def evaluate_text(request: TextEvaluationRequest):
     # --------------------------------------------------------------------------------------------
     # Optimized cosine similarity-based classification with threshold
     # --------------------------------------------------------------------------------------------
-
-    # Batch embed all test dataset quotes
-    test_embeddings = embedding_model.encode(test_dataset["quote"], batch_size=128)
     
+    # Convert "quote" key into embeddings
+    def embed_quote(example):
+        example["quote_embedding"] = embedding_model.encode(example["quote"]).tolist()
+        return example
+    
+    test_dataset = test_dataset.map(embed_quote, batched=True)
+    
+    # Convert test embeddings to numpy array
+    test_embeddings = np.array(test_dataset["quote_embedding"])
+
     # Compute cosine similarity in a single operation
     similarity_matrix = np.dot(test_embeddings, class_embeddings.T)  # Efficient matrix multiplication
     best_indices = similarity_matrix.argmax(axis=1)  # Get index of highest similarity for each test sample
     best_similarities = similarity_matrix.max(axis=1)  # Get max similarity values
 
-    # Apply threshold (0.9) for classification
+    # Apply threshold (0.8) for classification
     predictions = [
         LABEL_MAPPING[class_labels[idx]] if sim > 0.8 else LABEL_MAPPING["0_not_relevant"]
         for idx, sim in zip(best_indices, best_similarities)