Update tasks/text.py

tasks/text.py

@@ -2,30 +2,63 @@ from fastapi import APIRouter
 from datetime import datetime
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
-import random
+import torch
+from transformers import AutoTokenizer, RobertaForSequenceClassification
+from torch.utils.data import Dataset, DataLoader
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = "RoBERTa Climate Disinformation Classifier"
 ROUTE = "/text"
 
-@router.post(ROUTE, tags=["Text Task"], 
-             description=DESCRIPTION)
+class FrugalDataClass(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]
+
+        encodings = self.tokenizer(
+            text,
+            max_length=self.max_len,
+            padding='max_length',
+            truncation=True,
+            return_tensors="pt"
+        )
+
+        return {
+            'input_ids': encodings['input_ids'].flatten(),
+            'attention_mask': encodings['attention_mask'].flatten(),
+            'labels': torch.tensor(label, dtype=torch.long)
+        }
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+tokenizer = AutoTokenizer.from_pretrained("roberta-base")
+model = RobertaForSequenceClassification.from_pretrained(
+    "roberta-base",
+    num_labels=8
+)
+model.load_state_dict(torch.load('best_roberta_model.pth', map_location=device))
+model.to(device)
+model.eval()
+
+@router.post(ROUTE, 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
+    Evaluate text classification for climate disinformation detection using RoBERTa.
     """
-    # Get space info
     username, space_url = get_space_info()
 
-    # Define the label mapping
     LABEL_MAPPING = {
         "0_not_relevant": 0,
         "1_not_happening": 1,
@@ -37,41 +70,35 @@ 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"]]})
 
-    # Split dataset
     train_test = dataset["train"].train_test_split(test_size=request.test_size, seed=request.test_seed)
     test_dataset = train_test["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.
-    #--------------------------------------------------------------------------------------------   
-    
-    # Make random predictions (placeholder for actual model inference)
+    test_texts = test_dataset["quote"]
     true_labels = test_dataset["label"]
-    predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
-
-    #--------------------------------------------------------------------------------------------
-    # YOUR MODEL INFERENCE STOPS HERE
-    #--------------------------------------------------------------------------------------------   
-
     
-    # Stop tracking emissions
+    test_dataset = FrugalDataClass(test_texts, true_labels, tokenizer)
+    test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False)
+    
+    predictions = []
+    with torch.no_grad():
+        for batch in test_loader:
+            input_ids = batch['input_ids'].to(device)
+            attention_mask = batch['attention_mask'].to(device)
+            outputs = model(input_ids, attention_mask=attention_mask)
+            preds = torch.argmax(outputs.logits, dim=1).cpu().numpy()
+            predictions.extend(preds)
+    
     emissions_data = tracker.stop_task()
     
-    # Calculate accuracy
     accuracy = accuracy_score(true_labels, predictions)
     
-    # Prepare results dictionary
     results = {
         "username": username,
         "space_url": space_url,
@@ -89,4 +116,4 @@ async def evaluate_text(request: TextEvaluationRequest):
         }
     }
     
-    return results 
+    return results