revert to direct model loading

requirements.txt

@@ -9,4 +9,5 @@ python-dotenv==1.0.0
 requests==2.31.0
 numpy==1.24.3
 pydantic==2.4.2
-accelerate
+accelerate
+huggingface-hub

tasks/text.py

@@ -7,7 +7,7 @@ import os
 from concurrent.futures import ThreadPoolExecutor
 from typing import List, Dict, Tuple
 import torch
-from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info, start_tracking, stop_tracking
@@ -26,12 +26,11 @@ class TextClassifier:
         max_retries = 3
         for attempt in range(max_retries):
             try:
-                # Initialize using pipeline instead
-                self.classifier = pipeline(
-                    "text-classification",
-                    model="Tonic/climate-guard-toxic-agent",
-                    device=self.device
-                )
+                # Initialize tokenizer and model separately
+                self.tokenizer = AutoTokenizer.from_pretrained("Tonic/climate-guard-toxic-agent")
+                self.model = AutoModelForSequenceClassification.from_pretrained("Tonic/climate-guard-toxic-agent")
+                self.model.to(self.device)
+                self.model.eval()
                 print("Model initialized successfully")
                 break
             except Exception as e:
@@ -43,11 +42,20 @@ class TextClassifier:
     def predict_single(self, text: str) -> int:
         """Predict single text instance"""
         try:
-            result = self.classifier(text)
-            # Extract the label index from the result
-            # Assuming the model outputs label indices directly
-            label = int(result[0]['label'].split('_')[0])
-            return label
+            # Tokenize and prepare input
+            inputs = self.tokenizer(
+                text,
+                return_tensors="pt",
+                truncation=True,
+                max_length=512,
+                padding=True
+            ).to(self.device)
+            
+            # Get prediction
+            with torch.no_grad():
+                outputs = self.model(**inputs)
+                predictions = outputs.logits.argmax(-1)
+                return predictions.item()
         except Exception as e:
             print(f"Error in single prediction: {str(e)}")
             return 0  # Return default prediction on error