Update tasks/text.py

tasks/text.py

@@ -2,23 +2,23 @@ from fastapi import APIRouter
 from datetime import datetime
 from datasets import load_dataset
 from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.svm import SVC
+from sklearn.linear_model import LogisticRegression
+from sklearn.model_selection import GridSearchCV
 from sklearn.metrics import accuracy_score
-from sklearn.model_selection import train_test_split
-import numpy as np
+from sklearn.pipeline import Pipeline
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 router = APIRouter()
 
-DESCRIPTION = "TF-IDF + SVM Classifier"
+DESCRIPTION = "TF-IDF + Logistic Regression"
 ROUTE = "/text"
 
 @router.post(ROUTE, tags=["Text Task"], description=DESCRIPTION)
 async def evaluate_text(request: TextEvaluationRequest):
     """
-    Evaluate text classification for climate disinformation detection using TF-IDF and SVM.
+    Evaluate text classification for climate disinformation detection using TF-IDF and Logistic Regression.
     """
     # Get space info
     username, space_url = get_space_info()
@@ -45,7 +45,6 @@ async def evaluate_text(request: TextEvaluationRequest):
     train_data = dataset["train"]
     test_data = dataset["test"]
 
-    # Extract text and labels
     train_texts, train_labels = train_data["text"], train_data["label"]
     test_texts, test_labels = test_data["text"], test_data["label"]
 
@@ -53,17 +52,27 @@ async def evaluate_text(request: TextEvaluationRequest):
     tracker.start()
     tracker.start_task("inference")
 
-    # TF-IDF Vectorization
-    vectorizer = TfidfVectorizer(max_features=10000, ngram_range=(1, 2), stop_words="english")
-    X_train = vectorizer.fit_transform(train_texts)
-    X_test = vectorizer.transform(test_texts)
+    # Define the pipeline with TF-IDF and Logistic Regression
+    pipeline = Pipeline([
+        ('tfidf', TfidfVectorizer(max_features=10000, ngram_range=(1, 2), stop_words="english")),
+        ('clf', LogisticRegression(max_iter=1000, random_state=42))
+    ])
+
+    # Set up GridSearchCV for hyperparameter tuning
+    param_grid = {
+        'tfidf__max_features': [5000, 10000, 15000],
+        'tfidf__ngram_range': [(1, 1), (1, 2)],
+        'clf__C': [0.1, 1, 10]  # Regularization strength
+    }
+
+    grid_search = GridSearchCV(pipeline, param_grid, cv=3, scoring='accuracy', verbose=2)
+    grid_search.fit(train_texts, train_labels)
 
-    # Train SVM Classifier
-    svm_model = SVC(kernel="linear", probability=True)
-    svm_model.fit(X_train, train_labels)
+    # Get best estimator from GridSearch
+    best_model = grid_search.best_estimator_
 
     # Model Inference
-    predictions = svm_model.predict(X_test)
+    predictions = best_model.predict(test_texts)
 
     # Stop tracking emissions
     emissions_data = tracker.stop_task()
@@ -85,7 +94,8 @@ async def evaluate_text(request: TextEvaluationRequest):
         "dataset_config": {
             "dataset_name": request.dataset_name,
             "test_size": len(test_data),
-        }
+        },
+        "best_params": grid_search.best_params_
     }
 
     return results