Put train.py

train.py

@@ -17,6 +17,7 @@ import re
 
 import torch
 from datasets import load_dataset
+from huggingface_hub import HfApi, create_repo
 from peft import LoraConfig
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from trl import SFTConfig, SFTTrainer
@@ -65,23 +66,20 @@ def main():
     ap.add_argument("--grad_accum", type=int, default=4)
     ap.add_argument("--lr", type=float, default=5e-5)
     ap.add_argument("--gc", action="store_true", help="enable gradient checkpointing")
-    ap.add_argument("--report_lengths", action="store_true")
     ap.add_argument("--output", default="victor/functiongemma-270m-agent-sft-lora")
     args = ap.parse_args()
 
+    tokenizer = AutoTokenizer.from_pretrained(BASE, trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    print("pad_token:", tokenizer.pad_token, "| pad_id:", tokenizer.pad_token_id)
+
     if not args.smoke:
         # Fail fast if the push token is missing/invalid.
-        from huggingface_hub import HfApi, create_repo
-
         HfApi().whoami(token=os.environ["HF_TOKEN"])
         create_repo(args.output, token=os.environ["HF_TOKEN"], exist_ok=True)
         print("auth OK; output repo ready:", args.output)
 
-    tokenizer = AutoTokenizer.from_pretrained(BASE, trust_remote_code=True)
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-    print("pad_token:", tokenizer.pad_token, "| pad_id:", tokenizer.pad_token_id)
-
     ds = load_dataset("victor/functiongemma-agent-sft", split="train")
     print("total rows:", len(ds))
     if args.smoke:
@@ -152,10 +150,21 @@ def main():
         print("SMOKE OK")
         return
 
-    trainer.save_model("./fg-final")
-    model.push_to_hub(args.output, token=os.environ["HF_TOKEN"])
-    tokenizer.push_to_hub(args.output, token=os.environ["HF_TOKEN"])
-    print("Pushed LoRA adapter to", args.output)
+    token = os.environ["HF_TOKEN"]
+    # Remove files from the earlier buggy run that pushed the UNTRAINED base
+    # model into the output repo, so the repo holds only the real adapter.
+    for f in ["model.safetensors", "config.json", "generation_config.json", "README.md"]:
+        try:
+            HfApi().delete_file(path_in_repo=f, repo_id=args.output, token=token)
+            print("deleted stale file:", f)
+        except Exception as e:
+            print("skip delete", f, "->", e)
+
+    # Push the TRAINED LoRA adapter (trainer.model is the PEFT-wrapped model).
+    trainer.model.save_pretrained("./fg-adapter", safe_serialization=True)
+    trainer.model.push_to_hub(args.output, token=token)
+    tokenizer.push_to_hub(args.output, token=token)
+    print("Pushed trained LoRA adapter to", args.output)
 
 
 if __name__ == "__main__":