Update tasks/text.py

tasks/text.py

@@ -1,10 +1,9 @@
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from fastapi import APIRouter
 from datetime import datetime
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
-from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.linear_model import LogisticRegression
-from sklearn.pipeline import Pipeline
+import torch
 import numpy as np
 
 from .utils.evaluation import TextEvaluationRequest
@@ -12,27 +11,9 @@ from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 router = APIRouter()
 
-DESCRIPTION = "Climate Disinformation Detection - TF-IDF + LogReg"
+DESCRIPTION = "Efficient Climate Disinformation Detection"
 ROUTE = "/text"
 
-def create_pipeline():
-    """Create an efficient text classification pipeline"""
-    return Pipeline([
-        ('tfidf', TfidfVectorizer(
-            max_features=10000,  # Limit features for efficiency
-            ngram_range=(1, 2),  # Use unigrams and bigrams
-            stop_words='english',
-            min_df=2,  # Remove very rare terms
-            max_df=0.95  # Remove very common terms
-        )),
-        ('classifier', LogisticRegression(
-            C=1.0,
-            multi_class='multinomial',
-            max_iter=200,
-            n_jobs=-1  # Use all CPU cores
-        ))
-    ])
-
 @router.post(ROUTE, tags=["Text Task"], description=DESCRIPTION)
 async def evaluate_text(request: TextEvaluationRequest):
     """
@@ -53,48 +34,74 @@ async def evaluate_text(request: TextEvaluationRequest):
         "7_fossil_fuels_needed": 7
     }
 
+    # Load and prepare the dataset
+    dataset = load_dataset(request.dataset_name)
+    dataset = dataset.map(lambda x: {"label": LABEL_MAPPING[x["label"]]})
+    train_test = dataset["train"].train_test_split(test_size=request.test_size, seed=request.test_seed)
+    test_dataset = train_test["test"]
+    
     # Start tracking emissions
     tracker.start()
     tracker.start_task("inference")
 
     try:
-        # Load and prepare the dataset
-        dataset = load_dataset(request.dataset_name)
-        
-        # Convert string labels to integers
-        dataset = dataset.map(lambda x: {"label": LABEL_MAPPING[x["label"]]})
-        
-        # Split dataset
-        train_test = dataset["train"].train_test_split(
-            test_size=request.test_size, 
-            seed=request.test_seed
-        )
-        
-        train_dataset = train_test["train"]
-        test_dataset = train_test["test"]
+        # Model configuration
+        model_name = "distilbert-base-uncased"  # Lighter model than MobileBERT
+        BATCH_SIZE = 64  # Increased batch size
+        MAX_LENGTH = 128  # Reduced sequence length
 
-        # Create and train pipeline
-        pipeline = create_pipeline()
-        
-        # Train the model
-        pipeline.fit(
-            train_dataset["quote"],
-            train_dataset["label"]
+        # Initialize tokenizer and model
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForSequenceClassification.from_pretrained(
+            model_name,
+            num_labels=8,
+            problem_type="single_label_classification",
         )
-        
-        # Make predictions
-        predictions = pipeline.predict(test_dataset["quote"])
-        
-        # Get true labels
-        true_labels = test_dataset["label"]
 
-        # Stop tracking emissions
+        # Enable mixed precision training if available
+        if torch.cuda.is_available():
+            model = model.half()  # Convert to FP16
+
+        # Move model to device
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model = model.to(device)
+        model.eval()
+
+        # Get test texts
+        test_texts = test_dataset["quote"]
+        predictions = []
+
+        # Process in efficient batches
+        for i in range(0, len(test_texts), BATCH_SIZE):
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+            batch_texts = test_texts[i:i + BATCH_SIZE]
+            
+            # Efficient tokenization
+            inputs = tokenizer(
+                batch_texts,
+                padding=True,
+                truncation=True,
+                max_length=MAX_LENGTH,
+                return_tensors="pt"
+            )
+            
+            # Move inputs to device efficiently
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+
+            # Inference with optimizations
+            with torch.no_grad(), torch.cuda.amp.autocast(enabled=torch.cuda.is_available()):
+                outputs = model(**inputs)
+                batch_preds = torch.argmax(outputs.logits, dim=1)
+                predictions.extend(batch_preds.cpu().numpy())
+
+        # Get true labels and compute accuracy
+        true_labels = test_dataset['label']
         emissions_data = tracker.stop_task()
-        
-        # Calculate accuracy
         accuracy = accuracy_score(true_labels, predictions)
         
-        # Prepare results dictionary
+        # Prepare results
         results = {
             "username": username,
             "space_url": space_url,
@@ -115,6 +122,5 @@ async def evaluate_text(request: TextEvaluationRequest):
         return results
         
     except Exception as e:
-        # Stop tracking in case of error
         tracker.stop_task()
         raise e