revert inference code

tasks/text.py

@@ -1,151 +1,201 @@
-# tasks/text.py
-from fastapi import APIRouter, HTTPException
+from fastapi import APIRouter
 from datetime import datetime
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
+import random
+from transformers import pipeline, AutoConfig
+import os
+from concurrent.futures import ThreadPoolExecutor
+from typing import List, Dict, Tuple
+import numpy as np
 import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from torch.utils.data import Dataset, DataLoader
-import logging
 
 from .utils.evaluation import TextEvaluationRequest
-from .utils.emissions import start_tracking, stop_tracking, clean_emissions_data, get_space_info
+from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+# Disable torch compile
+os.environ["TORCH_COMPILE_DISABLE"] = "1"
 
 router = APIRouter()
 
-DESCRIPTION = "Climate Guard Toxic Agent Model"
+DESCRIPTION = "Random Baseline"
 ROUTE = "/text"
 
-class TextDataset(Dataset):
-    def __init__(self, texts, labels, tokenizer, max_len=128):
-        self.texts = texts
-        self.labels = labels
-        self.tokenizer = tokenizer
-        self.max_len = max_len
-
-    def __len__(self):
-        return len(self.texts)
-
-    def __getitem__(self, idx):
-        text = str(self.texts[idx])
-        label = self.labels[idx]
-        
-        encoding = self.tokenizer(
-            text,
-            max_length=self.max_len,
-            padding='max_length',
-            truncation=True,
-            return_tensors="pt"
-        )
-        
-        return {
-            'input_ids': encoding['input_ids'].squeeze(0),
-            'attention_mask': encoding['attention_mask'].squeeze(0),
-            'labels': torch.tensor(label, dtype=torch.long)
-        }
+class TextClassifier:
+    def __init__(self):
+        # Add retry mechanism for model initialization
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                self.config = AutoConfig.from_pretrained("Tonic/climate-guard-toxic-agent")
+                self.label2id = self.config.label2id
+                self.classifier = pipeline(
+                    "text-classification",
+                    "Tonic/climate-guard-toxic-agent",
+                    device="cpu",
+                    batch_size=16
+                )
+                print("Model initialized successfully")
+                break
+            except Exception as e:
+                if attempt == max_retries - 1:
+                    raise Exception(f"Failed to initialize model after {max_retries} attempts: {str(e)}")
+                print(f"Attempt {attempt + 1} failed, retrying...")
+                time.sleep(1)
+
+    def process_batch(self, batch: List[str], batch_idx: int) -> Tuple[List[int], int]:
+        """Process a batch of texts and return their predictions"""
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                print(f"Processing batch {batch_idx} with {len(batch)} items (attempt {attempt + 1})")
+                # Process texts one by one in case of errors
+                predictions = []
+                for text in batch:
+                    try:
+                        pred = self.classifier(text)
+                        pred_label = self.label2id[pred[0]["label"]]
+                        predictions.append(pred_label)
+                    except Exception as e:
+                        print(f"Error processing text in batch {batch_idx}: {str(e)}")
+                
+                if not predictions:
+                    raise Exception("No predictions generated for batch")
+                
+                print(f"Completed batch {batch_idx} with {len(predictions)} predictions")
+                return predictions, batch_idx
+            
+            except Exception as e:
+                if attempt == max_retries - 1:
+                    print(f"Final error in batch {batch_idx}: {str(e)}")
+                    return [0] * len(batch), batch_idx  # Return default predictions instead of empty list
+                print(f"Error in batch {batch_idx} (attempt {attempt + 1}): {str(e)}")
+                time.sleep(1)
+
 
