Update tasks/text.py

tasks/text.py

@@ -2,7 +2,9 @@ from fastapi import APIRouter
 from datetime import datetime
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
-import random
+from sklearn.linear_model import LogisticRegression
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.model_selection import train_test_split
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
@@ -47,6 +49,12 @@ async def evaluate_text(request: TextEvaluationRequest):
     train_test = dataset["train"].train_test_split(test_size=request.test_size, seed=request.test_seed)
     test_dataset = train_test["test"]
     
+    tfidf_vect = TfidfVectorizer(stop_words = 'english')
+    
+    tfidf_train = tfidf_vect.fit_transform(train_dataset['quote'])
+    tfidf_test = tfidf_vect.transform(test_dataset['quote'])
+
+
     # Start tracking emissions
     tracker.start()
     tracker.start_task("inference")
@@ -58,8 +66,11 @@ async def evaluate_text(request: TextEvaluationRequest):
     
     # Make random predictions (placeholder for actual model inference)
     true_labels = test_dataset["label"]
-    predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
-
+    LR = LogisticRegression(class_weight='balanced', max_iter=20, random_state=1234,
+                   solver='liblinear')
+    LR.fit(pd.DataFrame.sparse.from_spmatrix(tfidf_train), pd.DataFrame(y_train_v))
+    predictions=LR.predict(pd.DataFrame.sparse.from_spmatrix(tfidf_test))
+    
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE
     #--------------------------------------------------------------------------------------------