Update train.py

train.py

@@ -1,24 +1,26 @@
-from datasets import load_dataset, Dataset
-from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
-import os
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
 
-# Load dataset
-dataset = load_dataset("json", data_files="python.jsonl")
+# Load dataset from jsonl
+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
+tokenizer.pad_token = tokenizer.eos_token  # Fix for padding error
 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)
+# 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
 
-tokenized_dataset = dataset["train"].map(tokenize_function)
+# Apply preprocessing
+tokenized_dataset = dataset["train"].map(preprocess_function)
 
-# Training arguments
+# Training configuration
 training_args = TrainingArguments(
     output_dir="trained_model",
     learning_rate=2e-5,
@@ -29,27 +31,12 @@ training_args = TrainingArguments(
     logging_steps=1
 )
 
-# Trainer
+# Trainer setup
 trainer = Trainer(
     model=model,
     args=training_args,
-    train_dataset=tokenized_dataset,
+    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"
-)
+# Start training
+trainer.train()