Update train.py

train.py

@@ -1,42 +1,44 @@
 from datasets import load_dataset
-from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
+from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
 
-# Load dataset from jsonl
+# Load your dataset
 dataset = load_dataset("json", data_files="python.jsonl")
 
 # Load tokenizer and model
-model_name = "distilgpt2"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-tokenizer.pad_token = tokenizer.eos_token  # Fix for padding error
-model = AutoModelForCausalLM.from_pretrained(model_name)
-
-# Tokenization and label setup for causal LM
-def preprocess_function(examples):
-    full_text = examples["prompt"] + examples["completion"]
-    model_inputs = tokenizer(full_text, truncation=True, padding="max_length", max_length=512)
-    model_inputs["labels"] = model_inputs["input_ids"].copy()  # Important for loss calculation
-    return model_inputs
-
-# Apply preprocessing
-tokenized_dataset = dataset["train"].map(preprocess_function)
-
-# Training configuration
+tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
+model = AutoModelForCausalLM.from_pretrained("distilgpt2")
+
+# Add padding token if missing
+tokenizer.pad_token = tokenizer.eos_token
+model.config.pad_token_id = tokenizer.pad_token_id
+
+# Tokenize data
+def tokenize_function(example):
+    full_text = example["prompt"] + example["completion"]
+    return tokenizer(full_text, truncation=True, padding="max_length", max_length=512)
+
+tokenized_dataset = dataset["train"].map(tokenize_function)
+
+# Training config
 training_args = TrainingArguments(
-    output_dir="trained_model",
-    learning_rate=2e-5,
-    per_device_train_batch_size=4,
-    num_train_epochs=5,
-    weight_decay=0.01,
-    save_total_limit=1,
-    logging_steps=1
+    output_dir="./results",
+    per_device_train_batch_size=2,
+    num_train_epochs=1,
+    logging_steps=10,
+    save_strategy="no",
 )
 
-# Trainer setup
+# Trainer
 trainer = Trainer(
     model=model,
     args=training_args,
-    train_dataset=tokenized_dataset
+    train_dataset=tokenized_dataset,
+    tokenizer=tokenizer,
 )
 
-# Start training
-trainer.train()
+# Train
+trainer.train()
+
+# ✅ Save model & tokenizer
+trainer.save_model("trained_model")
+tokenizer.save_pretrained("trained_model")