Update tasks/text.py

tasks/text.py

@@ -11,7 +11,7 @@ from .utils.emissions import tracker, clean_emissions_data, get_space_info
 from sentence_transformers import SentenceTransformer
 from sklearn.preprocessing import MinMaxScaler
 import numpy as np
-import pickle
+import skops.io as sio
 
 router = APIRouter()
 
@@ -66,8 +66,8 @@ async def evaluate_text(request: TextEvaluationRequest):
     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)
+    trusted_types = ['sklearn.feature_selection._univariate_selection.f_classif']
+    disp = sio.load('logistic_regression_model.skops',trusted=trusted_types)
     
     ## Data prep
     embeddings = model.encode(list(test_dataset['quote']), prompt_name=query_prompt_name)
@@ -81,11 +81,11 @@ async def evaluate_text(request: TextEvaluationRequest):
     # YOUR MODEL INFERENCE STOPS HERE
     #--------------------------------------------------------------------------------------------   
 
-    
     # Stop tracking emissions
     emissions_data = tracker.stop_task()
     
     # Calculate accuracy
+    true_labels = test_dataset["label"]
     accuracy = accuracy_score(true_labels, predictions)
     
     # Prepare results dictionary