update imports , textclassifier

tasks/text.py

@@ -8,7 +8,7 @@ from concurrent.futures import ThreadPoolExecutor
 from typing import List, Dict, Tuple
 import torch
 import torch.nn as nn
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoConfig
+from transformers import AutoTokenizer, pipeline
 from huggingface_hub import login
 from dotenv import load_dotenv
 
@@ -42,34 +42,22 @@ class TextClassifier:
                 # Initialize tokenizer first
                 self.tokenizer = AutoTokenizer.from_pretrained(
                     model_name,
-                    model_max_length=8192,
+                    model_max_length=512,  # Reduced from 8192
                     padding_side='right',
                     truncation_side='right'
                 )
                 
-                # Load base config
-                self.config = AutoConfig.from_pretrained(
-                    model_name,
-                    num_labels=8,
-                    problem_type="single_label_classification"
-                )
-                
-                # Set required attributes
-                self.config.hidden_size = 768
-                self.config.num_attention_heads = 12
-                self.config.num_hidden_layers = 12
-                self.config.norm_eps = 1e-5
-                
-                # Initialize model with basic config
-                self.model = AutoModelForSequenceClassification.from_pretrained(
-                    model_name,
-                    config=self.config,
-                    ignore_mismatched_sizes=True
+                # Use pipeline for simpler initialization
+                self.classifier = pipeline(
+                    "text-classification",
+                    model=model_name,
+                    tokenizer=self.tokenizer,
+                    device=self.device,
+                    max_length=512,
+                    truncation=True,
+                    batch_size=32
                 )
                 
-                # Move model to appropriate device
-                self.model = self.model.to(self.device)
-                self.model.eval()
                 print("Model initialized successfully")
                 break
                 
@@ -84,22 +72,9 @@ class TextClassifier:
         try:
             print(f"Processing batch {batch_idx} with {len(batch)} items")
             
-            # Tokenize with padding and truncation
-            inputs = self.tokenizer(
-                batch,
-                return_tensors="pt",
-                truncation=True,
-                max_length=512,
-                padding=True
-            )
-            
-            # Move inputs to device
-            inputs = {k: v.to(self.device) for k, v in inputs.items()}
-            
-            # Get predictions
-            with torch.no_grad():
-                outputs = self.model(**inputs)
-                predictions = torch.argmax(outputs.logits, dim=-1).cpu().tolist()
+            # Use pipeline for prediction
+            outputs = self.classifier(batch)
+            predictions = [int(output['label'].split('_')[0]) for output in outputs]
             
             print(f"Completed batch {batch_idx} with {len(predictions)} predictions")
             return predictions, batch_idx
@@ -110,11 +85,10 @@ class TextClassifier:
 
     def __del__(self):
         # Clean up CUDA memory
-        if hasattr(self, 'model'):
-            del self.model
+        if hasattr(self, 'classifier'):
+            del self.classifier
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
-    
 
 @router.post(ROUTE, tags=["Text Task"], description=DESCRIPTION)
 async def evaluate_text(request: TextEvaluationRequest):