Update tasks/text.py

tasks/text.py

@@ -61,6 +61,10 @@ async def evaluate_text(request: TextEvaluationRequest):
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE CODE HERE
 
+    # Set the device to MPS (if available)
+    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+    print(f"Using device: {device}")
+
     #Load the embedding model
     #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
@@ -91,7 +95,7 @@ async def evaluate_text(request: TextEvaluationRequest):
 
     # Make predictions
     with torch.no_grad():
-        outputs = model_nn(text_embeddings)
+        outputs = model_nn(text_embeddings.to(device))
         _, predicted = torch.max(outputs, 1)  # Get the class with the highest score
     
     # Decode the predictions back to original labels using label_encoder