Update tasks/text.py

tasks/text.py

@@ -3,13 +3,12 @@ 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 .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
-import numpy as np
-import torch
-
 router = APIRouter()
 
 DESCRIPTION = "FrugalDisinfoHunter Model"
@@ -51,46 +50,58 @@ async def evaluate_text(request: TextEvaluationRequest):
 
     try:
         # Model configuration
-        model_name = "Zen0/FrugalDisinfoHunter"  # Model path
-        tokenizer_name = "google/mobilebert-uncased"  # Base MobileBERT tokenizer
-        BATCH_SIZE = 32  # Batch size for efficient processing
-        MAX_LENGTH = 128  # Maximum sequence length
+        model_name = "google/mobilebert-uncased"  # Base model
+        local_weights = "model/model.pt"  # Path to our trained weights
+        BATCH_SIZE = 32
+        MAX_LENGTH = 256  # Increased from 128
 
-        # Initialize model and tokenizer
+        # Initialize tokenizer and model
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
         model = AutoModelForSequenceClassification.from_pretrained(
             model_name,
             num_labels=8,
-            output_hidden_states=True,
             problem_type="single_label_classification"
         )
-        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+        
+        # Load our trained weights
+        try:
+            state_dict = torch.load(local_weights, map_location='cpu')
+            model.load_state_dict(state_dict)
+        except Exception as e:
+            print(f"Error loading weights: {e}")
+            # Continue with base model if weights fail to load
+            pass
 
         # Move model to appropriate device
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         model = model.to(device)
-        model.eval()  # Set model to evaluation mode
+        model.eval()  # Set to evaluation mode
 
-        # Get test texts
+        # Get test texts and process in batches
         test_texts = test_dataset["quote"]
         predictions = []
 
         # Process in batches
         for i in range(0, len(test_texts), BATCH_SIZE):
+            # Clear CUDA cache if using GPU
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
             batch_texts = test_texts[i:i + BATCH_SIZE]
             
-            # Tokenize batch
+            # Tokenize with padding and attention masks
             inputs = tokenizer(
                 batch_texts,
                 padding=True,
                 truncation=True,
-                return_tensors="pt",
-                max_length=MAX_LENGTH
+                max_length=MAX_LENGTH,
+                return_tensors="pt"
             )
             
             # Move inputs to device
-            inputs = {key: val.to(device) for key, val in inputs.items()}
+            inputs = {k: v.to(device) for k, v in inputs.items()}
 
-            # Run inference
+            # Run inference with no gradient computation
             with torch.no_grad():
                 outputs = model(**inputs)
                 batch_preds = torch.argmax(outputs.logits, dim=1)