mf_inference.py

#!/usr/bin/env python3
"""
mf_inference.py — sample script that loads the trained MF model bundle and runs model output.

Hugging Face ecosystem used here:
  * `datasets`        -> loads the preference CSV into a HF Dataset (falls back to pandas)
  * `huggingface_hub` -> optional `--push_to_hub`: creates + uploads the bundle as a HF model repo

Examples
--------
  # score specific (user, prompt) pairs
  python mf_inference.py --model_dir mf_bundle --user_id U0007 \
      --prompt_ids P0001,P0400,P0572 --csv preference_data_synthetic.csv

  # rank all known prompts for a user (top cloud / top local)
  python mf_inference.py --model_dir mf_bundle --user_id U0007 --top_k 5 \
      --csv preference_data_synthetic.csv

  # push the bundle to the Hugging Face Hub (needs HF_TOKEN or huggingface-cli login)
  python mf_inference.py --model_dir mf_bundle --user_id U0007 --top_k 3 \
      --csv preference_data_synthetic.csv --push_to_hub --repo_id your-org/cloud-local-mf
"""
import argparse
import json
import os
import sys
from pathlib import Path

import numpy as np

try:  # HF ecosystem (optional but preferred)
    from datasets import Dataset
    HAVE_DATASETS = True
except ImportError:
    HAVE_DATASETS = False

try:
    from huggingface_hub import HfApi, upload_folder
    HAVE_HUB = True
except ImportError:
    HAVE_HUB = False

MODEL_KEYS = ("mu", "bu", "bi", "P", "Q", "user_ids", "prompt_ids")


def sigmoid(z):
    return 1.0 / (1.0 + np.exp(-np.clip(z, -30, 30)))


class SimpleMF:
    """HF-style loader for the bundle written by the training notebook.

    from_pretrained() reads config.json + mf_params.npz so inference never
    depends on the training code or kernel state.
    """

    def __init__(self, config, params):
        self.config = config
        self.mu = float(params["mu"])
        self.bu = params["bu"]
        self.bi = params["bi"]
        self.P = params["P"]
        self.Q = params["Q"]
        self.user_ids = [str(x) for x in params["user_ids"]]
        self.prompt_ids = [str(x) for x in params["prompt_ids"]]
        self._uidx = {u: i for i, u in enumerate(self.user_ids)}
        self._pidx = {p: i for i, p in enumerate(self.prompt_ids)}

    @classmethod
    def from_pretrained(cls, model_dir):
        model_dir = Path(model_dir)
        config = json.loads((model_dir / "config.json").read_text())
        params = np.load(model_dir / "mf_params.npz", allow_pickle=True)
        missing = [k for k in MODEL_KEYS if k not in params.files]
        if missing:
            raise ValueError(f"bundle {model_dir} is missing: {missing}")
        return cls(config, params)

    # ---- scoring ------------------------------------------------------
    def score_ids(self, user_ids, prompt_ids):
        """Raw scores r_hat for lists of string ids (both must be known)."""
        u = np.array([self._uidx[x] for x in user_ids])
        i = np.array([self._pidx[x] for x in prompt_ids])
        return self.mu + self.bu[u] + self.bi[i] + (self.P[u] * self.Q[i]).sum(1)

    def predict(self, user_ids, prompt_ids):
        """P(cloud preferred) in [0, 1] for (user, prompt) pairs."""
        return sigmoid(self.score_ids(user_ids, prompt_ids))

    def rank_for_user(self, user_id, top_k=5):
        """Score every known prompt for one user; returns (desc, asc) arrays of rows."""
        if user_id not in self._uidx:
            raise KeyError(f"unknown user '{user_id}' — bundle knows {len(self.user_ids)} users")
        u = self._uidx[user_id]
        r = self.mu + self.bu[u] + self.bi + (self.P[u] * self.Q).sum(1)
        p = sigmoid(r)
        order = np.argsort(-p)
        def rows(idx):
            return [{"prompt_id": self.prompt_ids[j], "p_cloud": float(p[j]),
                     "choice": "cloud" if p[j] >= 0.5 else "local"} for j in idx]
        return rows(order[:top_k]), rows(order[-top_k:][::-1])


