Update train.py

train.py

@@ -1,53 +1,45 @@
-import argparse, os
+import argparse, os, traceback
 from datasets import load_dataset
 from transformers import (
     AutoTokenizer, AutoModelForCausalLM,
     DataCollatorForLanguageModeling, Trainer, TrainingArguments
 )
 
+DONE = "TRAIN_DONE"
+ERRF = "TRAIN_ERROR"
+
 def parse_args():
     ap = argparse.ArgumentParser()
-    ap.add_argument("--dataset", required=True, help="JSON/JSONL file (or folder with shards)")
+    ap.add_argument("--dataset", required=True)
     ap.add_argument("--output",  default="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)
     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, help="Use first N rows for quick runs")
+    ap.add_argument("--subset", type=int, default=0)
     return ap.parse_args()
 
 def main():
     a = parse_args()
     print(f"📦 Loading dataset from: {a.dataset}", flush=True)
-
-    if os.path.isdir(a.dataset):
-        pattern = os.path.join(a.dataset, "/.jsonl")
-        ds = load_dataset("json", data_files=pattern, split="train")
-    else:
-        ds = load_dataset("json", data_files=a.dataset, split="train")
-
+    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)
 
-    print(f"🧠 Loading model: {a.model_name}", 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
     model = AutoModelForCausalLM.from_pretrained(a.model_name)
 
-    # ✅ batched=True passes dict-of-lists
     def tokenize(batch):
         if "text" in batch:
             texts = batch["text"]
         elif "prompt" in batch and "completion" in batch:
-            prompts = batch["prompt"]
-            completions = batch["completion"]
-            texts = [(str(p).rstrip() + "\n" + str(c)) for p, c in zip(prompts, completions)]
+            texts = [str(p).rstrip() + "\n" + str(c) for p, c in zip(batch["prompt"], batch["completion"])]
         else:
             raise ValueError("Dataset must have 'text' or 'prompt' + 'completion'.")
         return tok(texts, padding="max_length", truncation=True, max_length=a.block_size)
@@ -63,8 +55,7 @@ def main():
         num_train_epochs=a.epochs,
         learning_rate=a.learning_rate,
         logging_steps=5,
-        save_steps=200,
-        save_total_limit=1,
+        save_strategy="no",
         report_to=[],
         fp16=False,
     )
@@ -80,7 +71,12 @@ 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
     print("✅ Done.", flush=True)
 
-if __name__ == "__main__":
-    main()
+if __name___ == "__main__":
+    try:
+        main()
+    except Exception:
+        open(ERRF, "w").write(traceback.format_exc())
+        raise