-@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
     """
-    try:
-        logger.info("Starting evaluation")
-        username, space_url = get_space_info()
-
-        # Label mapping
-        LABEL_MAPPING = {
-            "0_not_relevant": 0,
-            "1_not_happening": 1,
-            "2_not_human": 2,
-            "3_not_bad": 3,
-            "4_solutions_harmful_unnecessary": 4,
-            "5_science_unreliable": 5,
-            "6_proponents_biased": 6,
-            "7_fossil_fuels_needed": 7
+    # Get space info
+    username, space_url = get_space_info()
+
+    # Define the label mapping
+    LABEL_MAPPING = {
+        "0_not_relevant": 0,
+        "1_not_happening": 1,
+        "2_not_human": 2,
+        "3_not_bad": 3,
+        "4_solutions_harmful_unnecessary": 4,
+        "5_science_unreliable": 5,
+        "6_proponents_biased": 6,
+        "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"]]})
+
+    # Split dataset
+    train_test = dataset["train"]
+    test_dataset = dataset["test"]
+    
+    # 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.
+    #--------------------------------------------------------------------------------------------   
+    
+    true_labels = test_dataset["label"]
+    
+    # Initialize the model once
+    classifier = TextClassifier()
+
+    # Prepare batches
+    batch_size = 32
+    quotes = test_dataset["quote"]
+    num_batches = len(quotes) // batch_size + (1 if len(quotes) % batch_size != 0 else 0)
+    batches = [
+        quotes[i * batch_size:(i + 1) * batch_size]
+        for i in range(num_batches)
+    ]
+
+    # Initialize batch_results before parallel processing
+    batch_results = [[] for _ in range(num_batches)]
+    
+    # Process batches in parallel
+    max_workers = min(os.cpu_count(), 4)  # Limit to 4 workers or CPU count
+    print(f"Processing with {max_workers} workers")
+    
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        # Submit all batches for processing
+        future_to_batch = {
+            executor.submit(
+                classifier.process_batch, 
+                batch, 
+                idx
+            ): idx for idx, batch in enumerate(batches)
         }
 
-        logger.info("Loading dataset")
-        dataset = load_dataset(request.dataset_name)
-        dataset = dataset.map(lambda x: {"label": LABEL_MAPPING[x["label"]]})
-        test_dataset = dataset["test"]
-        
-        logger.info("Starting emissions tracking")
-        start_tracking()
-        
-        try:
-            logger.info("Loading model and tokenizer")
-            model_name = "Tonic/climate-guard-toxic-agent"
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            model = AutoModelForSequenceClassification.from_pretrained(
-                model_name, 
-                num_labels=len(LABEL_MAPPING)
-            )
-            
-            logger.info("Preparing dataset")
-            test_data = TextDataset(
-                texts=test_dataset["text"],
-                labels=test_dataset["label"],
-                tokenizer=tokenizer
-            )
-            
-            test_loader = DataLoader(test_data, batch_size=16)
-            
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-            logger.info(f"Using device: {device}")
-            model = model.to(device)
-            model.eval()
-            
-            predictions = []
-            ground_truth = []
-            
-            logger.info("Running inference")
-            with torch.no_grad():
-                for batch in test_loader:
-                    input_ids = batch['input_ids'].to(device)
-                    attention_mask = batch['attention_mask'].to(device)
-                    labels = batch['labels'].to(device)
-                    
-                    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
-                    _, predicted = torch.max(outputs.logits, 1)
-                    
-                    predictions.extend(predicted.cpu().numpy())
-                    ground_truth.extend(labels.cpu().numpy())
-            
-            accuracy = accuracy_score(ground_truth, predictions)
-            logger.info(f"Accuracy: {accuracy}")
-            
-            emissions_data = stop_tracking()
-            
-            results = {
-                "username": username,
-                "space_url": space_url,
-                "submission_timestamp": datetime.now().isoformat(),
-                "model_description": DESCRIPTION,
-                "accuracy": float(accuracy),
-                "energy_consumed_wh": float(emissions_data.energy_consumed * 1000),
-                "emissions_gco2eq": float(emissions_data.emissions * 1000),
-                "emissions_data": clean_emissions_data(emissions_data.__dict__),
-                "api_route": ROUTE,
-                "dataset_config": {
-                    "dataset_name": request.dataset_name,
-                    "test_size": request.test_size,
-                    "test_seed": request.test_seed
-                }
-            }
-            
-            logger.info("Evaluation completed successfully")
-            return results
-            
-        except Exception as e:
-            logger.error(f"Error during evaluation: {str(e)}")
-            stop_tracking()
-            raise HTTPException(status_code=500, detail=str(e))
-            
-    except Exception as e:
-        logger.error(f"Error in evaluate_text: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
+        # Collect results in order
+        for future in future_to_batch:
+            batch_idx = future_to_batch[future]
+            try:
+                predictions, idx = future.result()
+                if predictions:  # Only store non-empty predictions
+                    batch_results[idx] = predictions
+                    print(f"Stored results for batch {idx} ({len(predictions)} predictions)")
+            except Exception as e:
+                print(f"Failed to get results for batch {batch_idx}: {e}")
+                # Use default predictions instead of empty list
+                batch_results[batch_idx] = [0] * len(batches[batch_idx])
+
+    # Flatten predictions while maintaining order
+    predictions = []
+    for batch_preds in batch_results:
+        if batch_preds is not None:
+            predictions.extend(batch_preds)
+
+    #--------------------------------------------------------------------------------------------
+    # YOUR MODEL INFERENCE STOPS HERE
+    #--------------------------------------------------------------------------------------------   
+    
+    # Stop tracking emissions
+    emissions_data = tracker.stop_task()
+    
+    # Calculate accuracy
+    accuracy = accuracy_score(true_labels, predictions)
+    print("accuracy : ", accuracy)
+    
+    # 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
+        }
+    }
+
+    print("results : ", results)
+    
+    return results