add model inference code

tasks/text.py

@@ -3,6 +3,9 @@ from datetime import datetime
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score
 import random
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from torch.utils.data import Dataset, DataLoader
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
@@ -51,14 +54,57 @@ async def evaluate_text(request: TextEvaluationRequest):
     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.
-    #--------------------------------------------------------------------------------------------   
+    # Load the model and tokenizer
+    model_name = "Tonic/climate-guard-toxic-agent"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name)
+
+    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()
+    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)
     
-    # Make random predictions (placeholder for actual model inference)
-    true_labels = test_dataset["label"]
-    predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
+            predictions.extend(predicted.cpu().numpy())
+            ground_truth.extend(labels.cpu().numpy())
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE
@@ -69,7 +115,7 @@ async def evaluate_text(request: TextEvaluationRequest):
     emissions_data = tracker.stop_task()
     
     # Calculate accuracy
-    accuracy = accuracy_score(true_labels, predictions)
+    accuracy = accuracy_score(ground_truth, predictions)
     
     # Prepare results dictionary
     results = {