subhash4face/cloud-local-mf

1#!/usr/bin/env python32"""3mf_inference.py — sample script that loads the trained MF model bundle and runs model output.4 5Hugging Face ecosystem used here:6  * `datasets`        -> loads the preference CSV into a HF Dataset (falls back to pandas)7  * `huggingface_hub` -> optional `--push_to_hub`: creates + uploads the bundle as a HF model repo8 9Examples10--------11  # score specific (user, prompt) pairs12  python mf_inference.py --model_dir mf_bundle --user_id U0007 \13      --prompt_ids P0001,P0400,P0572 --csv preference_data_synthetic.csv14 15  # rank all known prompts for a user (top cloud / top local)16  python mf_inference.py --model_dir mf_bundle --user_id U0007 --top_k 5 \17      --csv preference_data_synthetic.csv18 19  # push the bundle to the Hugging Face Hub (needs HF_TOKEN or huggingface-cli login)20  python mf_inference.py --model_dir mf_bundle --user_id U0007 --top_k 3 \21      --csv preference_data_synthetic.csv --push_to_hub --repo_id your-org/cloud-local-mf22"""23import argparse24import json25import os26import sys27from pathlib import Path28 29import numpy as np30 31try:  # HF ecosystem (optional but preferred)32    from datasets import Dataset33    HAVE_DATASETS = True34except ImportError:35    HAVE_DATASETS = False36 37try:38    from huggingface_hub import HfApi, upload_folder39    HAVE_HUB = True40except ImportError:41    HAVE_HUB = False42 43MODEL_KEYS = ("mu", "bu", "bi", "P", "Q", "user_ids", "prompt_ids")44 45 46def sigmoid(z):47    return 1.0 / (1.0 + np.exp(-np.clip(z, -30, 30)))48 49 50class SimpleMF:51    """HF-style loader for the bundle written by the training notebook.52 53    from_pretrained() reads config.json + mf_params.npz so inference never54    depends on the training code or kernel state.55    """56 57    def __init__(self, config, params):58        self.config = config59        self.mu = float(params["mu"])60        self.bu = params["bu"]61        self.bi = params["bi"]62        self.P = params["P"]63        self.Q = params["Q"]64        self.user_ids = [str(x) for x in params["user_ids"]]65        self.prompt_ids = [str(x) for x in params["prompt_ids"]]66        self._uidx = {u: i for i, u in enumerate(self.user_ids)}67        self._pidx = {p: i for i, p in enumerate(self.prompt_ids)}68 69    @classmethod70    def from_pretrained(cls, model_dir):71        model_dir = Path(model_dir)72        config = json.loads((model_dir / "config.json").read_text())73        params = np.load(model_dir / "mf_params.npz", allow_pickle=True)74        missing = [k for k in MODEL_KEYS if k not in params.files]75        if missing:76            raise ValueError(f"bundle {model_dir} is missing: {missing}")77        return cls(config, params)78 79    # ---- scoring ------------------------------------------------------80    def score_ids(self, user_ids, prompt_ids):81        """Raw scores r_hat for lists of string ids (both must be known)."""82        u = np.array([self._uidx[x] for x in user_ids])83        i = np.array([self._pidx[x] for x in prompt_ids])84        return self.mu + self.bu[u] + self.bi[i] + (self.P[u] * self.Q[i]).sum(1)85 86    def predict(self, user_ids, prompt_ids):87        """P(cloud preferred) in [0, 1] for (user, prompt) pairs."""88        return sigmoid(self.score_ids(user_ids, prompt_ids))89 90    def rank_for_user(self, user_id, top_k=5):91        """Score every known prompt for one user; returns (desc, asc) arrays of rows."""92        if user_id not in self._uidx:93            raise KeyError(f"unknown user '{user_id}' — bundle knows {len(self.user_ids)} users")94        u = self._uidx[user_id]95        r = self.mu + self.bu[u] + self.bi + (self.P[u] * self.Q).sum(1)96        p = sigmoid(r)97        order = np.argsort(-p)98        def rows(idx):99            return [{"prompt_id": self.prompt_ids[j], "p_cloud": float(p[j]),100                     "choice": "cloud" if p[j] >= 0.5 else "local"} for j in idx]101        return rows(order[:top_k]), rows(order[-top_k:][::-1])102 103 104def load_catalog(path):105    """Load the preference file; returns dict prompt_id -> {topic, text} (best effort)."""106    catalog = {}107    if path is None or not Path(path).exists():108        return catalog109    df = Dataset.from_csv(path) if HAVE_DATASETS else _pandas_read(path)110    for row in df:111        pid = str(row.get("prompt_id", row.get("prompt", "")))112        if pid:113            catalog[pid] = {"topic": str(row.get("topic", "")),114                            "text": str(row.get("prompt_text", row.get("prompt", "")))}115    return catalog116 117 118def _pandas_read(path):119    import pandas as pd120    return pd.read_csv(path)121 122 123def print_report(rows, catalog, title):124    print(f"\n{title}")125    print(f"{'prompt_id':<10}{'p(cloud)':>9}  {'choice':<6}  topic / prompt")126    print("-" * 78)127    for r in rows:128        meta = catalog.get(r["prompt_id"], {})129        topic = meta.get("topic", "?")130        text = meta.get("text", "")131        text = text[:46] + "…" if len(text) > 46 else text132        print(f"{r['prompt_id']:<10}{r['p_cloud']:>9.3f}  {r['choice']:<6}  {topic:<18} {text}")133 134 135def main():136    ap = argparse.ArgumentParser(description="Load the MF bundle and run model output")137    ap.add_argument("--model_dir", default="mf_bundle", help="path to the saved bundle")138    ap.add_argument("--user_id", default="U0007")139    ap.add_argument("--prompt_ids", help="comma-separated prompt ids to score")140    ap.add_argument("--top_k", type=int, default=0, help="rank top-k prompts for the user")141    ap.add_argument("--csv", help="preference CSV (for topic/text display)")142    ap.add_argument("--push_to_hub", action="store_true", help="upload bundle as a HF model repo")143    ap.add_argument("--repo_id", default=None, help="HF repo id, e.g. your-org/cloud-local-mf")144    args = ap.parse_args()145 146    libs = [f"numpy {np.__version__}"]147    if HAVE_DATASETS:148        import datasets149        libs.append(f"datasets {datasets.__version__}")150    if HAVE_HUB:151        import huggingface_hub152        libs.append(f"huggingface_hub {huggingface_hub.__version__}")153    print("python libs:", ", ".join(libs))154 155    model = SimpleMF.from_pretrained(args.model_dir)156    print(f"loaded bundle: {Path(args.model_dir).resolve()} "157          f"({len(model.user_ids)} users x {len(model.prompt_ids)} prompts, k={model.P.shape[1]})")158    print(f"model config : {model.config.get('model')}")159 160    catalog = load_catalog(args.csv)161 162    if args.prompt_ids:163        pids = [p.strip() for p in args.prompt_ids.split(",") if p.strip()]164        unknown = [p for p in pids if p not in model._pidx]165        if unknown:166            print(f"ERROR: unknown prompt ids {unknown} — bundle knows {len(model.prompt_ids)} prompts")167            sys.exit(2)168        p = model.predict([args.user_id] * len(pids), pids)169        rows = [{"prompt_id": pid, "p_cloud": float(pi), "choice": "cloud" if pi >= 0.5 else "local"}170                for pid, pi in zip(pids, p)]171        print_report(rows, catalog, f"model output — choices for user {args.user_id}")172 173    if args.top_k:174        top, bottom = model.rank_for_user(args.user_id, args.top_k)175        print_report(top, catalog, f"user {args.user_id} — top {args.top_k} prompts → cloud")176        print_report(bottom, catalog, f"user {args.user_id} — top {args.top_k} prompts → local")177 178    # optional Hub push -------------------------------------------------179    if args.push_to_hub:180        if not HAVE_HUB:181            print("huggingface_hub not installed — cannot push to the Hub")182            sys.exit(1)183        if not args.repo_id:184            print("--push_to_hub requires --repo_id, e.g. your-org/cloud-local-mf")185            sys.exit(2)186        token = os.environ.get("HF_TOKEN", None)187        if not token:188            print("No HF_TOKEN in environment. Run `huggingface-cli login` (or set HF_TOKEN) "189                  "and re-run to push.")190            sys.exit(0)191        api = HfApi()192        api.whoami(token=token)193        api.create_repo(repo_id=args.repo_id, token=token, repo_type="model", exist_ok=True)194        upload_folder(folder_path=str(Path(args.model_dir).resolve()),195                      repo_id=args.repo_id, token=token,196                      commit_message="Add cloud-vs-local MF preference model bundle")197        print(f"pushed bundle -> https://huggingface.co/{args.repo_id}")198 199    print("\nmodel output complete.")200 201 202if __name__ == "__main__":203    main()204