backend/src/deploy_model.py

"""Ship a trained checkpoint to the Hub, with a rollback point and gates.

    python -m backend.src.deploy_model models/production_v2/final            # dry run
    python -m backend.src.deploy_model models/production_v2/final --confirm  # ship it

Production reads one repo, VivDubs/eumora-emotion-classifier, from two places:
the Vercel /api/emotion route (user text, via the HF inference router) and the
HF Space's /api/score (lyrics, loaded at container start). Pushing to that repo
therefore changes both, so this refuses to ship anything that fails the checks
this project learned the hard way:

  - the checkpoint must not be collapsed. Training reports a representation
    health check; a model that emits the same output for every input looks like
    underfitting in the metrics and would be invisible in production.
  - it must carry folded calibration, because the Space only disables the
    legacy sarcasm hacks in predict.py for calibrated checkpoints. Shipping an
    uncalibrated model would silently re-enable them.
  - the label space must match the eight classes both callers map by name.

Before uploading it tags the current head as a rollback point, and afterwards it
restarts the Space so /api/score loads the new weights instead of the ones it
cached at startup.

Needs HF_TOKEN with write access to the repo, read from the environment or .env.
Never pass the token on the command line.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from datetime import datetime
from pathlib import Path

MODEL_REPO = os.environ.get("HF_MODEL_REPO", "VivDubs/eumora-emotion-classifier")
SPACE_REPO = os.environ.get("HF_SPACE_REPO", "VivDubs/eumora-api")
EXPECTED_LABELS = ["sadness", "joy", "love", "anger", "fear", "surprise", "neutral", "sarcasm"]
# Only the files a stock sequence classifier needs. Keeps training_results.json
# and anything else local out of the public-facing model card area.
SHIP_FILES = {
    "config.json", "model.safetensors", "tokenizer.json", "tokenizer_config.json",
    "special_tokens_map.json", "spm.model", "added_tokens.json",
}


def load_env() -> None:
    for candidate in (Path(__file__).resolve().parents[2] / ".env",
                      Path(__file__).resolve().parents[1] / ".env"):
        if candidate.exists():
            for line in candidate.read_text().splitlines():
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    k, v = line.split("=", 1)
                    os.environ.setdefault(k.strip(), v.strip().strip("\"'"))


def check_checkpoint(path: Path) -> list:
    """Return a list of blocking problems; empty means shippable."""
    problems = []
    cfg_path = path / "config.json"
    if not cfg_path.exists():
        return [f"{cfg_path} not found"]
    cfg = json.loads(cfg_path.read_text())

    labels = [cfg.get("id2label", {}).get(str(i)) for i in range(len(cfg.get("id2label", {})))]
    if labels != EXPECTED_LABELS:
        problems.append(f"label space {labels} != {EXPECTED_LABELS}")

    if not (cfg.get("calibration") or {}).get("folded"):
        problems.append("calibration is not folded into the weights — the Space would "
                        "re-apply the legacy sarcasm hacks on top of it")

    results = path / "training_results.json"
    if not results.exists():
        problems.append("training_results.json missing — cannot verify representation health")
    else:
        health = json.loads(results.read_text()).get("representation_health") or {}
        if health.get("collapsed", True):
            problems.append(f"representation health says collapsed or unknown: {health}")

    if not (path / "model.safetensors").exists():
        problems.append("model.safetensors missing")
    return problems


def check_ranking(candidate_json: Path, live_json: Path) -> list:
    """Refuse a model that does not beat the live one on the product metric.

    Classifier metrics do not predict this. The 20k improved checkpoint scored
    68.5% in-domain macro-F1 and still tied the live model on ranking AUC
    (0.779 vs 0.785, paired difference -0.006 [-0.120, +0.093]) while losing the
    confidence signal the UI depends on. Ship on ranking, not on F1.
    """
    import json as _json

    problems = []
    for path in (candidate_json, live_json):
        if not path.exists():
            problems.append(f"{path} not found — produce it with bench_production --ranking")
    if problems:
        return problems
    cand = _json.loads(candidate_json.read_text()).get("weighted_auc")
    live = _json.loads(live_json.read_text()).get("weighted_auc")
    if cand is None or live is None:
        return ["ranking results are missing weighted_auc"]
    if cand <= live:
        problems.append(f"candidate ranking AUC {cand:.3f} does not beat live {live:.3f} — "
                        f"deploying would not improve what listeners see")
    return problems


def check_system(compare_json: Path) -> list:
    """Accept a full-system comparison (train_fusion compare) as evidence.

    Since the fusion (FINDINGS section 9), lyric scores reach listeners only
    through the stacked VAD match, so lyric-only ranking is no longer the right
    test. `compare` runs the whole system with the candidate text model against
    the live one on 4Q, PMEmo and gold and records whether its ship rule held.
    """
    if not compare_json.exists():
        return [f"{compare_json} not found — run `train_fusion lyrics` then `train_fusion compare`"]
    report = json.loads(compare_json.read_text())
    if not report.get("ship"):
        return [f"system comparison says do not ship: {report}"]
    return []


def main() -> int:
    parser = argparse.ArgumentParser(description="Deploy a checkpoint to the production repo")
    parser.add_argument("checkpoint")
    parser.add_argument("--confirm", action="store_true", help="Actually upload (default: dry run)")
    parser.add_argument("--no-restart", action="store_true", help="Do not restart the Space")
    parser.add_argument("--ranking", type=str, default=None,
                        help="bench_production --ranking --local JSON for this checkpoint")
    parser.add_argument("--live-ranking", type=str, default="results/ranking_live.json",
                        help="bench_production --ranking JSON for the live model")
    parser.add_argument("--override-ranking", action="store_true",
                        help="Ship despite no ranking improvement (state why in the commit)")
    parser.add_argument("--system", type=str, default=None,
                        help="results/fusion_compare_<tag>.json from train_fusion compare; replaces "
                             "--ranking as evidence and installs configs/fusion_<tag>.json with the model")
    args = parser.parse_args()

    load_env()
    ckpt = Path(args.checkpoint)
    print(f"checkpoint : {ckpt}")
    print(f"model repo : {MODEL_REPO}")
    print(f"space      : {SPACE_REPO}\n")

    problems = check_checkpoint(ckpt)
    fusion_swap = None
    if args.system:
        problems += check_system(Path(args.system))
        tag = Path(args.system).stem.replace("fusion_compare_", "")
        fusion_swap = Path(__file__).resolve().parents[1] / "configs" / f"fusion_{tag}.json"
        if not fusion_swap.exists():
            problems.append(f"{fusion_swap} missing — the fusion must be refit on this model's lyric scores")
    elif not args.override_ranking:
        if not args.ranking:
            problems.append("no ranking evidence — run `bench_production --ranking --local "
                            f"{ckpt} --json <out>` and pass it with --ranking")
        else:
            problems += check_ranking(Path(args.ranking), Path(args.live_ranking))
    if problems:
        print("REFUSING TO DEPLOY:")
        for p in problems:
            print(f"  - {p}")
        return 1
    results = json.loads((ckpt / "training_results.json").read_text())
    tm, probe = results.get("test_metrics", {}), results.get("probe", {})
    print("checks passed")
    print(f"  test macro-F1 {tm.get('test_f1_macro')}   probe overall {probe.get('probe_overall')}")
    print(f"  health        {results.get('representation_health')}\n")

    token = os.environ.get("HF_TOKEN") or os.environ.get("HF_API_TOKEN")
    if not token:
        print("No HF_TOKEN in the environment or .env. Add a token with *write* access to")
        print(f"{MODEL_REPO} to .env as HF_TOKEN=... (do not paste it into chat), then re-run.")
        return 1

    from huggingface_hub import HfApi

    api = HfApi(token=token)
    who = api.whoami()
    role = (who.get("auth", {}).get("accessToken", {}) or {}).get("role")
    print(f"authenticated as {who.get('name')} (token role: {role})")

    files = sorted(f.name for f in ckpt.iterdir() if f.name in SHIP_FILES)
    tag = "pre-" + datetime.now().strftime("%Y%m%d-%H%M")
    print(f"will tag current head as '{tag}' (rollback point)")
    print(f"will upload: {', '.join(files)}")
    print(f"will restart space: {not args.no_restart}")

    if not args.confirm:
        print("\ndry run — nothing changed. Re-run with --confirm to deploy.")
        return 0

    head = api.model_info(MODEL_REPO).sha
    api.create_tag(MODEL_REPO, tag=tag, revision=head,
                   tag_message=f"State before deploying {ckpt.name}")
    print(f"\ntagged {head[:10]} as {tag}")

    commit = api.upload_folder(
        folder_path=str(ckpt), repo_id=MODEL_REPO, repo_type="model",
        allow_patterns=files,
        commit_message=f"Deploy calibrated checkpoint ({ckpt.parent.name})",
    )
    print(f"uploaded: {commit.commit_url if hasattr(commit, 'commit_url') else commit}")

    # Verify what the repo now serves, rather than trusting the upload.
    from huggingface_hub import hf_hub_download
    served = json.loads(Path(hf_hub_download(MODEL_REPO, "config.json", token=token,
                                             force_download=True)).read_text())
    if not (served.get("calibration") or {}).get("folded"):
        print("WARNING: served config does not show folded calibration — investigate")
        return 1
    print("verified: served config carries folded calibration")

    if fusion_swap is not None:
        live_cfg = fusion_swap.parent / "fusion.json"
        backup = fusion_swap.parent / f"fusion.before-{tag}.json"
        if live_cfg.exists():
            backup.write_text(live_cfg.read_text())
        live_cfg.write_text(fusion_swap.read_text())
        print(f"installed {fusion_swap.name} as configs/fusion.json (previous kept as {backup.name}).")
        print("Commit and push it now: the push rebuilds the Space, which then loads the new weights")
        print("and the matching fusion together. Until then the Space runs the new text model with the")
        print("old fusion, so do not restart it separately.")
    elif not args.no_restart:
        api.restart_space(SPACE_REPO)
        print(f"restarted {SPACE_REPO} so /api/score loads the new weights")

    print(f"\nrollback: python -c \"from huggingface_hub import HfApi; "
          f"HfApi().upload_folder(...)\" from tag '{tag}', or in the Hub UI revert to {head[:10]}")
    return 0


if __name__ == "__main__":
    sys.exit(main())