Update tasks/text.py

tasks/text.py

@@ -8,9 +8,11 @@ from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
 
 import numpy as np
+from climate_model import ModelWrapper
+from preprocessing import ClimateTextPreprocessor
+import torch
 
 
 router = APIRouter()
@@ -61,14 +63,13 @@ async def evaluate_text(request: TextEvaluationRequest):
     #--------------------------------------------------------------------------------------------   
 
 
-    # Model and Tokenizer
+        # Model and Tokenizer
     model_name = "Zen0/FrugalDisinfoHunter"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForSequenceClassification.from_pretrained(model_name)
-
     
     # Tokenize the test data
-    test_texts = test_dataset["text"]  # Extracting the 'text' column (quotes)
+    test_texts = test_dataset["quote"]  # Changed from "text" to "quote"
     inputs = tokenizer(test_texts, padding=True, truncation=True, return_tensors="pt", max_length=512)
     
     # Move model and inputs to GPU if available
@@ -81,9 +82,9 @@ async def evaluate_text(request: TextEvaluationRequest):
         outputs = model(**inputs)
         logits = outputs.logits
     
-    # Get predictions from the logits (choose the class with the highest logit)
+    # Get predictions from the logits
     predictions = torch.argmax(logits, dim=-1).cpu().numpy()
-
+    
     true_labels = test_dataset['label']
 
     #--------------------------------------------------------------------------------------------