Update train.py

train.py

@@ -1,37 +1,38 @@
 from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
+import os
 
-# Load your dataset
+# Load dataset
 dataset = load_dataset("json", data_files="python.jsonl")
 
 # Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
-model = AutoModelForCausalLM.from_pretrained("distilgpt2")
-
-# Add padding token if missing
 tokenizer.pad_token = tokenizer.eos_token
+model = AutoModelForCausalLM.from_pretrained("distilgpt2")
 model.config.pad_token_id = tokenizer.pad_token_id
 
-# Tokenize function
+# Tokenization 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()
     return tokens
 
-# Tokenize
+# Tokenize dataset
 tokenized_dataset = dataset["train"].map(tokenize_function)
 
-# Training configuration
+# Training arguments
 training_args = TrainingArguments(
     output_dir="./results",
     per_device_train_batch_size=2,
     num_train_epochs=1,
     logging_steps=10,
-    save_strategy="no",
+    save_strategy="epoch",
+    logging_dir="./logs",
+    report_to="none"
 )
 
-# Trainer
+# Trainer setup
 trainer = Trainer(
     model=model,
     args=training_args,
@@ -39,9 +40,9 @@ trainer = Trainer(
     tokenizer=tokenizer,
 )
 
-# Train
+# Start training
 trainer.train()
 
-# Save model and tokenizer
+# Save final model
 trainer.save_model("trained_model")
 tokenizer.save_pretrained("trained_model")