Update train.py

train.py

@@ -1,19 +1,20 @@
-import argparse, os, traceback
+# train.py
+import argparse, os, json
 from pathlib import Path
 from datasets import load_dataset
 from transformers import (
     AutoTokenizer, AutoModelForCausalLM,
     DataCollatorForLanguageModeling, Trainer, TrainingArguments
 )
+import zipfile
 
-ROOT = Path(__file__).resolve().parent          # /home/user/app
-DONE = ROOT / "TRAIN_DONE"                      # <- write here
-ERRF = ROOT / "TRAIN_ERROR"
+ROOT = Path(_file_).resolve().parent
 
 def parse_args():
     ap = argparse.ArgumentParser()
     ap.add_argument("--dataset", required=True)
     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")
     ap.add_argument("--epochs", type=float, default=1.0)
     ap.add_argument("--batch_size", type=int, default=2)
@@ -24,6 +25,9 @@ def parse_args():
 
 def main():
     a = parse_args()
+    out_dir  = Path(a.output).resolve()
+    zip_path = Path(a.zip_path).resolve()
+
     print(f"📦 Loading dataset from: {a.dataset}", flush=True)
     ds = load_dataset("json", data_files=a.dataset, split="train")
     cols = ds.column_names
@@ -52,7 +56,7 @@ def main():
     collator = DataCollatorForLanguageModeling(tokenizer=tok, mlm=False)
 
     args = TrainingArguments(
-        output_dir=a.output,
+        output_dir=str(out_dir),
         overwrite_output_dir=True,
         per_device_train_batch_size=a.batch_size,
         num_train_epochs=a.epochs,
@@ -70,18 +74,20 @@ def main():
     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)
-    DONE.write_text("ok")                       # <- SIGNAL!
+    print(f"💾 Saving to {out_dir}", flush=True)
+    os.makedirs(out_dir, exist_ok=True)
+    trainer.save_model(out_dir)
+    tok.save_pretrained(out_dir)
+
+    # Zip the folder ourselves (no flags, no UI dependency)
+    if zip_path.exists():
+        zip_path.unlink()
+    print(f"📦 Zipping → {zip_path.name}", flush=True)
+    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as z:
+        for p in out_dir.rglob("*"):
+            z.write(p, arcname=p.relative_to(out_dir))
+
     print("✅ Done.", flush=True)
 
 if __name__ == "__main__":
-    try:
-        DONE.unlink(missing_ok=True)
-        ERRF.unlink(missing_ok=True)
-        main()
-    except Exception:
-        ERRF.write_text(traceback.format_exc())
-        raise
+    main()