Update train.py

train.py

@@ -1,47 +1,105 @@
 import argparse
+import os
+import sys
+from typing import List
+
 from datasets import load_dataset
-from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
-import torch
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    DataCollatorForLanguageModeling,
+    TrainingArguments,
+    Trainer,
+)
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--dataset", required=True, help="Path to a JSON/JSONL (or .gz) file.")
+    p.add_argument("--output", default="trained_model", help="Folder to save the fine-tuned model.")
+    p.add_argument("--model_name", default="distilgpt2", help="Base model.")
+    p.add_argument("--epochs", type=float, default=1.0)
+    p.add_argument("--batch_size", type=int, default=2)
+    p.add_argument("--block_size", type=int, default=256)
+    p.add_argument("--learning_rate", type=float, default=5e-5)
+    return p.parse_args()
 
 def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--dataset", type=str, required=True)
-    args = parser.parse_args()
+    args = parse_args()
+    print(f"📥 Loading dataset: {args.dataset}", flush=True)
+    ds = load_dataset("json", data_files=args.dataset, split="train")
+
+    cols = ds.column_names
+    print(f"🧾 Columns: {cols}", flush=True)
 
-    print("📥 Loading dataset...")
-    dataset = load_dataset("json", data_files=args.dataset, split="train")
+    print(f"🧠 Loading model & tokenizer: {args.model_name}", flush=True)
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token  # GPT-2 family has no pad token
 
-    tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
-    tokenizer.pad_token = tokenizer.eos_token
+    model = AutoModelForCausalLM.from_pretrained(args.model_name)
 
-    def tokenize_function(examples):
-        return tokenizer(examples["prompt"], truncation=True, padding="max_length", max_length=256)
+    def build_texts(batch) -> List[str]:
+        if "text" in batch:
+            return [str(t) for t in batch["text"]]
+        if "prompt" in batch and "completion" in batch:
+            # simple join: prompt + newline + completion
+            return [f"{str(p).rstrip()}\n{str(c)}" for p, c in zip(batch["prompt"], batch["completion"])]
+        raise ValueError("Dataset must contain 'text' OR both 'prompt' and 'completion' fields.")
 
-    tokenized_dataset = dataset.map(tokenize_function, batched=True)
+    def tokenize(batch):
+        texts = build_texts(batch)
+        return tokenizer(
+            texts,
+            padding="max_length",
+            truncation=True,
+            max_length=args.block_size,
+        )
 
-    print("📦 Loading model...")
-    model = AutoModelForCausalLM.from_pretrained("distilgpt2")
+    print("🔁 Tokenizing…", flush=True)
+    tokenized = ds.map(
+        tokenize,
+        batched=True,
+        remove_columns=cols,  # keep only tokenized fields
+    )
+
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 
+    print("⚙ Preparing Trainer…", flush=True)
     training_args = TrainingArguments(
-        output_dir="./trained_model",
+        output_dir=args.output,
         overwrite_output_dir=True,
-        num_train_epochs=1,
-        per_device_train_batch_size=2,
-        save_strategy="epoch",
-        logging_dir="./logs",
+        per_device_train_batch_size=args.batch_size,
+        num_train_epochs=args.epochs,
+        learning_rate=args.learning_rate,
         logging_steps=10,
-        no_cuda=not torch.cuda.is_available()
+        save_steps=200,           # frequent-ish checkpoints (kept to 1)
+        save_total_limit=1,
+        report_to=[],
+        gradient_accumulation_steps=1,
+        fp16=False,               # CPU-friendly; enable if GPU has fp16
     )
 
     trainer = Trainer(
         model=model,
         args=training_args,
-        train_dataset=tokenized_dataset
+        train_dataset=tokenized,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
     )
 
-    print("🚀 Starting training...")
+    print("🚀 Training…", flush=True)
     trainer.train()
-    print("✅ Training finished. Model saved to ./trained_model")
 
-if __name__ == "__main__":
-    main()
+    print(f"💾 Saving to: {args.output}", flush=True)
+    os.makedirs(args.output, exist_ok=True)
+    trainer.save_model(args.output)
+    tokenizer.save_pretrained(args.output)
+    print("✅ Done.", flush=True)
+
+if _name_ == "_main_":
+    try:
+        main()
+    except Exception as e:
+        # Make sure a failure returns non-zero so your app can detect it
+        print(f"❌ Training failed: {e}", flush=True)
+        sys.exit(1)