def load_catalog(path):
    """Load the preference file; returns dict prompt_id -> {topic, text} (best effort)."""
    catalog = {}
    if path is None or not Path(path).exists():
        return catalog
    df = Dataset.from_csv(path) if HAVE_DATASETS else _pandas_read(path)
    for row in df:
        pid = str(row.get("prompt_id", row.get("prompt", "")))
        if pid:
            catalog[pid] = {"topic": str(row.get("topic", "")),
                            "text": str(row.get("prompt_text", row.get("prompt", "")))}
    return catalog


def _pandas_read(path):
    import pandas as pd
    return pd.read_csv(path)


def print_report(rows, catalog, title):
    print(f"\n{title}")
    print(f"{'prompt_id':<10}{'p(cloud)':>9}  {'choice':<6}  topic / prompt")
    print("-" * 78)
    for r in rows:
        meta = catalog.get(r["prompt_id"], {})
        topic = meta.get("topic", "?")
        text = meta.get("text", "")
        text = text[:46] + "…" if len(text) > 46 else text
        print(f"{r['prompt_id']:<10}{r['p_cloud']:>9.3f}  {r['choice']:<6}  {topic:<18} {text}")


def main():
    ap = argparse.ArgumentParser(description="Load the MF bundle and run model output")
    ap.add_argument("--model_dir", default="mf_bundle", help="path to the saved bundle")
    ap.add_argument("--user_id", default="U0007")
    ap.add_argument("--prompt_ids", help="comma-separated prompt ids to score")
    ap.add_argument("--top_k", type=int, default=0, help="rank top-k prompts for the user")
    ap.add_argument("--csv", help="preference CSV (for topic/text display)")
    ap.add_argument("--push_to_hub", action="store_true", help="upload bundle as a HF model repo")
    ap.add_argument("--repo_id", default=None, help="HF repo id, e.g. your-org/cloud-local-mf")
    args = ap.parse_args()

    libs = [f"numpy {np.__version__}"]
    if HAVE_DATASETS:
        import datasets
        libs.append(f"datasets {datasets.__version__}")
    if HAVE_HUB:
        import huggingface_hub
        libs.append(f"huggingface_hub {huggingface_hub.__version__}")
    print("python libs:", ", ".join(libs))

    model = SimpleMF.from_pretrained(args.model_dir)
    print(f"loaded bundle: {Path(args.model_dir).resolve()} "
          f"({len(model.user_ids)} users x {len(model.prompt_ids)} prompts, k={model.P.shape[1]})")
    print(f"model config : {model.config.get('model')}")

    catalog = load_catalog(args.csv)

    if args.prompt_ids:
        pids = [p.strip() for p in args.prompt_ids.split(",") if p.strip()]
        unknown = [p for p in pids if p not in model._pidx]
        if unknown:
            print(f"ERROR: unknown prompt ids {unknown} — bundle knows {len(model.prompt_ids)} prompts")
            sys.exit(2)
        p = model.predict([args.user_id] * len(pids), pids)
        rows = [{"prompt_id": pid, "p_cloud": float(pi), "choice": "cloud" if pi >= 0.5 else "local"}
                for pid, pi in zip(pids, p)]
        print_report(rows, catalog, f"model output — choices for user {args.user_id}")

    if args.top_k:
        top, bottom = model.rank_for_user(args.user_id, args.top_k)
        print_report(top, catalog, f"user {args.user_id} — top {args.top_k} prompts → cloud")
        print_report(bottom, catalog, f"user {args.user_id} — top {args.top_k} prompts → local")

    # optional Hub push -------------------------------------------------
    if args.push_to_hub:
        if not HAVE_HUB:
            print("huggingface_hub not installed — cannot push to the Hub")
            sys.exit(1)
        if not args.repo_id:
            print("--push_to_hub requires --repo_id, e.g. your-org/cloud-local-mf")
            sys.exit(2)
        token = os.environ.get("HF_TOKEN", None)
        if not token:
            print("No HF_TOKEN in environment. Run `huggingface-cli login` (or set HF_TOKEN) "
                  "and re-run to push.")
            sys.exit(0)
        api = HfApi()
        api.whoami(token=token)
        api.create_repo(repo_id=args.repo_id, token=token, repo_type="model", exist_ok=True)
        upload_folder(folder_path=str(Path(args.model_dir).resolve()),
                      repo_id=args.repo_id, token=token,
                      commit_message="Add cloud-vs-local MF preference model bundle")
        print(f"pushed bundle -> https://huggingface.co/{args.repo_id}")

    print("\nmodel output complete.")


if __name__ == "__main__":
    main()