presentation/scripts/upload_to_hf.py

"""Publish the exported classifier and the Gradio demo to Hugging Face (private model repo + private Space).

NATHAN RUNS THIS HIMSELF, AFTER EXPLICIT APPROVAL. No agent, hook or workflow may invoke it with
``--yes``. The default is a dry run that prints the exact plan (files, sizes, repos, hardware,
visibility, secret and variable names) and performs no network call at all; ``--check-auth`` adds a
single read-only ``whoami`` to the dry run. Anything that writes to the Hub requires ``--yes``.

What a live run does, in order (huggingface_hub 1.33.0 API, verified in docs/research):

1. ``HfApi().whoami()`` -- the token comes from ``HF_TOKEN`` or the ``hf auth login`` cache and is never
   printed; only the account name is logged.
2. ``create_repo(model_repo, repo_type="model", private=True, exist_ok=True)``.
3. ``upload_folder(model_dir -> model_repo, allow_patterns=[model.safetensors, config.json,
   tokenizer.json, README.md])``: the bf16 safetensors export, its config, the tokenizer and the model card.
4. ``create_repo(space_repo, repo_type="space", space_sdk="gradio", space_hardware=<--hardware>,
   private=True, exist_ok=True)``.
5. ``add_space_secret(space_repo, "HF_TOKEN", $HF_TOKEN_FOR_SPACE)``: a READ token the Space uses to download
   the private weights at startup (``preload_from_hub`` cannot read private repos). The value must come from
   the environment variable ``HF_TOKEN_FOR_SPACE``; the script never prompts and never prints it. Use
   ``--skip-secret`` if the secret is already set.
6. ``add_space_variable(space_repo, "PHISH_MODEL_REPO", model_repo)`` so app.py knows where to download from.
7. ``upload_folder(app_dir -> space_repo)``, by default ignoring ``*.safetensors``, ``tokenizer.json``,
   ``config.json``, ``__pycache__/*`` and ``*.pyc`` so the Space fetches weights from the model repo at
   runtime; ``--bundle-weights`` uploads them into the Space instead (then no secret is needed).
8. Print both URLs and ``get_space_runtime(space_repo).hardware`` to confirm the ZeroGPU flavour stuck.

Example (dry run, the default):   .venv/bin/python scripts/upload_to_hf.py --model-dir model_export/<run>
Example (live, after approval):   HF_TOKEN_FOR_SPACE=hf_... .venv/bin/python scripts/upload_to_hf.py \\
                                      --model-dir model_export/<run> --yes
"""

from __future__ import annotations

import argparse
import logging
import os
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Sequence

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))

import common  # noqa: E402

LOG = logging.getLogger("upload_to_hf")

MODEL_ALLOW_PATTERNS: list[str] = ["model.safetensors", "config.json", "tokenizer.json", "README.md"]
SPACE_IGNORE_PATTERNS: list[str] = ["*.safetensors", "tokenizer.json", "config.json", "__pycache__/*", "*.pyc"]
SPACE_REQUIRED_FILES: tuple[str, ...] = ("app.py", "model.py", "requirements.txt", "README.md")
SECRET_ENV = "HF_TOKEN_FOR_SPACE"


@dataclass
class Plan:
    model_dir: Path
    app_dir: Path
    model_repo: str
    space_repo: str
    hardware: str
    private: bool
    bundle_weights: bool
    skip_secret: bool
    model_files: list[tuple[str, int]] = field(default_factory=list)
    space_files: list[tuple[str, int]] = field(default_factory=list)
    problems: list[str] = field(default_factory=list)

    @property
    def space_ignore(self) -> list[str]:
        return ["__pycache__/*", "*.pyc"] if self.bundle_weights else list(SPACE_IGNORE_PATTERNS)


def _rel_files(folder: Path) -> list[str]:
    return sorted(str(p.relative_to(folder)) for p in folder.rglob("*") if p.is_file() and ".git" not in p.parts)


def _select(folder: Path, allow: Sequence[str] | None, ignore: Sequence[str] | None) -> list[tuple[str, int]]:
    """Apply the same include/exclude logic ``upload_folder`` uses, so the dry run lists exactly what would ship."""
    from huggingface_hub.utils import filter_repo_objects

    chosen = filter_repo_objects(_rel_files(folder), allow_patterns=list(allow) if allow else None,
                                 ignore_patterns=list(ignore) if ignore else None)
    return [(rel, (folder / rel).stat().st_size) for rel in chosen]


