Update train.py

train.py

@@ -12,11 +12,11 @@ model = AutoModelForCausalLM.from_pretrained("distilgpt2")
 tokenizer.pad_token = tokenizer.eos_token
 model.config.pad_token_id = tokenizer.pad_token_id
 
-# Tokenize function: provide input_ids + labels (needed for loss)
+# Tokenize function
 def tokenize_function(example):
     full_text = example["prompt"] + example["completion"]
     tokens = tokenizer(full_text, truncation=True, padding="max_length", max_length=512)
-    tokens["labels"] = tokens["input_ids"].copy()  # 👈 labels = input_ids for language modeling
+    tokens["labels"] = tokens["input_ids"].copy()
     return tokens
 
 # Tokenize