fix: pass tokenizer explicitly to AbuseDataset and safeguard evaluation step

train_abuse_model.py

@@ -54,7 +54,7 @@ logger.info("PyTorch version:", torch.__version__)
 # Custom Dataset class
 
 class AbuseDataset(Dataset):
-    def __init__(self, texts, labels):
+    def __init__(self, texts, labels, tokenizer):
         self.encodings = tokenizer(texts, truncation=True, padding=True, max_length=512)
         self.labels = labels
 
@@ -223,10 +223,9 @@ def run_training():
                 param.requires_grad = False
             
 
-        train_dataset = AbuseDataset(train_texts, train_labels)
-        val_dataset = AbuseDataset(val_texts, val_labels)
-        test_dataset = AbuseDataset(test_texts, test_labels)
-
+        train_dataset = AbuseDataset(train_texts, train_labels,tokenizer)
+        val_dataset = AbuseDataset(val_texts, val_labels,tokenizer)
+        test_dataset = AbuseDataset(test_texts, test_labels,tokenizer)
 
         # TrainingArguments for HuggingFace Trainer (logging, saving)
         training_args = TrainingArguments(
@@ -270,9 +269,10 @@ def run_training():
 
     # Evaluation
     try:
-        label_map = {0.0: "no", 0.5: "plausibly", 1.0: "yes"}
-        evaluate_model_with_thresholds(trainer, test_dataset)
-        logger.info("Evaluation completed")
+        if 'trainer' in locals():
+            label_map = {0.0: "no", 0.5: "plausibly", 1.0: "yes"}
+            evaluate_model_with_thresholds(trainer, test_dataset)
+            logger.info("Evaluation completed")
     except Exception as e:
         logger.exception(f"Evaluation failed: {e}")
     log_buffer.seek(0)