Update tasks/text.py

tasks/text.py

@@ -1,32 +1,20 @@
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from fastapi import FastAPI, APIRouter
-from fastapi.middleware.cors import CORSMiddleware
+from fastapi import APIRouter
 from datetime import datetime
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
 import torch
-import numpy as np
+from torch.utils.data import Dataset, DataLoader
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
-# Initialize FastAPI app and router
-app = FastAPI()
 router = APIRouter()
 
-# Add CORS middleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-DESCRIPTION = "Efficient Climate Disinformation Detection"
+DESCRIPTION = "Climate Disinformation Detection"
 ROUTE = "/text"
 
-@router.post("/text", 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.
@@ -48,7 +36,11 @@ async def evaluate_text(request: TextEvaluationRequest):
 
     # 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"].train_test_split(test_size=request.test_size, seed=request.test_seed)
     test_dataset = train_test["test"]
     
@@ -57,65 +49,69 @@ async def evaluate_text(request: TextEvaluationRequest):
     tracker.start_task("inference")
 
     try:
-        # Model configuration
-        model_name = "distilbert-base-uncased"
-        BATCH_SIZE = 64
-        MAX_LENGTH = 128
-
-        # Initialize tokenizer and model
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForSequenceClassification.from_pretrained(
-            model_name,
-            num_labels=8,
-            problem_type="single_label_classification"
-        )
-
-        # Enable mixed precision if available
-        if torch.cuda.is_available():
-            model = model.half()
-
-        # Move model to device
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        model = model.to(device)
+        # Get texts and labels
+        texts = test_dataset["quote"]
+        labels = test_dataset["label"]
+        
+        # Load model and tokenizer from local directory
+        model_dir = "./"
+        tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        model = AutoModelForSequenceClassification.from_pretrained(model_dir)
+        
+        # Define dataset class
+        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 = self.texts[idx]
+                label = self.labels[idx]
+                encodings = self.tokenizer(
+                    text,
+                    max_length=self.max_len,
+                    padding='max_length',
+                    truncation=True,
+                    return_tensors="pt"
+                )
+                return {
+                    'input_ids': encodings['input_ids'].squeeze(0),
+                    'attention_mask': encodings['attention_mask'].squeeze(0),
+                    'labels': torch.tensor(label, dtype=torch.long)
+                }
+
+        # Create dataset and dataloader
+        test_dataset = TextDataset(texts, labels, tokenizer)
+        test_loader = DataLoader(test_dataset, batch_size=16)
+
+        # Model inference
         model.eval()
-
-        # Get test texts
-        test_texts = test_dataset["quote"]
         predictions = []
+        ground_truth = []
+        device = 'cpu'
+        
+        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)
 
-        # Process in batches
-        for i in range(0, len(test_texts), BATCH_SIZE):
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-
-            batch_texts = test_texts[i:i + BATCH_SIZE]
-            
-            # Tokenize batch
-            inputs = tokenizer(
-                batch_texts,
-                padding=True,
-                truncation=True,
-                max_length=MAX_LENGTH,
-                return_tensors="pt"
-            )
-            
-            # Move inputs to device
-            inputs = {k: v.to(device) for k, v in inputs.items()}
-
-            # Run inference
-            with torch.no_grad(), torch.cuda.amp.autocast(enabled=torch.cuda.is_available()):
-                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']
+                predictions.extend(predicted.cpu().numpy())
+                ground_truth.extend(labels.cpu().numpy())
 
         # Stop tracking emissions
         emissions_data = tracker.stop_task()
         
         # Calculate accuracy
-        accuracy = accuracy_score(true_labels, predictions)
+        accuracy = accuracy_score(test_dataset["label"], predictions)
         
         # Prepare results
         results = {
@@ -138,13 +134,6 @@ async def evaluate_text(request: TextEvaluationRequest):
         return results
         
     except Exception as e:
+        # Stop tracking in case of error
         tracker.stop_task()
-        raise e
-
-# Include the router
-app.include_router(router)
-
-# Add a health check endpoint
-@app.get("/health")
-async def health_check():
-    return {"status": "healthy"}
+        raise e