Update train.py

train.py

@@ -1,48 +1,46 @@
 from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
-import os
 
-# Load dataset
+# Load your dataset
 dataset = load_dataset("json", data_files="python.jsonl")
 
-# Load tokenizer and model
+# Load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
-tokenizer.pad_token = tokenizer.eos_token
 model = AutoModelForCausalLM.from_pretrained("distilgpt2")
+
+# Add pad token if needed
+tokenizer.pad_token = tokenizer.eos_token
 model.config.pad_token_id = tokenizer.pad_token_id
 
-# Tokenization function
-def tokenize_function(example):
+# Tokenize function
+def tokenize(example):
     full_text = example["prompt"] + example["completion"]
-    tokens = tokenizer(full_text, truncation=True, padding="max_length", max_length=512)
+    tokens = tokenizer(full_text, padding="max_length", truncation=True, max_length=512)
     tokens["labels"] = tokens["input_ids"].copy()
     return tokens
 
 # Tokenize dataset
-tokenized_dataset = dataset["train"].map(tokenize_function)
+tokenized_dataset = dataset["train"].map(tokenize)
 
-# Training arguments
+# Training settings
 training_args = TrainingArguments(
-    output_dir="./results",
-    per_device_train_batch_size=2,
+    output_dir="trained_model",
     num_train_epochs=1,
+    per_device_train_batch_size=2,
     logging_steps=10,
     save_strategy="epoch",
-    logging_dir="./logs",
-    report_to="none"
+    push_to_hub=False,
 )
 
-# Trainer setup
+# Trainer
 trainer = Trainer(
     model=model,
+    tokenizer=tokenizer,
     args=training_args,
     train_dataset=tokenized_dataset,
-    tokenizer=tokenizer,
 )
 
-# Start training
+# Train and save
 trainer.train()
-
-# Save final model
 trainer.save_model("trained_model")
 tokenizer.save_pretrained("trained_model")