Update train.py

train.py

@@ -12,14 +12,17 @@ model = AutoModelForCausalLM.from_pretrained("distilgpt2")
 tokenizer.pad_token = tokenizer.eos_token
 model.config.pad_token_id = tokenizer.pad_token_id
 
-# Tokenize data
+# Tokenize function: provide input_ids + labels (needed for loss)
 def tokenize_function(example):
     full_text = example["prompt"] + example["completion"]
-    return tokenizer(full_text, truncation=True, padding="max_length", max_length=512)
+    tokens = tokenizer(full_text, truncation=True, padding="max_length", max_length=512)
+    tokens["labels"] = tokens["input_ids"].copy()  # 👈 labels = input_ids for language modeling
+    return tokens
 
+# Tokenize
 tokenized_dataset = dataset["train"].map(tokenize_function)
 
-# Training config
+# Training configuration
 training_args = TrainingArguments(
     output_dir="./results",
     per_device_train_batch_size=2,
@@ -39,6 +42,6 @@ trainer = Trainer(
 # Train
 trainer.train()
 
-# ✅ Save model & tokenizer
+# Save model and tokenizer
 trainer.save_model("trained_model")
 tokenizer.save_pretrained("trained_model")