tasks/text.py

# tasks/text.py
from fastapi import APIRouter
from datetime import datetime
from datasets import load_dataset
from sklearn.metrics import accuracy_score
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 tracker, clean_emissions_data, get_space_info

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

router = APIRouter()

DESCRIPTION = "Climate Guard Toxic Agent Model"
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)
        }

@router.post(ROUTE, tags=["Text Task"], description=DESCRIPTION)
async def evaluate_text(request: TextEvaluationRequest):
    """
    Evaluate text classification for climate disinformation detection.
    """
    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
        }

        logger.info("Loading dataset")
        # Load 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"]
        
        logger.info("Starting emissions tracking")
        # Start tracking emissions
        tracker.start()
        
        try:
            # Load model and tokenizer
            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))
            
            # Prepare dataset
            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)
            
            # Model inference
            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())
            
            # Calculate accuracy
            accuracy = accuracy_score(ground_truth, predictions)
            logger.info(f"Accuracy: {accuracy}")
            
            # Stop tracking emissions
            emissions_data = tracker.stop()
            
            # Prepare results
            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),
                "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)}")
            tracker.stop()
            raise e
            
    except Exception as e:
        logger.error(f"Error in evaluate_text: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))