Update train.py

train.py

@@ -7,13 +7,11 @@ from transformers import (
 )
 import zipfile
 
-ROOT = Path(__file__).resolve().parent
-
 def parse_args():
     ap = argparse.ArgumentParser()
-    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("--dataset", required=True, help="Path to .jsonl")
+    ap.add_argument("--output",  required=True, help="Output model folder")
+    ap.add_argument("--zip_path", required=True, help="Path to write .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)
@@ -25,6 +23,7 @@ def main():
     a = parse_args()
     out_dir  = Path(a.output).resolve()
     zip_path = Path(a.zip_path).resolve()
+    out_dir.parent.mkdir(parents=True, exist_ok=True)
 
     print(f"📦 Loading dataset from: {a.dataset}", flush=True)
     ds = load_dataset("json", data_files=a.dataset, split="train")
@@ -36,13 +35,15 @@ def main():
         tok.pad_token = tok.eos_token
     model = AutoModelForCausalLM.from_pretrained(a.model_name)
 
-    def tokenize(batch):
+    def to_text(batch):
         if "text" in batch:
-            texts = batch["text"]
-        elif "prompt" in batch and "completion" in batch:
-            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 batch["text"]
+        if "prompt" in batch and "completion" in batch:
+            return [str(p).rstrip() + "\n" + str(c) for p, c in zip(batch["prompt"], batch["completion"])]
+        raise ValueError("Dataset must have 'text' or 'prompt' + 'completion'.")
+
+    def tokenize(batch):
+        texts = to_text(batch)
         return tok(texts, padding="max_length", truncation=True, max_length=a.block_size)
 
     print("🔁 Tokenizing…", flush=True)