Update tasks/text.py

tasks/text.py

@@ -7,11 +7,6 @@ import random
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
-import os
-print('######################    ----     ##################################')
-print(os.getcwd())
-
-
 ## add-on imports
 from sentence_transformers import SentenceTransformer
 from sklearn.preprocessing import MinMaxScaler
@@ -68,9 +63,11 @@ async def evaluate_text(request: TextEvaluationRequest):
     #--------------------------------------------------------------------------------------------   
     
     ## Models loading
+    # Embedding model
     query_prompt_name = "s2s_query"
     model = SentenceTransformer("dunzhang/stella_en_400M_v5",trust_remote_code=True).cuda()
-    
+
+    # Pre-trained Logistic Regression model
     trusted_types = ['sklearn.feature_selection._univariate_selection.f_classif']
     disp = sio.load('./tasks/logistic_regression_model.skops',trusted=trusted_types)
     
@@ -80,7 +77,7 @@ async def evaluate_text(request: TextEvaluationRequest):
     X_scaled = scaler.fit_transform(embeddings)
     
     ## Predictions
-    predictions= disp.predict(X_scaled)
+    predictions = disp.predict(X_scaled)
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE