Update tasks/text.py

tasks/text.py

@@ -7,6 +7,14 @@ import random
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
+
+
+#
+from sentence_transformers import SentenceTransformer
+from xgboost import XGBClassifier
+import pickle
+
+
 router = APIRouter()
 
 DESCRIPTION = "Random Baseline"
@@ -52,13 +60,27 @@ async def evaluate_text(request: TextEvaluationRequest):
     tracker.start_task("inference")
 
     #--------------------------------------------------------------------------------------------
+
+    #load 
+    # Step 1: Use Sentence-BERT to convert text to embeddings
+    model = SentenceTransformer("dunzhang/stella_en_400M_v5",trust_remote_code=True)
+ 
+    # Convert each sentence into a vector representation (embedding)
+    embeddings = model.encode(test_dataset['quote'])
     # YOUR MODEL INFERENCE CODE HERE
     # 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.
     #--------------------------------------------------------------------------------------------   
+
+    #load model
+    with open("models/stella_400_xgb_500.pkl","rb") as f:
+        xgb = pickle.load(f)
+
+    #predictions = xgb.predict(embeddings)
+        
     
     # Make random predictions (placeholder for actual model inference)
     true_labels = test_dataset["label"]
-    predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
+    predictions = xgb.predict(embeddings)
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE