Update tasks/text.py

tasks/text.py

@@ -11,6 +11,7 @@ from .utils.emissions import tracker, clean_emissions_data, get_space_info
 from sentence_transformers import SentenceTransformer
 from xgboost import XGBClassifier
 import pickle
+import torch
 
 
 router = APIRouter()
@@ -61,10 +62,13 @@ async def evaluate_text(request: TextEvaluationRequest):
     # YOUR MODEL INFERENCE CODE HERE
 
     #Load the embedding model
-    model = SentenceTransformer("dunzhang/stella_en_400M_v5",trust_remote_code=True) 
+    #model = SentenceTransformer("dunzhang/stella_en_400M_v5",trust_remote_code=True) 
+    model_name = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"  # You can use other Sentence Transformers models as needed
+    sentence_model = SentenceTransformer(model_name)
 
     # Convert each sentence into a vector representation (embedding)
     embeddings = model.encode(test_dataset['quote'].tolist())
+    embeddings = embeddings.cpu()
 
     # Update the code below to replace the random baseline by your model inference within the inference pass where the energy consumption and emissions are tracked.
     #--------------------------------------------------------------------------------------------   
@@ -74,11 +78,24 @@ async def evaluate_text(request: TextEvaluationRequest):
 
 
     #load the xgboost model
-    with open("models/stella_400_xgb_500.pkl",'rb') as f:
-        xgbclassifier = pickle.load(f)
+    #with open("models/stella_400_xgb_500.pkl",'rb') as f:
+    #    xgbclassifier = pickle.load(f)
+
+    model_nn = torch.load("models/model_nn.pth")
+
+    # Set the model to evaluation mode
+    model_nn.eval()
 
     #make inference
-    predictions = xgbclassifier.predict(embeddings)
+    #predictions = xgbclassifier.predict(embeddings)
+
+    # Make predictions
+    with torch.no_grad():
+        outputs = model_nn(text_embeddings)
+        _, predicted = torch.max(outputs, 1)  # Get the class with the highest score
+    
+    # Decode the predictions back to original labels using label_encoder
+    predicted_labels = label_encoder.inverse_transform(predicted.cpu().numpy())
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE