Create train.py

train.py

@@ -0,0 +1,55 @@
+from datasets import load_dataset, Dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
+import os
+
+# Load dataset
+dataset = load_dataset("json", data_files="python_train_100.jsonl")
+
+# Load tokenizer and model
+model_name = "distilgpt2"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# Tokenize function
+def tokenize_function(example):
+    full_text = f"### Prompt:\n{example['prompt']}\n### Completion:\n{example['completion']}"
+    return tokenizer(full_text, truncation=True, padding="max_length", max_length=512)
+
+tokenized_dataset = dataset["train"].map(tokenize_function)
+
+# Training arguments
+training_args = TrainingArguments(
+    output_dir="trained_model",
+    evaluation_strategy="no",
+    learning_rate=2e-5,
+    per_device_train_batch_size=4,
+    num_train_epochs=5,
+    weight_decay=0.01,
+    save_total_limit=1,
+    logging_dir="./logs",
+)
+
+# Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_dataset,
+)
+
+# Train
+trainer.train()
+
+# Save and push model to hub
+repo_name = "Percy3822/python_coder_100"
+trainer.save_model(repo_name)
+tokenizer.save_pretrained(repo_name)
+
+# Optional: push to hub
+from huggingface_hub import HfApi
+api = HfApi()
+api.upload_folder(
+    folder_path=repo_name,
+    path_in_repo="",
+    repo_id=repo_name,
+    repo_type="model"
+)