Create train.py

train.py

@@ -0,0 +1,75 @@
+import argparse, os, sys
+from typing import List
+from datasets import load_dataset
+from transformers import (
+    AutoTokenizer, AutoModelForCausalLM,
+    DataCollatorForLanguageModeling, TrainingArguments, Trainer
+)
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--dataset", required=True, help="JSONL (.jsonl or .jsonl.gz)")
+    p.add_argument("--output", default="trained_model")
+    p.add_argument("--model_name", default="distilgpt2")   # tiny & quick
+    p.add_argument("--epochs", type=float, default=0.5)    # short run
+    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():
+    a = parse_args()
+    print(f"📥 Loading dataset: {a.dataset}", flush=True)
+    ds = load_dataset("json", data_files=a.dataset, split="train")
+    cols = ds.column_names
+    print("🧾 Columns:", cols, flush=True)
+
+    tok = AutoTokenizer.from_pretrained(a.model_name)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+    model = AutoModelForCausalLM.from_pretrained(a.model_name)
+
+    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:
+            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'.")
+
+    def tokenize(batch):
+        texts = build_texts(batch)
+        return tok(texts, padding="max_length", truncation=True, max_length=a.block_size)
+
+    print("🔁 Tokenizing…", flush=True)
+    tokds = ds.map(tokenize, batched=True, remove_columns=cols)
+    collator = DataCollatorForLanguageModeling(tokenizer=tok, mlm=False)
+
+    print("⚙ Trainer…", flush=True)
+    args = TrainingArguments(
+        output_dir=a.output,
+        overwrite_output_dir=True,
+        per_device_train_batch_size=a.batch_size,
+        num_train_epochs=a.epochs,
+        learning_rate=a.learning_rate,
+        logging_steps=10,
+        save_steps=200,
+        save_total_limit=1,
+        report_to=[],
+        fp16=False,
+    )
+    trainer = Trainer(model=model, args=args, train_dataset=tokds, tokenizer=tok, data_collator=collator)
+
+    print("🚀 Training…", flush=True)
+    trainer.train()
+    print(f"💾 Saving to {a.output}", flush=True)
+    os.makedirs(a.output, exist_ok=True)
+    trainer.save_model(a.output)
+    tok.save_pretrained(a.output)
+    print("✅ Done.", flush=True)
+
+if _name_ == "_main_":
+    try:
+        main()
+    except Exception as e:
+        print(f"❌ Training failed: {e}", flush=True)
+        sys.exit(1)