add custom layernorm patch

tasks/text.py

@@ -7,6 +7,7 @@ import os
 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 huggingface_hub import login
 from dotenv import load_dotenv
@@ -30,6 +31,11 @@ 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"
@@ -38,27 +44,31 @@ class TextClassifier:
         
         for attempt in range(max_retries):
             try:
-                # Load config and modify it
+                # Load config
                 self.config = AutoConfig.from_pretrained(model_name)
                 
-                # Remove problematic bias parameters
-                if hasattr(self.config, 'norm_bias'):
-                    delattr(self.config, 'norm_bias')
-                
                 # Initialize tokenizer
                 self.tokenizer = AutoTokenizer.from_pretrained(
                     model_name,
-                    model_max_length=2048,
+                    model_max_length=512,
                     padding_side='right',
                     truncation_side='right'
                 )
                 
-                # Initialize model with modified config
-                self.model = AutoModelForSequenceClassification.from_pretrained(
-                    model_name,
-                    config=self.config,
-                    ignore_mismatched_sizes=True
-                )
+                # Patch LayerNorm
+                original_layernorm = nn.LayerNorm
+                nn.LayerNorm = CustomLayerNorm
+                
+                try:
+                    # Initialize model with patched LayerNorm
+                    self.model = AutoModelForSequenceClassification.from_pretrained(
+                        model_name,
+                        config=self.config,
+                        ignore_mismatched_sizes=True
+                    )
+                finally:
+                    # Restore original LayerNorm
+                    nn.LayerNorm = original_layernorm
                 
                 self.model.to(self.device)
                 self.model.eval()
@@ -81,7 +91,7 @@ class TextClassifier:
                 batch,
                 return_tensors="pt",
                 truncation=True,
-                max_length=2024,
+                max_length=512,
                 padding='max_length'
             ).to(self.device)