Update train.py

train.py

@@ -1,46 +1,48 @@
 from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
+import os
+import shutil
 
-# Load your dataset
+# Load dataset
 dataset = load_dataset("json", data_files="python.jsonl")
 
-# Load model and tokenizer
+# Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
 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
 
-# Tokenize function
+# Tokenization
 def tokenize(example):
     full_text = example["prompt"] + example["completion"]
-    tokens = tokenizer(full_text, padding="max_length", truncation=True, max_length=512)
+    tokens = tokenizer(full_text, truncation=True, padding="max_length", max_length=512)
     tokens["labels"] = tokens["input_ids"].copy()
     return tokens
 
-# Tokenize dataset
-tokenized_dataset = dataset["train"].map(tokenize)
+tokenized = dataset["train"].map(tokenize)
 
-# Training settings
-training_args = TrainingArguments(
-    output_dir="trained_model",
-    num_train_epochs=1,
+# Training
+args = TrainingArguments(
+    output_dir="./results",
     per_device_train_batch_size=2,
+    num_train_epochs=1,
     logging_steps=10,
-    save_strategy="epoch",
-    push_to_hub=False,
+    save_strategy="no",
 )
 
-# Trainer
 trainer = Trainer(
     model=model,
+    args=args,
+    train_dataset=tokenized,
     tokenizer=tokenizer,
-    args=training_args,
-    train_dataset=tokenized_dataset,
 )
 
-# Train and save
 trainer.train()
+
+# Save model
 trainer.save_model("trained_model")
-tokenizer.save_pretrained("trained_model")
+tokenizer.save_pretrained("trained_model")
+
+# Zip the model
+shutil.make_archive("trained_model", 'zip', "trained_model")
+print("✅ Training complete. Model zipped to 'trained_model.zip'")