Update tasks/text.py

tasks/text.py

@@ -1,26 +1,24 @@
 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.metrics import accuracy_score
-import random
+from sklearn.model_selection import train_test_split
+import numpy as np
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = "TF-IDF + SVM Classifier"
 ROUTE = "/text"
 
-@router.post(ROUTE, tags=["Text Task"], 
-             description=DESCRIPTION)
+@router.post(ROUTE, tags=["Text Task"], description=DESCRIPTION)
 async def evaluate_text(request: TextEvaluationRequest):
     """
-    Evaluate text classification for climate disinformation detection.
-    
-    Current Model: Random Baseline
-    - Makes random predictions from the label space (0-7)
-    - Used as a baseline for comparison
+    Evaluate text classification for climate disinformation detection using TF-IDF and SVM.
     """
     # Get space info
     username, space_url = get_space_info()
@@ -43,34 +41,36 @@ async def evaluate_text(request: TextEvaluationRequest):
     # Convert string labels to integers
     dataset = dataset.map(lambda x: {"label": LABEL_MAPPING[x["label"]]})
 
-    # Split dataset
-    train_test = dataset["train"]
-    test_dataset = dataset["test"]
-    
+    # Split dataset into training and testing sets
+    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"]
+
     # Start tracking emissions
     tracker.start()
     tracker.start_task("inference")
 
-    #--------------------------------------------------------------------------------------------
-    # YOUR MODEL INFERENCE CODE HERE
-    # 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))]
+    # 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)
+
+    # Train SVM Classifier
+    svm_model = SVC(kernel="linear", probability=True)
+    svm_model.fit(X_train, train_labels)
 
-    #--------------------------------------------------------------------------------------------
-    # YOUR MODEL INFERENCE STOPS HERE
-    #--------------------------------------------------------------------------------------------   
+    # Model Inference
+    predictions = svm_model.predict(X_test)
 
-    
     # Stop tracking emissions
     emissions_data = tracker.stop_task()
-    
+
     # Calculate accuracy
-    accuracy = accuracy_score(true_labels, predictions)
-    
+    accuracy = accuracy_score(test_labels, predictions)
+
     # Prepare results dictionary
     results = {
         "username": username,
@@ -84,9 +84,8 @@ async def evaluate_text(request: TextEvaluationRequest):
         "api_route": ROUTE,
         "dataset_config": {
             "dataset_name": request.dataset_name,
-            "test_size": request.test_size,
-            "test_seed": request.test_seed
+            "test_size": len(test_data),
         }
     }
-    
-    return results 
+
+    return results