fix model initialization with explicit loading

tasks/text.py

@@ -34,57 +34,67 @@ TOKENIZER_NAME = "answerdotai/ModernBERT-base"
 class TextClassifier:
     def __init__(self):
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        max_retries = 3
         
-        for attempt in range(max_retries):
-            try:
-                # Initialize tokenizer
-                self.tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
-                
-                # Initialize model
-                self.model = AutoModelForSequenceClassification.from_pretrained(
-                    MODEL_NAME,
-                    num_labels=8,
-                    trust_remote_code=True
-                ).to(self.device)
-                
-                # Convert to half precision
-                self.model = self.model.half()
-                self.model.eval()
-                
-                print("Model initialized successfully")
-                break
+        try:
+            # Initialize tokenizer
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                TOKENIZER_NAME,
+                model_max_length=8192,
+                padding_side='right',
+                truncation_side='right'
+            )
+            
+            # Load model configuration
+            model_config = {
+                "architectures": ["ModernBertForSequenceClassification"],
+                "model_type": "modernbert",
+                "num_labels": 8,
+                "problem_type": "single_label_classification",
+                "hidden_size": 768,
+                "num_attention_heads": 12,
+                "num_hidden_layers": 22,
+                "intermediate_size": 1152,
+                "max_position_embeddings": 8192,
+                "torch_dtype": "float32",
+                "transformers_version": "4.48.3",
+                "layer_norm_eps": 1e-05
+            }
+            
+            # Initialize model
+            self.model = AutoModelForSequenceClassification.from_pretrained(
+                MODEL_NAME,
+                config=model_config,
+                ignore_mismatched_sizes=True,
+                trust_remote_code=True
+            ).to(self.device)
+            
+            # Convert to half precision
+            self.model = self.model.half()
+            self.model.eval()
+            
+            print("Model initialized successfully")
                 
-            except Exception as e:
-                if attempt == max_retries - 1:
-                    raise Exception(f"Failed to initialize model after {max_retries} attempts: {str(e)}")
-                print(f"Attempt {attempt + 1} failed, retrying... Error: {str(e)}")
-                time.sleep(1)
+        except Exception as e:
+            print(f"Error initializing model: {str(e)}")
+            raise
 
-    def process_batch(self, texts: List[str]) -> List[int]:
+    def process_batch(self, batch):
         """Process a batch of texts and return their predictions"""
         try:
-            # Tokenize
-            inputs = self.tokenizer(
-                texts,
-                padding=True,
-                truncation=True,
-                return_tensors="pt"
-            )
-            
-            # Move inputs to device
-            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+            # Move batch to device
+            input_ids = batch['input_ids'].to(self.device)
+            attention_mask = batch['attention_mask'].to(self.device)
             
             # Get predictions
             with torch.no_grad():
-                outputs = self.model(**inputs)
+                outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
                 predictions = torch.argmax(outputs.logits, dim=-1)
                 
             return predictions.cpu().numpy().tolist()
             
         except Exception as e:
             print(f"Error in batch processing: {str(e)}")
-            return [0] * len(texts)
+            return [0] * len(batch['input_ids'])
 
     def __del__(self):
         if hasattr(self, 'model'):
@@ -121,35 +131,41 @@ async def evaluate_text(request: TextEvaluationRequest):
         tracker.start()
         tracker.start_task("inference")
 
-        # Get true labels
-        true_labels = test_dataset["label"]
-        
         # Initialize model
         classifier = TextClassifier()
         
-        # Process in batches
-        batch_size = 16
-        data_collator = DataCollatorWithPadding(tokenizer=classifier.tokenizer)
+        # Prepare tokenization function
+        def preprocess_function(examples):
+            return classifier.tokenizer(
+                examples["quote"],
+                truncation=True,
+                padding=True,
+                max_length=512
+            )
+        
+        # Tokenize dataset
+        tokenized_test = test_dataset.map(preprocess_function, batched=True)
+        tokenized_test.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
         
         # Create DataLoader
+        data_collator = DataCollatorWithPadding(tokenizer=classifier.tokenizer)
         test_loader = DataLoader(
-            test_dataset,
-            batch_size=batch_size,
+            tokenized_test,
+            batch_size=16,
             collate_fn=data_collator
         )
         
         # Get predictions
         all_predictions = []
         for batch in test_loader:
-            batch_texts = batch["quote"]
-            batch_preds = classifier.process_batch(batch_texts)
+            batch_preds = classifier.process_batch(batch)
             all_predictions.extend(batch_preds)
         
         # Stop tracking emissions
         emissions_data = tracker.stop_task()
         
         # Calculate accuracy
-        accuracy = accuracy_score(true_labels, all_predictions)
+        accuracy = accuracy_score(test_dataset["label"], all_predictions)
         
         # Prepare results
         results = {