scripts/upload_model.py

"""Upload the saved DistilBERT model to Hugging Face Hub.

Usage:
    huggingface-cli login
    python scripts/upload_model.py --repo zkmine/dialectica-bloom
"""

import argparse
import os


def main():
    """Push local model folder to a Hub model repo."""
    parser = argparse.ArgumentParser(description="Upload model to HF Hub")
    parser.add_argument("--repo", required=True,
                        help="target repo id, e.g. zkmine/dialectica-bloom")
    parser.add_argument("--local", default="models/distilbert_bloom",
                        help="local model directory to upload")
    parser.add_argument("--private", action="store_true",
                        help="create the repo as private")
    args = parser.parse_args()

    from huggingface_hub import HfApi, create_repo

    if not os.path.isdir(args.local):
        raise SystemExit(f"Local model dir not found: {args.local}")

    create_repo(args.repo, repo_type="model", private=args.private,
                exist_ok=True)
    api = HfApi()
    api.upload_folder(
        folder_path=args.local,
        repo_id=args.repo,
        repo_type="model",
    )
    print(f"Uploaded {args.local} -> https://huggingface.co/{args.repo}")


if __name__ == "__main__":
    main()