Update train.py

train.py

@@ -34,16 +34,22 @@ tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
-# ✅ Tokenize Data
+# ✅ Tokenize Data (Fixed)
 def preprocess_function(examples):
     prompt = examples.get("prompt", "")
     response = examples.get("response", "")
     inputs = f"Medical Q&A: {prompt} {response}"
+    
     model_inputs = tokenizer(inputs, padding="max_length", truncation=True, max_length=512)
-    model_inputs["labels"] = model_inputs["input_ids"].copy()
-    return model_inputs
 
-tokenized_dataset = dataset.map(preprocess_function, batched=True)
+    # ✅ Ensure labels have the same length
+    model_inputs["labels"] = model_inputs["input_ids"]
+
+    return {key: [val] for key, val in model_inputs.items()}  # ✅ Wrap values in lists
+
+# ✅ Apply tokenization
+tokenized_dataset = dataset.map(preprocess_function, batched=True, remove_columns=dataset.column_names)
+
 
 # ✅ Load Model with LoRA (Optimized for Falcon)
 model = AutoModelForCausalLM.from_pretrained(