Update tasks/text.py

tasks/text.py

@@ -7,6 +7,12 @@ import random
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
+## add-on imports
+from sentence_transformers import SentenceTransformer
+from sklearn.preprocessing import MinMaxScaler
+import numpy as np
+import pickle
+
 router = APIRouter()
 
 DESCRIPTION = "Random test"
@@ -56,9 +62,20 @@ async def evaluate_text(request: TextEvaluationRequest):
     # 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.
     #--------------------------------------------------------------------------------------------   
     
-    # Make random predictions (placeholder for actual model inference)
-    true_labels = test_dataset["label"]
-    predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
+    ## Models loading
+    query_prompt_name = "s2s_query"
+    model = SentenceTransformer("dunzhang/stella_en_400M_v5",trust_remote_code=True).cuda()
+    
+    with open('logistic_regression_model.pkl', 'rb') as file:
+        disp = pickle.load(file)
+    
+    ## Data prep
+    embeddings = model.encode(list(test_dataset['quote']), prompt_name=query_prompt_name)
+    scaler = MinMaxScaler()
+    X_scaled = scaler.fit_transform(embeddings)
+    
+    ## Predictions
+    predictions= disp.predict(X_scaled)
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE