revert to model loading

tasks/text.py

@@ -27,7 +27,7 @@ os.environ["TORCH_COMPILE_DISABLE"] = "1"
 
 router = APIRouter()
 
-DESCRIPTION = "ModernBERT fine-tuned for climate disinformation detection"
+DESCRIPTION = "Climate Guard Toxic Agent is a ModernBERT fine-tuned for climate disinformation detection"
 ROUTE = "/text"
 MODEL_NAME = "Tonic/climate-guard-toxic-agent"
 
@@ -38,7 +38,7 @@ class TextClassifier:
         
         for attempt in range(max_retries):
             try:
-                # Initialize tokenizer first
+                # Initialize tokenizer
                 self.tokenizer = AutoTokenizer.from_pretrained(
                     MODEL_NAME,
                     model_max_length=512,
@@ -46,17 +46,18 @@ class TextClassifier:
                     truncation_side='right'
                 )
                 
-                # Use pipeline directly without modifying config
-                self.classifier = pipeline(
-                    "text-classification",
-                    model=MODEL_NAME,
-                    tokenizer=self.tokenizer,
-                    device=self.device,
-                    max_length=512,
-                    truncation=True,
-                    batch_size=16
+                # Initialize model with basic configuration
+                self.model = AutoModelForSequenceClassification.from_pretrained(
+                    MODEL_NAME,
+                    num_labels=8,
+                    problem_type="single_label_classification",
+                    ignore_mismatched_sizes=True,
+                    trust_remote_code=True
                 )
                 
+                # Move model to device
+                self.model = self.model.to(self.device)
+                
                 print("Model initialized successfully")
                 break
                 
@@ -68,42 +69,34 @@ class TextClassifier:
 
     def process_batch(self, batch: List[str], batch_idx: int) -> Tuple[List[int], int]:
         """Process a batch of texts and return their predictions"""
-        max_retries = 3
-        for attempt in range(max_retries):
-            try:
-                print(f"Processing batch {batch_idx} with {len(batch)} items")
-                
-                # Process texts with error handling
-                predictions = []
-                for text in batch:
-                    try:
-                        result = self.classifier(text)
-                        # Extract the numeric label from the prediction
-                        label_str = result[0]['label']
-                        # Handle both numeric and string label formats
-                        if '_' in label_str:
-                            pred_label = int(label_str.split('_')[0])
-                        else:
-                            pred_label = int(label_str)
-                        predictions.append(pred_label)
-                    except Exception as e:
-                        print(f"Error processing text in batch {batch_idx}: {str(e)}")
-                        predictions.append(0)  # Default prediction
-                
-                print(f"Completed batch {batch_idx} with {len(predictions)} predictions")
-                return predictions, batch_idx
-                
-            except Exception as e:
-                if attempt == max_retries - 1:
-                    print(f"Final error in batch {batch_idx}: {str(e)}")
-                    return [0] * len(batch), batch_idx
-                print(f"Error in batch {batch_idx} (attempt {attempt + 1}): {str(e)}")
-                time.sleep(1)
+        try:
+            print(f"Processing batch {batch_idx} with {len(batch)} items")
+            
+            # Tokenize texts
+            inputs = self.tokenizer(
+                batch,
+                padding=True,
+                truncation=True,
+                max_length=512,
+                return_tensors="pt"
+            ).to(self.device)
+            
+            # Get predictions
+            with torch.no_grad():
+                outputs = self.model(**inputs)
+                predictions = torch.argmax(outputs.logits, dim=1).cpu().numpy()
+            
+            print(f"Completed batch {batch_idx} with {len(predictions)} predictions")
+            return predictions.tolist(), batch_idx
+            
+        except Exception as e:
+            print(f"Error in batch {batch_idx}: {str(e)}")
+            return [0] * len(batch), batch_idx
 
     def __del__(self):
         # Clean up CUDA memory
-        if hasattr(self, 'classifier'):
-            del self.classifier
+        if hasattr(self, 'model'):
+            del self.model
         if torch.cuda.is_available():
             torch.cuda.empty_cache()