add pipeline

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
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info, start_tracking, stop_tracking
@@ -26,14 +26,12 @@ class TextClassifier:
         max_retries = 3
         for attempt in range(max_retries):
             try:
-                # Load model and tokenizer directly instead of using pipeline
-                self.model = AutoModelForSequenceClassification.from_pretrained(
-                    "Tonic/climate-guard-toxic-agent"
-                ).to(self.device)
-                self.tokenizer = AutoTokenizer.from_pretrained(
-                    "Tonic/climate-guard-toxic-agent"
+                # Initialize using pipeline instead
+                self.classifier = pipeline(
+                    "text-classification",
+                    model="Tonic/climate-guard-toxic-agent",
+                    device=self.device
                 )
-                self.model.eval()  # Set to evaluation mode
                 print("Model initialized successfully")
                 break
             except Exception as e:
@@ -45,18 +43,11 @@ class TextClassifier:
     def predict_single(self, text: str) -> int:
         """Predict single text instance"""
         try:
-            inputs = self.tokenizer(
-                text,
-                return_tensors="pt",
-                truncation=True,
-                max_length=512,
-                padding=True
-            ).to(self.device)
-            
-            with torch.no_grad():
-                outputs = self.model(**inputs)
-                predictions = outputs.logits.argmax(-1)
-                return predictions.item()
+            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
         except Exception as e:
             print(f"Error in single prediction: {str(e)}")
             return 0  # Return default prediction on error
@@ -114,15 +105,13 @@ async def evaluate_text(request: TextEvaluationRequest):
     # Start tracking emissions
     start_tracking()
 
-    # tracker.start_task("inference")
-
     true_labels = test_dataset["label"]
     
     # Initialize the model once
     classifier = TextClassifier()
 
     # Prepare batches
-    batch_size = 16  # Reduced batch size for better memory management
+    batch_size = 16
     quotes = test_dataset["quote"]
     num_batches = len(quotes) // batch_size + (1 if len(quotes) % batch_size != 0 else 0)
     batches = [
@@ -162,7 +151,6 @@ async def evaluate_text(request: TextEvaluationRequest):
     
     # Stop tracking emissions
     emissions_data = stop_tracking()
-#    emissions_data = tracker.stop_task()
     
     # Calculate accuracy
     accuracy = accuracy_score(true_labels, predictions)