Spaces:
Running
Running
Download train_head.py from SharpMeow/vesuvius-radar: direct link, hf CLI and curl.
- Browser
- Download file 6.7 kB
-
https://huggingface.co/spaces/SharpMeow/vesuvius-radar/resolve/main/train_head.py
- Command line
-
hf download hf://spaces/SharpMeow/vesuvius-radar/train_head.py
-
curl -L -o train_head.py https://huggingface.co/spaces/SharpMeow/vesuvius-radar/resolve/main/train_head.py
6.7 kB
| """Fit the radar's learned relevance head (head.json) from blind ratings of collected items. | |
| Input: a JSON list of items as the radar stores them (GET /api/items or the store file) and a ratings file mapping each | |
| item id to one or more raters' {"value": 0..100}. The target is the mean rating. Two ridge heads are fitted: "full" | |
| (bge-small embedding of the item text plus metadata) and "meta" (metadata only, for items with no embedding). Each head | |
| is calibrated to 0..100 by isotonic regression on out-of-fold predictions, so the mapping is monotone and the bands | |
| mean what the ratings mean. Nested cross-validation reports what to expect on new items. Training needs scikit-learn | |
| and scipy; the Space itself only needs numpy to apply the head. | |
| python train_head.py --items items.json --ratings ratings.json --out head.json | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import datetime as dt | |
| import json | |
| import numpy as np | |
| import scanner | |
| def mean_rating(r) -> float: | |
| vals = [v["value"] for v in (r.values() if isinstance(r, dict) and "value" not in r else [r])] | |
| return float(np.mean(vals)) | |
| def fit(X, y, alphas): | |
| from sklearn.linear_model import RidgeCV | |
| return RidgeCV(alphas=alphas).fit(X, y) | |
| def oof(X, y, alphas, seed, folds=5): | |
| from sklearn.model_selection import KFold | |
| pred = np.zeros(len(y)) | |
| for tr, te in KFold(folds, shuffle=True, random_state=seed).split(X): | |
| pred[te] = fit(X[tr], y[tr], alphas).predict(X[te]) | |
| return pred | |
| def calibrator(pred, y, how="quantile"): | |
| """Map raw predictions to 0..100. "quantile" (default) gives the k-th percentile prediction the k-th percentile | |
| rating, so each band holds about as many items as the ratings put there; "isotonic" maps to the expected rating, | |
| which keeps fewer, surer items in the high band.""" | |
| if how == "isotonic": | |
| from sklearn.isotonic import IsotonicRegression | |
| iso = IsotonicRegression(y_min=0, y_max=100, out_of_bounds="clip").fit(pred, y) | |
| return iso.X_thresholds_, iso.y_thresholds_ | |
| q = np.linspace(0, 1, 41) | |
| cx, cy = np.quantile(pred, q), np.quantile(y, q) | |
| keep = np.r_[True, np.diff(cx) > 1e-9] # np.interp needs increasing x | |
| return cx[keep], cy[keep] | |
| def band(v, cfg): | |
| return "high" if v >= cfg["bands"]["high"] else "medium" if v >= cfg["bands"]["medium"] else "low" | |
| def nested(X, y, alphas, cfg, reps=5, how="quantile"): | |
| """Outer 5-fold: fit ridge and its calibration on the training part only, score the held-out part.""" | |
| from scipy.stats import spearmanr | |
| from sklearn.model_selection import KFold | |
| rho, prec, rec, n_hi = [], [], [], [] | |
| for seed in range(reps): | |
| out = np.zeros(len(y)) | |
| for tr, te in KFold(5, shuffle=True, random_state=100 + seed).split(X): | |
| cx, cy = calibrator(oof(X[tr], y[tr], alphas, seed), y[tr], how) | |
| out[te] = np.interp(fit(X[tr], y[tr], alphas).predict(X[te]), cx, cy) | |
| rho.append(spearmanr(out, y)[0]) | |
| hp = np.array([band(v, cfg) == "high" for v in out]) | |
| ht = np.array([band(v, cfg) == "high" for v in y]) | |
| prec.append((hp & ht).sum() / max(1, hp.sum())) | |
| rec.append((hp & ht).sum() / max(1, ht.sum())) | |
| n_hi.append(int(hp.sum())) | |
| return {"spearman": round(float(np.mean(rho)), 3), "spearman_sd": round(float(np.std(rho)), 3), | |
| "high_precision": round(float(np.mean(prec)), 2), "high_recall": round(float(np.mean(rec)), 2), | |
| "high_items": round(float(np.mean(n_hi)), 1)} | |
| def main(argv=None) -> int: | |
| from scipy.stats import spearmanr | |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) | |
| ap.add_argument("--items", required=True) | |
| ap.add_argument("--ratings", required=True) | |
| ap.add_argument("--out", default=str(scanner.HERE / "head.json")) | |
| ap.add_argument("--models", default=None, help="fastembed cache directory") | |
| ap.add_argument("--calibration", choices=("quantile", "isotonic"), default="quantile") | |
| a = ap.parse_args(argv) | |
| cfg = scanner.load_config() | |
| items = {i["id"]: i for i in json.load(open(a.items))} | |
| ratings = json.load(open(a.ratings)) | |
| ids = sorted(k for k in ratings if k in items) | |
| y = np.array([mean_rating(ratings[k]) for k in ids]) | |
| sources = sorted({items[k]["source"] for k in ids}) | |
| kinds = sorted({items[k]["kind"] for k in ids}) | |
| M = np.array([scanner.Head.features(items[k], sources, kinds) for k in ids], "float32") | |
| from fastembed import TextEmbedding | |
| model = TextEmbedding(cfg["semantic"]["model"], cache_dir=a.models) | |
| texts = [scanner.embed_text(cfg, items[k]["title"], items[k].get("summary") or "") for k in ids] | |
| E = np.array(list(model.embed(texts, batch_size=64)), "float32") | |
| E /= np.maximum(np.linalg.norm(E, axis=1, keepdims=True), 1e-9) | |
| alphas = np.logspace(-2, 3, 20) | |
| rule = np.array([items[k]["score"] for k in ids]) | |
| report = {"rule_score": round(float(spearmanr(rule, y)[0]), 3), | |
| "keyword": round(float(spearmanr([items[k]["keyword"] for k in ids], y)[0]), 3)} | |
| out = {"version": 1, "trained": dt.date.today().isoformat(), "items": len(ids), "target": "mean blind rating 0..100", | |
| "embedding": cfg["semantic"]["model"], "meta_features": scanner.META_FEATURES, "sources": sources, | |
| "kinds": kinds, "calibration": a.calibration} | |
| for name, X in (("full", np.hstack([E, M])), ("meta", M)): | |
| report[name] = nested(X, y, alphas, cfg, how=a.calibration) | |
| r = fit(X, y, alphas) | |
| cx, cy = calibrator(oof(X, y, alphas, 0), y, a.calibration) | |
| coef = r.coef_.astype(float) | |
| h = {"alpha": float(r.alpha_), "intercept": round(float(r.intercept_), 6), | |
| "coef_meta": [round(c, 6) for c in coef[-M.shape[1]:]], | |
| "cal_x": [round(float(v), 4) for v in cx], "cal_y": [round(float(v), 2) for v in cy]} | |
| if name == "full": | |
| h["coef_emb"] = [round(c, 6) for c in coef[:E.shape[1]]] | |
| out[name] = h | |
| rule_hi = np.array([band(v, cfg) == "high" for v in rule]) | |
| true_hi = np.array([band(v, cfg) == "high" for v in y]) | |
| report["rule_high_precision"] = round(float((rule_hi & true_hi).sum() / max(1, rule_hi.sum())), 2) | |
| report["rule_high_recall"] = round(float((rule_hi & true_hi).sum() / max(1, true_hi.sum())), 2) | |
| report["rule_high_items"] = int(rule_hi.sum()) | |
| report["rated_high_items"] = int(true_hi.sum()) | |
| out["evaluation"] = report | |
| with open(a.out, "w") as f: | |
| json.dump(out, f, separators=(",", ":")) | |
| f.write("\n") | |
| print(json.dumps(report, indent=1)) | |
| return 0 | |
| if __name__ == "__main__": | |
| raise SystemExit(main()) | |