def build_plan(args: argparse.Namespace) -> Plan:
    plan = Plan(
        model_dir=Path(args.model_dir).resolve(), app_dir=Path(args.app_dir).resolve(),
        model_repo=args.model_repo, space_repo=args.space_repo, hardware=args.hardware,
        private=args.private, bundle_weights=args.bundle_weights, skip_secret=args.skip_secret,
    )
    if not plan.model_dir.is_dir():
        plan.problems.append(f"model dir missing: {plan.model_dir} (run export.py first)")
    else:
        plan.model_files = _select(plan.model_dir, MODEL_ALLOW_PATTERNS, None)
        present = {rel for rel, _ in plan.model_files}
        for f in MODEL_ALLOW_PATTERNS[:3]:
            if f not in present:
                plan.problems.append(f"model dir lacks {f}")
    if not plan.app_dir.is_dir():
        plan.problems.append(f"app dir missing: {plan.app_dir}")
    else:
        plan.space_files = _select(plan.app_dir, None, plan.space_ignore)
        present = {rel for rel, _ in plan.space_files}
        for f in SPACE_REQUIRED_FILES:
            if f not in present:
                plan.problems.append(f"app dir lacks {f}")
        if plan.bundle_weights:
            for f in ("model.safetensors", "config.json", "tokenizer.json"):
                if f not in present:
                    plan.problems.append(f"--bundle-weights but app dir lacks {f} (src/export.py run copies them into the app dir)")
        src_model = common.project_root() / "src" / "model.py"
        app_model = plan.app_dir / "model.py"
        if src_model.is_file() and app_model.is_file() and src_model.read_bytes() != app_model.read_bytes():
            plan.problems.append("app/model.py is not byte-identical to src/model.py")
    if not plan.bundle_weights and not plan.skip_secret and not os.environ.get(SECRET_ENV):
        plan.problems.append(
            f"{SECRET_ENV} is not set. The Space needs a READ token stored as the HF_TOKEN secret to download "
            f"the private weights at startup. Export it in the shell that runs this script (never commit it), "
            f"or pass --skip-secret if the secret already exists on the Space, or --bundle-weights."
        )
    return plan


def _fmt_size(n: int) -> str:
    for unit in ("B", "KiB", "MiB", "GiB"):
        if n < 1024 or unit == "GiB":
            return f"{n:,.0f} {unit}" if unit == "B" else f"{n:.1f} {unit}"
        n /= 1024
    return f"{n:.1f} GiB"


def print_plan(plan: Plan, dry_run: bool) -> None:
    vis = "private" if plan.private else "PUBLIC"
    print(f"\n{'DRY RUN - nothing will be sent' if dry_run else 'LIVE RUN'}")
    print(f"account token source: {'HF_TOKEN env' if os.environ.get('HF_TOKEN') else 'hf auth login cache'} (value never printed)")
    print(f"\n[1] model repo  {plan.model_repo}  ({vis}, repo_type=model, exist_ok)")
    print(f"    upload_folder({plan.model_dir}) allow_patterns={MODEL_ALLOW_PATTERNS}")
    for rel, size in plan.model_files:
        print(f"      {rel:<24} {_fmt_size(size):>12}")
    print(f"    total {_fmt_size(sum(s for _, s in plan.model_files))} in {len(plan.model_files)} files")
    print(f"\n[2] space repo  {plan.space_repo}  ({vis}, sdk=gradio, hardware={plan.hardware}, exist_ok)")
    if plan.bundle_weights:
        print("    secret HF_TOKEN: not needed (weights bundled)")
    elif plan.skip_secret:
        print("    secret HF_TOKEN: skipped (--skip-secret)")
    else:
        print(f"    secret HF_TOKEN  <- ${SECRET_ENV} ({'set' if os.environ.get(SECRET_ENV) else 'NOT SET'})")
    print(f"    variable PHISH_MODEL_REPO = {plan.model_repo}")
    print(f"    upload_folder({plan.app_dir}) ignore_patterns={plan.space_ignore}")
    for rel, size in plan.space_files:
        print(f"      {rel:<24} {_fmt_size(size):>12}")
    print(f"    total {_fmt_size(sum(s for _, s in plan.space_files))} in {len(plan.space_files)} files")
    print(f"\n    urls: https://huggingface.co/{plan.model_repo}  https://huggingface.co/spaces/{plan.space_repo}")
    if plan.problems:
        print("\nPROBLEMS (a live run refuses to start while any remains):")
        for p in plan.problems:
            print(f"  - {p}")
    print()


def check_auth() -> str:
    """Read-only identity check; logs the account name only."""
    from huggingface_hub import HfApi

    me = HfApi().whoami()
    name = me.get("name", "?")
    LOG.info("authenticated as %s (type=%s)", name, me.get("type"))
    return name


