upload_model.py

"""Upload the trained checkpoint to a public HF Model repo.

The Space's predictor.py will hf_hub_download() it on startup. This is HF's
canonical pattern for serving model weights in a Space — Spaces' LFS path is
inconsistent for direct uploads, but model repos handle large files cleanly.

Usage:
    export HF_TOKEN=hf_xxx
    ./venv/bin/python upload_model.py
"""

from __future__ import annotations

import os
import sys
from pathlib import Path

from huggingface_hub import HfApi, create_repo

PROJECT_ROOT = Path(__file__).resolve().parent
CHECKPOINT = PROJECT_ROOT / "model" / "saved" / "brain_tumor_model.pth"
REPO_ID = "momenalhamza/brain-tumor-classifier"  # model repo, not Space


def main() -> None:
    token = os.environ.get("HF_TOKEN")
    if not token:
        sys.exit(
            "HF_TOKEN not set. In your shell run:\n"
            "  export HF_TOKEN=hf_xxxxxxxxxxxxxx\n"
            "  ./venv/bin/python upload_model.py"
        )
    if not CHECKPOINT.exists():
        sys.exit(f"Checkpoint not found at {CHECKPOINT}")

    api = HfApi(token=token)

    print(f"→ Ensuring model repo exists: {REPO_ID}")
    create_repo(
        repo_id=REPO_ID,
        repo_type="model",
        token=token,
        exist_ok=True,
    )

    size_mb = CHECKPOINT.stat().st_size / 1024 / 1024
    print(f"→ Uploading {CHECKPOINT.name} ({size_mb:.1f} MB) via LFS ...")
    api.upload_file(
        path_or_fileobj=str(CHECKPOINT),
        path_in_repo="brain_tumor_model.pth",
        repo_id=REPO_ID,
        repo_type="model",
        commit_message="Upload EfficientNet-B3 brain tumor classifier (95% test acc)",
    )

    # Also push a README so the model repo has context.
    readme = (
        "---\nlicense: mit\ntags:\n- pytorch\n- image-classification\n- medical-imaging\n- "
        "brain-tumor\n- efficientnet\n---\n\n# Brain Tumor Classifier (EfficientNet-B3)\n\n"
        "PyTorch EfficientNet-B3 fine-tuned on the "
        "[Brain Tumor MRI Dataset](https://www.kaggle.com/datasets/masoudnickparvar/brain-tumor-mri-dataset) "
        "for 4-class classification: glioma, meningioma, notumor, pituitary.\n\n"
        "**Test accuracy: 95.00%** on 1,600 held-out MRI scans.\n\n"
        "Per-class F1: glioma 0.903 · meningioma 0.939 · notumor 0.966 · pituitary 0.989.\n\n"
        "Used by the [brain-tumor-classification Space]"
        f"(https://huggingface.co/spaces/{REPO_ID.split('/')[0]}/brain-tumor-classification).\n"
    )
    api.upload_file(
        path_or_fileobj=readme.encode("utf-8"),
        path_in_repo="README.md",
        repo_id=REPO_ID,
        repo_type="model",
        commit_message="Add model card",
    )

    print(
        f"\n✓ Done.\n"
        f"  Model: https://huggingface.co/{REPO_ID}\n"
        f"  The Space's predictor.py will hf_hub_download() this file on next boot."
    )


if __name__ == "__main__":
    main()