revert to template

tasks/text.py

@@ -1,92 +1,30 @@
 from fastapi import APIRouter
 from datetime import datetime
-import time
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
-import os
-from concurrent.futures import ThreadPoolExecutor
-from typing import List, Dict, Tuple
 import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, BertForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from torch.utils.data import DataLoader
 from transformers import DataCollatorWithPadding
-from huggingface_hub import login
-from dotenv import load_dotenv
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
-# Load environment variables
-load_dotenv()
-
-# Authenticate with Hugging Face
-HF_TOKEN = os.getenv('HF_TOKEN')
-if HF_TOKEN:
-    login(token=HF_TOKEN)
-
 router = APIRouter()
 
-DESCRIPTION = "Climate Guard Toxic Agent is a ModernBERT for Climate Disinformation Detection"
+DESCRIPTION = "ModernBERT for Climate Disinformation Detection"
 ROUTE = "/text"
-MODEL_NAME = "Tonic/climate-guard-toxic-agent"
-TOKENIZER_NAME = "answerdotai/ModernBERT-base"
-
-class TextClassifier:
-    def __init__(self):
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        
-        try:
-            # Initialize tokenizer
-            self.tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
-            
-            # Initialize model with auto class
-            self.model = AutoModelForSequenceClassification.from_pretrained(
-                MODEL_NAME,
-                trust_remote_code=True,
-                num_labels=8,
-                problem_type="single_label_classification",
-                ignore_mismatched_sizes=True
-            ).to(self.device)
-            
-            # Convert to half precision and eval mode
-            self.model = self.model.half()
-            self.model.eval()
-            
-            print("Model initialized successfully")
-                
-        except Exception as e:
-            print(f"Error initializing model: {str(e)}")
-            raise
-
-    def process_batch(self, batch):
-        try:
-            # Move batch to device
-            input_ids = batch['input_ids'].to(self.device)
-            attention_mask = batch['attention_mask'].to(self.device)
-            
-            # Get predictions
-            with torch.no_grad():
-                outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
-                predictions = torch.argmax(outputs.logits, dim=-1)
-                
-            return predictions.cpu().numpy().tolist()
-            
-        except Exception as e:
-            print(f"Error in batch processing: {str(e)}")
-            return [0] * len(batch['input_ids'])
-
-    def __del__(self):
-        if hasattr(self, 'model'):
-            del self.model
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
 
-@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."""
-    
+    """
+    Evaluate text classification for climate disinformation detection using ModernBERT.
+    """
+    # Get space info
     username, space_url = get_space_info()
 
+    # Define the label mapping
     LABEL_MAPPING = {
         "0_not_relevant": 0,
         "1_not_happening": 1,
@@ -98,29 +36,51 @@ async def evaluate_text(request: TextEvaluationRequest):
         "7_fossil_fuels_needed": 7
     }
 
+    # 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"]]})
+
+    # Get test dataset
+    test_dataset = dataset["test"]
+    
+    # Start tracking emissions
+    tracker.start()
+    tracker.start_task("inference")
+
+    #--------------------------------------------------------------------------------------------
+    # MODEL INFERENCE CODE
+    #--------------------------------------------------------------------------------------------   
+    
     try:
-        # Load dataset
-        dataset = load_dataset(request.dataset_name)
+        # Set device
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         
-        # Convert labels
-        dataset = dataset.map(lambda x: {"label": LABEL_MAPPING[x["label"]]})
-        test_dataset = dataset["test"]
+        # Initialize tokenizer
+        tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
         
-        # Start tracking emissions
-        tracker.start()
-        tracker.start_task("inference")
-
-        # Initialize model
-        classifier = TextClassifier()
+        # Initialize model with configuration that avoids bias parameter
+        model = AutoModelForSequenceClassification.from_pretrained(
+            "Tonic/climate-guard-toxic-agent",
+            trust_remote_code=True,
+            num_labels=8,
+            problem_type="single_label_classification",
+            ignore_mismatched_sizes=True,
+            torch_dtype=torch.float16  # Use float16 for efficiency
+        ).to(device)
+        
+        # Set model to evaluation mode
+        model.eval()
         
-        # Prepare tokenization function
+        # Tokenize function
         def preprocess_function(examples):
-            return classifier.tokenizer(
+            return tokenizer(
                 examples["quote"],
                 truncation=True,
                 padding=True,
                 max_length=512,
-                return_tensors=None  # Changed this to None for batched processing
+                return_tensors=None
             )
         
         # Tokenize dataset
@@ -134,7 +94,7 @@ async def evaluate_text(request: TextEvaluationRequest):
         tokenized_test.set_format("torch")
         
         # Create DataLoader
-        data_collator = DataCollatorWithPadding(tokenizer=classifier.tokenizer)
+        data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
         test_loader = DataLoader(
             tokenized_test,
             batch_size=16,
@@ -143,37 +103,54 @@ async def evaluate_text(request: TextEvaluationRequest):
         )
         
         # Get predictions
-        all_predictions = []
-        for batch in test_loader:
-            batch_preds = classifier.process_batch(batch)
-            all_predictions.extend(batch_preds)
-        
-        # Stop tracking emissions
-        emissions_data = tracker.stop_task()
-        
-        # Calculate accuracy
-        accuracy = accuracy_score(test_dataset["label"], all_predictions)
-        
-        # Prepare results
-        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
-            }
-        }
+        predictions = []
+        with torch.no_grad():
+            for batch in test_loader:
+                # Move batch to device
+                input_ids = batch['input_ids'].to(device)
+                attention_mask = batch['attention_mask'].to(device)
+                
+                # Get model outputs
+                outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+                preds = torch.argmax(outputs.logits, dim=-1)
+                
+                # Add batch predictions to list
+                predictions.extend(preds.cpu().numpy().tolist())
+                
+        # Clean up GPU memory
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            
+    except Exception as e:
+        print(f"Error during model inference: {str(e)}")
+        raise
 
-        return results
+    #--------------------------------------------------------------------------------------------
+    # MODEL INFERENCE ENDS HERE
+    #--------------------------------------------------------------------------------------------   
 
-    except Exception as e:
-        print(f"Error in evaluate_text: {str(e)}")
-        raise Exception(f"Failed to process request: {str(e)}")
+    # Stop tracking emissions
+    emissions_data = tracker.stop_task()
+    
+    # Calculate accuracy
+    accuracy = accuracy_score(test_dataset["label"], 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