| |
| """ |
| 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: |
| 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) |
|
|
| |
| 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") |
|
|
| |
| 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() |
|
|
