Update train.py

train.py

@@ -1,50 +1,47 @@
 import argparse
 from datasets import load_dataset
-from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer, DataCollatorForLanguageModeling
-
-parser = argparse.ArgumentParser()
-parser.add_argument("--dataset", required=True)
-parser.add_argument("--output", default="trained_model")
-args = parser.parse_args()
-
-print("📊 Loading dataset...")
-dataset = load_dataset("json", data_files=args.dataset, split="train")
-
-print("🧠 Loading model and tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
-tokenizer.pad_token = tokenizer.eos_token
-model = AutoModelForCausalLM.from_pretrained("distilgpt2")
-
-# ✅ Clean, batch-safe tokenize
-def tokenize(batch):
-    full_texts = [str(p) + str(c) for p, c in zip(batch["prompt"], batch["completion"])]
-    return tokenizer(full_texts, padding="max_length", truncation=True, max_length=256)
-
-print("🔁 Tokenizing...")
-tokenized = dataset.map(tokenize, batched=True)
-
-print("📦 Setting up trainer...")
-data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
-
-training_args = TrainingArguments(
-    output_dir=args.output,
-    per_device_train_batch_size=2,
-    num_train_epochs=1,
-    logging_steps=1,
-    save_steps=5,
-    save_total_limit=1,
-    report_to=[]
-)
-
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=tokenized,
-    tokenizer=tokenizer,
-    data_collator=data_collator,
-)
-
-print("🚀 Starting training...")
-trainer.train()
-trainer.save_model(args.output)
-print("✅ Done.")
+from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
+import torch
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset", type=str, required=True)
+    args = parser.parse_args()
+
+    print("📥 Loading dataset...")
+    dataset = load_dataset("json", data_files=args.dataset, split="train")
+
+    tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
+    tokenizer.pad_token = tokenizer.eos_token
+
+    def tokenize_function(examples):
+        return tokenizer(examples["prompt"], truncation=True, padding="max_length", max_length=256)
+
+    tokenized_dataset = dataset.map(tokenize_function, batched=True)
+
+    print("📦 Loading model...")
+    model = AutoModelForCausalLM.from_pretrained("distilgpt2")
+
+    training_args = TrainingArguments(
+        output_dir="./trained_model",
+        overwrite_output_dir=True,
+        num_train_epochs=1,
+        per_device_train_batch_size=2,
+        save_strategy="epoch",
+        logging_dir="./logs",
+        logging_steps=10,
+        no_cuda=not torch.cuda.is_available()
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_dataset
+    )
+
+    print("🚀 Starting training...")
+    trainer.train()
+    print("✅ Training finished. Model saved to ./trained_model")
+
+if __name__ == "__main__":
+    main()