Update train.py

train.py

@@ -8,22 +8,37 @@ from transformers import (
 
 def parse_args():
     p = argparse.ArgumentParser()
-    p.add_argument("--dataset", required=True, help="JSONL (.jsonl or .jsonl.gz)")
+    p.add_argument("--dataset", required=True, help="JSON/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("--model_name", default="distilgpt2")
+    p.add_argument("--epochs", type=float, default=0.5)
     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)
+    # quick mode:
+    p.add_argument("--quick", type=int, default=0)        # 1 => tiny model + fast
+    p.add_argument("--max_steps", type=int, default=0)    # >0 overrides epochs
+    p.add_argument("--subset", type=int, default=0)       # use first N rows
     return p.parse_args()
 
 def main():
     a = parse_args()
+
+    if a.quick:
+        a.model_name = "sshleifer/tiny-gpt2"  # ultra-tiny, very fast
+        if a.max_steps <= 0: a.max_steps = 8
+        if a.subset <= 0: a.subset = 32
+        a.epochs = 1.0
+
     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)
 
+    if a.subset and a.subset > 0:
+        ds = ds.select(range(min(a.subset, len(ds))))
+        print(f"✂ Using subset: {len(ds)} rows", flush=True)
+
     tok = AutoTokenizer.from_pretrained(a.model_name)
     if tok.pad_token is None:
         tok.pad_token = tok.eos_token
@@ -33,7 +48,7 @@ def main():
         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"])]
+            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):
@@ -49,13 +64,14 @@ def main():
         output_dir=a.output,
         overwrite_output_dir=True,
         per_device_train_batch_size=a.batch_size,
-        num_train_epochs=a.epochs,
+        num_train_epochs=a.epochs if a.max_steps == 0 else 1,
         learning_rate=a.learning_rate,
-        logging_steps=10,
-        save_steps=200,
+        logging_steps=1,
+        save_steps=50,
         save_total_limit=1,
         report_to=[],
         fp16=False,
+        max_steps=a.max_steps if a.max_steps > 0 else -1,
     )
     trainer = Trainer(model=model, args=args, train_dataset=tokds, tokenizer=tok, data_collator=collator)
 
@@ -67,7 +83,7 @@ def main():
     tok.save_pretrained(a.output)
     print("✅ Done.", flush=True)
 
-if __name__ == "__main__":
+if _name_ == "_main_":
     try:
         main()
     except Exception as e: