switch model loading technique

tasks/text.py

@@ -31,11 +31,6 @@ router = APIRouter()
 DESCRIPTION = "Climate Guard Toxic Agent Classifier"
 ROUTE = "/text"
 
-# Custom LayerNorm that ignores bias parameter
-class CustomLayerNorm(nn.LayerNorm):
-    def __init__(self, normalized_shape, eps=1e-5, **kwargs):
-        super().__init__(normalized_shape, eps=eps)
-
 class TextClassifier:
     def __init__(self):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -45,33 +40,34 @@ class TextClassifier:
         for attempt in range(max_retries):
             try:
                 # Load config
-                self.config = AutoConfig.from_pretrained(model_name)
+                self.config = AutoConfig.from_pretrained(
+                    model_name,
+                    trust_remote_code=True
+                )
                 
                 # Initialize tokenizer
                 self.tokenizer = AutoTokenizer.from_pretrained(
                     model_name,
                     model_max_length=2048,
                     padding_side='right',
-                    truncation_side='right'
+                    truncation_side='right',
+                    trust_remote_code=True
                 )
                 
-                # Patch LayerNorm
-                original_layernorm = nn.LayerNorm
-                nn.LayerNorm = CustomLayerNorm
+                # Initialize model
+                self.model = AutoModelForSequenceClassification.from_pretrained(
+                    model_name,
+                    config=self.config,
+                    trust_remote_code=True,
+                    torch_dtype=torch.float32,
+                    device_map="auto",
+                    low_cpu_mem_usage=True
+                )
                 
-                try:
-                    # Initialize model with patched LayerNorm
-                    self.model = AutoModelForSequenceClassification.from_pretrained(
-                        model_name,
-                        config=self.config,
-                        ignore_mismatched_sizes=True,
-                        low_cpu_mem_usage=True
-                    )
-                finally:
-                    # Restore original LayerNorm
-                    nn.LayerNorm = original_layernorm
+                # Force model to CPU if CUDA is not available
+                if not torch.cuda.is_available():
+                    self.model = self.model.cpu()
                 
-                self.model.to(self.device)
                 self.model.eval()
                 print("Model initialized successfully")
                 break
@@ -87,14 +83,17 @@ class TextClassifier:
         try:
             print(f"Processing batch {batch_idx} with {len(batch)} items")
             
-            # Tokenize
+            # Tokenize with smaller max length
             inputs = self.tokenizer(
                 batch,
                 return_tensors="pt",
                 truncation=True,
-                max_length=2048,
-                padding='max_length'
-            ).to(self.device)
+                max_length=512,  # Reduced max length
+                padding=True
+            )
+            
+            # Move inputs to device
+            inputs = {k: v.to(self.device) for k, v in inputs.items()}
             
             # Get predictions
             with torch.no_grad():
@@ -108,6 +107,13 @@ class TextClassifier:
             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, 'model'):
+            del self.model
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
 @router.post(ROUTE, tags=["Text Task"], description=DESCRIPTION)
 async def evaluate_text(request: TextEvaluationRequest):
     """Evaluate text classification for climate disinformation detection."""
@@ -128,8 +134,21 @@ async def evaluate_text(request: TextEvaluationRequest):
     }
 
     try:
-        # Load and prepare the dataset
-        dataset = load_dataset("QuotaClimat/frugalaichallenge-text-train", token=HF_TOKEN)
+        # Load and prepare the dataset with retry mechanism
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                dataset = load_dataset(
+                    "QuotaClimat/frugalaichallenge-text-train",
+                    token=HF_TOKEN,
+                    trust_remote_code=True
+                )
+                break
+            except Exception as e:
+                if attempt == max_retries - 1:
+                    raise Exception(f"Failed to load dataset after {max_retries} attempts: {str(e)}")
+                print(f"Dataset loading attempt {attempt + 1} failed, retrying... Error: {str(e)}")
+                time.sleep(2)
         
         # Convert string labels to integers
         dataset = dataset.map(lambda x: {"label": LABEL_MAPPING[x["label"]]})
@@ -146,8 +165,8 @@ async def evaluate_text(request: TextEvaluationRequest):
         # Initialize the model once
         classifier = TextClassifier()
 
-        # Prepare batches
-        batch_size = 24
+        # Prepare batches with smaller batch size
+        batch_size = 16  # Reduced batch size
         quotes = test_dataset["quote"]
         num_batches = len(quotes) // batch_size + (1 if len(quotes) % batch_size != 0 else 0)
         batches = [
@@ -158,8 +177,8 @@ async def evaluate_text(request: TextEvaluationRequest):
         # Initialize batch_results
         batch_results = [[] for _ in range(num_batches)]
         
-        # Process batches in parallel
-        max_workers = min(os.cpu_count(), 4)
+        # Process batches in parallel with fewer workers
+        max_workers = min(os.cpu_count(), 2)  # Reduced number of workers
         print(f"Processing with {max_workers} workers")
         
         with ThreadPoolExecutor(max_workers=max_workers) as executor:
@@ -192,6 +211,11 @@ async def evaluate_text(request: TextEvaluationRequest):
         accuracy = accuracy_score(true_labels, predictions)
         print("accuracy:", accuracy)
         
+        # Clean up
+        del classifier
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        
         # Prepare results
         results = {
             "username": username,