Update train.py

train.py

@@ -1,17 +1,19 @@
 import argparse, os, traceback
+from pathlib import Path
 from datasets import load_dataset
 from transformers import (
     AutoTokenizer, AutoModelForCausalLM,
     DataCollatorForLanguageModeling, Trainer, TrainingArguments
 )
 
-DONE = "TRAIN_DONE"
-ERRF = "TRAIN_ERROR"
+ROOT = Path(_file_).resolve().parent          # /home/user/app
+DONE = ROOT / "TRAIN_DONE"                      # <- write here
+ERRF = ROOT / "TRAIN_ERROR"
 
 def parse_args():
     ap = argparse.ArgumentParser()
     ap.add_argument("--dataset", required=True)
-    ap.add_argument("--output",  default="trained_model")
+    ap.add_argument("--output",  default=str(ROOT / "trained_model"))
     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)
@@ -26,6 +28,7 @@ def main():
     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"✂ Subset: {len(ds)} rows", flush=True)
@@ -71,12 +74,14 @@ def main():
     os.makedirs(a.output, exist_ok=True)
     trainer.save_model(a.output)
     tok.save_pretrained(a.output)
-    open(DONE, "w").write("ok")      # <—— signal file
+    DONE.write_text("ok")                       # <- SIGNAL!
     print("✅ Done.", flush=True)
 
 if __name__ == "__main__":
     try:
+        DONE.unlink(missing_ok=True)
+        ERRF.unlink(missing_ok=True)
         main()
     except Exception:
-        open(ERRF, "w").write(traceback.format_exc())
+        ERRF.write_text(traceback.format_exc())
         raise