#!/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()