vesuvius-radar / train_head.py
SharpMeow's picture
Learned relevance score, Newest default, 10 per page, filter dropdowns, Contributors, easter eggs (Scrolls d993ec8)
f37fe0f verified
Raw History Blame Contribute Delete
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())