Update backend/app/train.py

backend/app/train.py

@@ -5,7 +5,7 @@ from transformers import (
     AutoModelForSequenceClassification,
     TrainingArguments,
     Trainer,
-    DataCollatorWithPadding,
+    Dad taCollatorWithPadding,
 )
 import torch
 from datasets import Dataset
@@ -38,12 +38,15 @@ def load_and_prepare_dataset():
 
 # -------- Tokenization --------
 def tokenize_function(example, tokenizer):
-    return tokenizer(
+    tokens = tokenizer(
         example["text"],
         truncation=True,
         padding=True,
         max_length=256,
     )
+    tokens["label"] = example["label"]  # ✅ Keep label after tokenization
+    return tokens
+
 
 # -------- Main Training Function --------
 def train():