improve model loading

tasks/text.py

@@ -39,10 +39,12 @@ class TextClassifier:
             # Initialize tokenizer
             self.tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
             
-            # Initialize model
-            self.model = BertForSequenceClassification.from_pretrained(
+            # Initialize model with auto class
+            self.model = AutoModelForSequenceClassification.from_pretrained(
                 MODEL_NAME,
+                trust_remote_code=True,
                 num_labels=8,
+                problem_type="single_label_classification",
                 ignore_mismatched_sizes=True
             ).to(self.device)
             
@@ -117,19 +119,27 @@ async def evaluate_text(request: TextEvaluationRequest):
                 examples["quote"],
                 truncation=True,
                 padding=True,
-                max_length=512
+                max_length=512,
+                return_tensors=None  # Changed this to None for batched processing
             )
         
         # Tokenize dataset
-        tokenized_test = test_dataset.map(preprocess_function, batched=True)
-        tokenized_test.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
+        tokenized_test = test_dataset.map(
+            preprocess_function, 
+            batched=True,
+            remove_columns=test_dataset.column_names
+        )
+        
+        # Set format for pytorch
+        tokenized_test.set_format("torch")
         
         # Create DataLoader
         data_collator = DataCollatorWithPadding(tokenizer=classifier.tokenizer)
         test_loader = DataLoader(
             tokenized_test,
             batch_size=16,
-            collate_fn=data_collator
+            collate_fn=data_collator,
+            shuffle=False
         )
         
         # Get predictions