Update tasks/text.py

tasks/text.py

@@ -1,15 +1,18 @@
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from fastapi import APIRouter
 from datetime import datetime
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
-import random
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
+import numpy as np
+import torch
+
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = "FrugalDisinfoHunter Model"
 ROUTE = "/text"
 
 @router.post(ROUTE, tags=["Text Task"], description=DESCRIPTION)
@@ -46,32 +49,83 @@ async def evaluate_text(request: TextEvaluationRequest):
     tracker.start()
     tracker.start_task("inference")
 
-    # Get true labels
-    true_labels = test_dataset["label"]
-    predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
+    try:
+        # Model configuration
+        model_name = "Zen0/FrugalDisinfoHunter"  # Model path
+        tokenizer_name = "google/mobilebert-uncased"  # Base MobileBERT tokenizer
+        BATCH_SIZE = 32  # Batch size for efficient processing
+        MAX_LENGTH = 128  # Maximum sequence length
 
-    # Stop tracking emissions
-    emissions_data = tracker.stop_task()
-    
-    # Calculate accuracy
-    accuracy = accuracy_score(true_labels, predictions)
-    
-    # Prepare results dictionary
-    results = {
-        "username": username,
-        "space_url": space_url,
-        "submission_timestamp": datetime.now().isoformat(),
-        "model_description": DESCRIPTION,
-        "accuracy": float(accuracy),
-        "energy_consumed_wh": emissions_data.energy_consumed * 1000,
-        "emissions_gco2eq": emissions_data.emissions * 1000,
-        "emissions_data": clean_emissions_data(emissions_data),
-        "api_route": ROUTE,
-        "dataset_config": {
-            "dataset_name": request.dataset_name,
-            "test_size": request.test_size,
-            "test_seed": request.test_seed
+        # Initialize model and tokenizer
+        model = AutoModelForSequenceClassification.from_pretrained(
+            model_name,
+            num_labels=8,
+            output_hidden_states=True,
+            problem_type="single_label_classification"
+        )
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+
+        # Move model to appropriate device
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model = model.to(device)
+        model.eval()  # Set model to evaluation mode
+
+        # Get test texts
+        test_texts = test_dataset["quote"]
+        predictions = []
+
+        # Process in batches
+        for i in range(0, len(test_texts), BATCH_SIZE):
+            batch_texts = test_texts[i:i + BATCH_SIZE]
+            
+            # Tokenize batch
+            inputs = tokenizer(
+                batch_texts,
+                padding=True,
+                truncation=True,
+                return_tensors="pt",
+                max_length=MAX_LENGTH
+            )
+            
+            # Move inputs to device
+            inputs = {key: val.to(device) for key, val in inputs.items()}
+
+            # Run inference
+            with torch.no_grad():
+                outputs = model(**inputs)
+                batch_preds = torch.argmax(outputs.logits, dim=1)
+                predictions.extend(batch_preds.cpu().numpy())
+
+        # Get true labels
+        true_labels = test_dataset['label']
+
+        # Stop tracking emissions
+        emissions_data = tracker.stop_task()
+        
+        # Calculate accuracy
+        accuracy = accuracy_score(true_labels, predictions)
+        
+        # Prepare results dictionary
+        results = {
+            "username": username,
+            "space_url": space_url,
+            "submission_timestamp": datetime.now().isoformat(),
+            "model_description": DESCRIPTION,
+            "accuracy": float(accuracy),
+            "energy_consumed_wh": emissions_data.energy_consumed * 1000,
+            "emissions_gco2eq": emissions_data.emissions * 1000,
+            "emissions_data": clean_emissions_data(emissions_data),
+            "api_route": ROUTE,
+            "dataset_config": {
+                "dataset_name": request.dataset_name,
+                "test_size": request.test_size,
+                "test_seed": request.test_seed
+            }
         }
-    }
-    
-    return results
+        
+        return results
+        
+    except Exception as e:
+        # Stop tracking in case of error
+        tracker.stop_task()
+        raise e