Update train.py

train.py

@@ -1,5 +1,4 @@
-# train.py
-import argparse, os, json
+import argparse, os
 from pathlib import Path
 from datasets import load_dataset
 from transformers import (
@@ -8,11 +7,11 @@ from transformers import (
 )
 import zipfile
 
-ROOT = Path(__file__).resolve().parent
+ROOT = Path(_file_).resolve().parent
 
 def parse_args():
     ap = argparse.ArgumentParser()
-    ap.add_argument("--dataset", required=True)
+    ap.add_argument("--dataset", required=True, help="Path to .jsonl (or a folder you adapt later)")
     ap.add_argument("--output",  default=str(ROOT / "trained_model"))
     ap.add_argument("--zip_path", default=str(ROOT / "trained_model.zip"))
     ap.add_argument("--model_name", default="Salesforce/codegen-350M-multi")
@@ -20,7 +19,6 @@ def parse_args():
     ap.add_argument("--batch_size", type=int, default=2)
     ap.add_argument("--block_size", type=int, default=256)
     ap.add_argument("--learning_rate", type=float, default=5e-5)
-    ap.add_argument("--subset", type=int, default=0)
     return ap.parse_args()
 
 def main():
@@ -33,10 +31,6 @@ def main():
     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"✂ Subset: {len(ds)} rows", flush=True)
-
     tok = AutoTokenizer.from_pretrained(a.model_name, use_fast=True)
     if tok.pad_token is None and tok.eos_token is not None:
         tok.pad_token = tok.eos_token
@@ -79,7 +73,7 @@ def main():
     trainer.save_model(out_dir)
     tok.save_pretrained(out_dir)
 
-    # Zip the folder ourselves (no flags, no UI dependency)
+    # Zip the folder
     if zip_path.exists():
         zip_path.unlink()
     print(f"📦 Zipping → {zip_path.name}", flush=True)