Update tasks/text.py

tasks/text.py

@@ -1,10 +1,8 @@
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, MobileBertTokenizerFast
-
+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
@@ -12,21 +10,15 @@ from .utils.emissions import tracker, clean_emissions_data, get_space_info
 import numpy as np
 import torch
 
-
-
-
 router = APIRouter()
 
 DESCRIPTION = "FrugalDisinfoHunter Model"
 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: FrugalDisinfoHunter
     """
     # Get space info
     username, space_url = get_space_info()
@@ -57,69 +49,83 @@ async def evaluate_text(request: TextEvaluationRequest):
     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.
-    #--------------------------------------------------------------------------------------------   
-
-
-    # Model and Tokenizer
-    # Model and Tokenizer - use the same model name for both
-    model_name = "Zen0/FrugalDisinfoHunter"
-    model = AutoModelForSequenceClassification.from_pretrained(
-        model_name,
-        num_labels=8,
-        output_hidden_states=True,
-        problem_type="single_label_classification"
-    )
-    tokenizer = AutoTokenizer.from_pretrained(model_name)  # Use the same model_name
-    
-    # Tokenize the test data
-    test_texts = test_dataset["quote"] 
-    inputs = tokenizer(test_texts, padding=True, truncation=True, return_tensors="pt", max_length=512)
-    
-    # Move model and inputs to GPU if available
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
-    inputs = {key: val.to(device) for key, val in inputs.items()}
-    
-    # Run inference on the dataset using the model
-    with torch.no_grad():  # Disable gradient calculations
-        outputs = model(**inputs)
-        logits = outputs.logits
-    
-    # Get predictions from the logits
-    predictions = torch.argmax(logits, dim=-1).cpu().numpy()
-    
-    true_labels = test_dataset['label']
-
-    #--------------------------------------------------------------------------------------------
-    # YOUR MODEL INFERENCE STOPS HERE
-    #--------------------------------------------------------------------------------------------   
-
-    
-    # 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
+    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
+
+        # 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