Update tasks/text.py

tasks/text.py

@@ -62,6 +62,43 @@ async def evaluate_text(request: TextEvaluationRequest):
     tracker.start_task("inference")
 
     #--------------------------------------------------------------------------------------------
+    # Load a pre-trained Sentence-BERT model
+    model = SentenceTransformer('sentence-transformers/all-MPNET-base-v2', device='cpu')
+    # Generate sentence embeddings
+    sentence_embeddings = model.encode(test_dataset["quote"])
+
+    #load the models
+    with open("xgb_bin.pkl","rb") as f:
+        xgb_bin = pickle.load(f)
+
+    with open("xgb_multi.pkl","rb") as f:
+        xgb_multi = pickle.load(f)
+
+
+    X_train = test_dataset["quote"]
+    
+    y_train = test_dataset["label"].copy()
+
+    #binary
+    y_train_binary = y_train.copy()
+    y_train_binary[y_train_binary != 0] = 1
+
+
+    #multi class
+    X_train_multi = X_train[y_train != 0]
+
+    y_train_multi = y_train[y_train != 0]
+
+    #predictions
+    y_pred_bin = xgb_bin.predict(X_train)
+
+    y_pred_multi = xgb_multi.predict(X_train_multi) + 1
+
+    y_pred_bin[y_pred_bin==1] = y_pred_multi
+    
+    
+    
+
 
 
     #predictions = xgb.predict(embeddings)
@@ -80,7 +117,7 @@ async def evaluate_text(request: TextEvaluationRequest):
     emissions_data = tracker.stop_task()
     
     # Calculate accuracy
-    accuracy = accuracy_score(true_labels, np.array([1]*len(true_labels)))
+    accuracy = accuracy_score(true_labels, y_true))
     
     # Prepare results dictionary
     results = {