Upload push_model_to_hf.py with huggingface_hub

push_model_to_hf.py

@@ -0,0 +1,145 @@
+#!/usr/bin/env python3
+"""
+Push trained LUNA model to HuggingFace model repo.
+
+Target: https://huggingface.co/ASTERIZER/LUNA-100M
+
+Uploads:
+  - Model weights (lit_model.pth, latest.pt)
+  - Tokenizer files
+  - Training config
+  - Quantized GGUF files (if present)
+
+Usage:
+    HF_TOKEN=hf_xxx python push_model_to_hf.py
+    HF_TOKEN=hf_xxx python push_model_to_hf.py --model_dir out/pretrain/luna-100m-english-1b
+    HF_TOKEN=hf_xxx python push_model_to_hf.py --include_gguf
+"""
+
+import argparse
+import os
+from pathlib import Path
+
+from huggingface_hub import HfApi
+
+
+MODEL_REPO = "ASTERIZER/LUNA-100M"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Push LUNA model to HuggingFace")
+    parser.add_argument("--repo_id", default=MODEL_REPO)
+    parser.add_argument("--model_dir", default="out/pretrain/luna-100m-english-1b",
+                        help="Directory containing trained model")
+    parser.add_argument("--path_in_repo", default="english_1b_continued",
+                        help="Subfolder in HF repo")
+    parser.add_argument("--include_gguf", action="store_true",
+                        help="Also upload GGUF quantisations")
+    parser.add_argument("--include_tokenizer", action="store_true", default=True,
+                        help="Upload tokenizer files")
+    parser.add_argument("--private", action="store_true")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        raise RuntimeError("Set HF_TOKEN environment variable")
+
+    api = HfApi(token=token)
+
+    # Create model repo
+    api.create_repo(
+        repo_id=args.repo_id,
+        repo_type="model",
+        private=args.private,
+        exist_ok=True,
+    )
+    print(f"Model repo: https://huggingface.co/{args.repo_id}\n")
+
+    total = 0
+    model_dir = Path(args.model_dir)
+
+    # ── 1. Model weights ──────────────────────────────────────────────────
+    print("── Model weights ──")
+    if model_dir.exists():
+        # Upload the full model directory
+        api.upload_folder(
+            repo_id=args.repo_id,
+            repo_type="model",
+            folder_path=str(model_dir),
+            path_in_repo=args.path_in_repo,
+        )
+        file_count = len(list(model_dir.rglob("*")))
+        print(f"  Uploaded {file_count} files from {model_dir}")
+        total += file_count
+    else:
+        print(f"  SKIP: {model_dir} not found")
+        # Try common alternative paths
+        for alt in [
+            "Base/out/pretrain/luna_100m/final",
+            "Base/out/pretrain/custom-100m-english/final_raw",
+        ]:
+            alt_path = Path(alt)
+            if alt_path.exists():
+                print(f"  Found alternative: {alt_path}")
+                api.upload_folder(
+                    repo_id=args.repo_id,
+                    repo_type="model",
+                    folder_path=str(alt_path),
+                    path_in_repo="pretrained",
+                )
+                total += len(list(alt_path.rglob("*")))
+                break
+
+    # ── 2. Tokenizer ──────────────────────────────────────────────────────
+    if args.include_tokenizer:
+        print("\n── Tokenizer ──")
+        tok_dir = Path("Base/checkpoints/EleutherAI/pythia-160m")
+        tok_files = ["config.json", "tokenizer_config.json", "tokenizer.json"]
+        for tf in tok_files:
+            fpath = tok_dir / tf
+            if fpath.exists():
+                api.upload_file(
+                    path_or_fileobj=str(fpath),
+                    path_in_repo=f"tokenizer/{tf}",
+                    repo_id=args.repo_id,
+                    repo_type="model",
+                )
+                print(f"  OK: {fpath}")
+                total += 1
+
+    # ── 3. Training config ────────────────────────────────────────────────
+    print("\n── Config ──")
+    for cfg in ["train_continue_english_1b.yaml", "train_config.yaml"]:
+        if os.path.exists(cfg):
+            api.upload_file(
+                path_or_fileobj=cfg,
+                path_in_repo=f"config/{cfg}",
+                repo_id=args.repo_id,
+                repo_type="model",
+            )
+            print(f"  OK: {cfg}")
+            total += 1
+
+    # ── 4. GGUF quantisations ────────────────────────────────────────────
+    if args.include_gguf:
+        print("\n── GGUF quantisations ──")
+        gguf_dir = Path("quantisations")
+        if gguf_dir.exists():
+            for gf in gguf_dir.glob("*.gguf"):
+                api.upload_file(
+                    path_or_fileobj=str(gf),
+                    path_in_repo=f"gguf/{gf.name}",
+                    repo_id=args.repo_id,
+                    repo_type="model",
+                )
+                print(f"  OK: {gf}")
+                total += 1
+
+    print(f"\nDone! Uploaded {total} items to https://huggingface.co/{args.repo_id}")
+
+
+if __name__ == "__main__":
+    main()