def run_live(plan: Plan, commit_message: str) -> int:
    from huggingface_hub import HfApi

    api = HfApi()
    name = check_auth()
    owner = plan.model_repo.split("/")[0]
    if name != owner:
        LOG.warning("token belongs to %s but repos are under %s; the push will fail unless %s has write access", name, owner, name)

    LOG.info("create_repo %s (model, private=%s)", plan.model_repo, plan.private)
    url = api.create_repo(plan.model_repo, repo_type="model", private=plan.private, exist_ok=True)
    LOG.info("model repo: %s", url)
    LOG.info("upload_folder %s -> %s", plan.model_dir, plan.model_repo)
    info = api.upload_folder(
        folder_path=str(plan.model_dir), repo_id=plan.model_repo, repo_type="model",
        commit_message=commit_message, allow_patterns=MODEL_ALLOW_PATTERNS,
    )
    LOG.info("model commit: %s", getattr(info, "commit_url", info))

    LOG.info("create_repo %s (space, sdk=gradio, hardware=%s, private=%s)", plan.space_repo, plan.hardware, plan.private)
    url = api.create_repo(
        plan.space_repo, repo_type="space", space_sdk="gradio", space_hardware=plan.hardware,
        private=plan.private, exist_ok=True,
    )
    LOG.info("space repo: %s", url)
    if not plan.bundle_weights and not plan.skip_secret:
        api.add_space_secret(plan.space_repo, "HF_TOKEN", os.environ[SECRET_ENV],
                             description="read token for the private model repo")
        LOG.info("space secret HF_TOKEN set")
    api.add_space_variable(plan.space_repo, "PHISH_MODEL_REPO", plan.model_repo,
                           description="model repo app.py downloads weights from")
    LOG.info("space variable PHISH_MODEL_REPO=%s set", plan.model_repo)
    LOG.info("upload_folder %s -> %s", plan.app_dir, plan.space_repo)
    info = api.upload_folder(
        folder_path=str(plan.app_dir), repo_id=plan.space_repo, repo_type="space",
        commit_message=commit_message, ignore_patterns=plan.space_ignore,
    )
    LOG.info("space commit: %s", getattr(info, "commit_url", info))

    runtime = api.get_space_runtime(plan.space_repo)
    hardware = getattr(runtime, "hardware", None)
    requested = getattr(runtime, "requested_hardware", None)
    print(f"\nmodel: https://huggingface.co/{plan.model_repo}")
    print(f"space: https://huggingface.co/spaces/{plan.space_repo}")
    print(f"space runtime: stage={getattr(runtime, 'stage', '?')} hardware={hardware} requested={requested}")
    if hardware != plan.hardware and requested != plan.hardware:
        LOG.warning("hardware is %s, not %s: set it in Space settings or api.request_space_hardware()", hardware, plan.hardware)
    return 0


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0], formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--config", default="config.yaml")
    ap.add_argument("--model-dir", default=None, help="export dir (model_export/<run>); default: newest under model_export/")
    ap.add_argument("--app-dir", default=None, help="Space source dir (default: config paths.app_dir)")
    ap.add_argument("--model-repo", default=None, help="default: config hf.model_repo")
    ap.add_argument("--space-repo", default=None, help="default: config hf.space_repo")
    ap.add_argument("--hardware", default=None, help="Space hardware flavour (default: config hf.space_hardware)")
    ap.add_argument("--private", action=argparse.BooleanOptionalAction, default=True)
    ap.add_argument("--dry-run", action=argparse.BooleanOptionalAction, default=None,
                    help="print the plan only (DEFAULT). --no-dry-run alone is refused; add --yes")
    ap.add_argument("--yes", action="store_true", help="actually write to the Hub (Nathan only, after approval)")
    ap.add_argument("--check-auth", action="store_true", help="dry run: also call whoami() (read-only)")
    ap.add_argument("--bundle-weights", action="store_true", help="upload weights into the Space instead of the model repo download")
    ap.add_argument("--skip-secret", action="store_true", help="do not set the HF_TOKEN space secret (already set)")
    ap.add_argument("--commit-message", default="Upload BITS phishing detector")
    return ap.parse_args(argv)


def _newest_export(root: Path) -> Path | None:
    if not root.is_dir():
        return None
    runs = [p for p in root.iterdir() if p.is_dir()]
    return max(runs, key=lambda p: p.stat().st_mtime) if runs else None


def main(argv: Sequence[str] | None = None) -> int:
    args = parse_args(argv)
    cfg = common.load_config(args.config)
    global LOG
    LOG = common.setup_logging("upload_to_hf", cfg["paths"]["logs"])
    hf = cfg.get("hf", {})
    root = common.project_root()
    args.model_repo = args.model_repo or hf.get("model_repo", "BuildITSystems/phishing-detector-encoder")
    args.space_repo = args.space_repo or hf.get("space_repo", "BuildITSystems/phishing-detector-demo")
    args.hardware = args.hardware or hf.get("space_hardware", "zero-a10g")
    args.app_dir = args.app_dir or cfg["paths"].get("app_dir", str(root / "app"))
    if args.model_dir is None:
        newest = _newest_export(root / "model_export")
        args.model_dir = str(newest) if newest else str(root / "model_export" / "<run_name>")

    dry_run = args.dry_run if args.dry_run is not None else not args.yes
    if not dry_run and not args.yes:
        LOG.error("--no-dry-run requires --yes; refusing")
        return 2
    if dry_run and args.yes:
        LOG.info("--dry-run given together with --yes: dry run wins")

    plan = build_plan(args)
    print_plan(plan, dry_run)
    if dry_run:
        if args.check_auth:
            check_auth()
        LOG.info("dry run complete; no Hub write performed")
        return 0
    if plan.problems:
        LOG.error("refusing live run: %d problem(s) listed above", len(plan.problems))
        return 2
    return run_live(plan, args.commit_message)


if __name__ == "__main__":
    raise SystemExit(main())