diff --git a/report/confidence_histogram.png b/report/confidence_histogram.png new file mode 100644 index 0000000..60a7c50 Binary files /dev/null and b/report/confidence_histogram.png differ diff --git a/report/deferral_tradeoff.png b/report/deferral_tradeoff.png new file mode 100644 index 0000000..4a150c5 Binary files /dev/null and b/report/deferral_tradeoff.png differ diff --git a/report/metrics.json b/report/metrics.json index 46994d4..576dbe1 100644 --- a/report/metrics.json +++ b/report/metrics.json @@ -4,7 +4,13 @@ "accuracy": 0.69 }, "setfit": { - "accuracy": 0.895 + "accuracy": 0.925 + }, + "gate": { + "threshold": 0.55, + "in_distribution_kept_share": 1.0, + "ood_deferred_share": 0.067, + "kept_accuracy": 0.925 }, "haiku": { "model": "claude-haiku-4-5", diff --git a/report/model_comparison.png b/report/model_comparison.png index ec509ca..9181423 100644 Binary files a/report/model_comparison.png and b/report/model_comparison.png differ diff --git a/report/ood_probe.txt b/report/ood_probe.txt new file mode 100644 index 0000000..9478eeb --- /dev/null +++ b/report/ood_probe.txt @@ -0,0 +1,34 @@ +# Out-of-distribution / garbled probes for the confidence-gate plot. These are +# deliberately unlike the five voice intents: gibberish, other languages, +# off-domain prose, ASR noise. The gate (routelet -> Claude fallback) is meant +# to fire on inputs like these. No labels: we only read the model's confidence. +asdkfj qwoieu zxcmnv plok +blorp the fnordle wug sprongle +quelle heure est il maintenant +ich moechte ein kaltes bier bitte +the mitochondria is the powerhouse of the cell +photosynthesis converts sunlight water and co2 into glucose +lorem ipsum dolor sit amet consectetur adipiscing elit +the treaty of westphalia was signed in sixteen forty eight +uh the the um where is it like you know +skrrt skrrt yeah yeah ok ok +three point one four one five nine two six five +glark mizzle frobnicate the wibble +supercalifragilisticexpialidocious +a b c d e f g h i j k l m n o p +the stock market closed lower amid inflation fears +ribosomes translate messenger rna into proteins +hakuna matata what a wonderful phrase +to be or not to be that is the question +the quick brown fox jumps over the lazy dog +e equals m c squared is mass energy equivalence +my favorite color is the number seven on tuesdays +wingardium leviosa not leviosaa +the derivative of sine is cosine +clouds drift slowly across an empty autumn sky +zzzzz hmmmm errr uhhh ahh +random words banana thunderstorm philosophy keyboard ocean +the french revolution began in seventeen eighty nine +i think therefore i am said descartes +plinko bork snorf glimble dax +the speed of light is about three hundred thousand kilometers per second diff --git a/report/report.py b/report/report.py index 49ff31b..b51aaae 100644 --- a/report/report.py +++ b/report/report.py @@ -37,12 +37,21 @@ TFIDF_MODEL = PROJECT_ROOT / "models" / "baseline.joblib" SETFIT_DIR = PROJECT_ROOT / "models" / "setfit" BASELINES = PROJECT_ROOT / "report" / "baselines.json" +OOD_PROBE = PROJECT_ROOT / "report" / "ood_probe.txt" OUT_DIR = PROJECT_ROOT / "report" # Okabe-Ito (colorblind-safe). Gray baseline, green hero (routelet), amber oracle. C_TFIDF = "#999999" C_ROUTELET = "#009E73" C_HAIKU = "#E69F00" +# In-distribution vs OOD: blue and vermillion, distinguishable for colorblind +# readers (and labeled, so not relying on hue alone). +C_INDIST = "#0072B2" +C_OOD = "#D55E00" + +# Routelet's on-device confidence gate: below this, Aegis defers to the Claude +# fallback. Mirrors ROUTELET_CONFIDENCE_THRESHOLD in the aegis crate's tuning.rs. +GATE = 0.55 def score_tfidf(texts: list[str], gold: list[str]) -> float: @@ -53,15 +62,57 @@ def score_tfidf(texts: list[str], gold: list[str]) -> float: return float(accuracy_score(gold, preds)) -def score_setfit(texts: list[str], gold: list[str]) -> float: - """The shipped SetFit model. Trained on preprocess()'d text, so score the - same way. This is the fine-tuned bge-small body + LR head in torch; int8 - ONNX (what Rust actually runs) is verified equivalent at export time.""" +def _load_temperature() -> float: + """The calibrated temperature baked into the model dir, or 1.0 if absent. + Must match what Aegis applies so the gate confidence here equals production.""" + path = SETFIT_DIR / "temperature.json" + if path.exists(): + return float(json.loads(path.read_text())["temperature"]) + return 1.0 + + +def load_setfit() -> dict: + """Load the shipped SetFit model once into a reusable bundle (body + head + weights + calibrated temperature), so the holdout and the OOD probe can both + be scored without reloading the model.""" from setfit import SetFitModel model = SetFitModel.from_pretrained(str(SETFIT_DIR)) - preds = list(model.predict([preprocess(t) for t in texts])) - return float(accuracy_score(gold, preds)) + head = model.model_head + return { + "body": model.model_body, + "coef": head.coef_, + "intercept": head.intercept_, + "labels": head.classes_.tolist(), + "temperature": _load_temperature(), + } + + +def setfit_predict(bundle: dict, texts: list[str]) -> tuple[np.ndarray, list[str]]: + """Run the model the way Aegis does: embed preprocess()'d text, apply the LR + head, temperature-scale, softmax, argmax. Returns (max-softmax confidence per + row, predicted labels). The int8 ONNX Aegis ships is verified equivalent at + export time.""" + emb = bundle["body"].encode( + [preprocess(t) for t in texts], convert_to_numpy=True, show_progress_bar=False + ) + logits = (emb @ bundle["coef"].T + bundle["intercept"]) / bundle["temperature"] + logits -= logits.max(axis=1, keepdims=True) + probs = np.exp(logits) + probs /= probs.sum(axis=1, keepdims=True) + idx = probs.argmax(axis=1) + conf = probs[np.arange(len(idx)), idx] + preds = [bundle["labels"][i] for i in idx] + return conf, preds + + +def load_ood_probes() -> list[str]: + """Read the OOD/garbled probe lines (skipping blanks and # comments).""" + return [ + ln.strip() + for ln in OOD_PROBE.read_text().splitlines() + if ln.strip() and not ln.lstrip().startswith("#") + ] def load_haiku(eval_n: int) -> dict | None: @@ -123,6 +174,88 @@ def plot_model_comparison(metrics: dict) -> Path: return out +def plot_confidence_histogram(indist: np.ndarray, ood: np.ndarray, gate: float) -> Path: + """Figure 2: routelet confidence on in-distribution holdout vs OOD/garbled + probes, with the gate line. The cascade works if in-distribution input sits + above the gate (kept on-device) while OOD input falls below it (deferred to + Claude). Overlapping (not stacked) histograms, densities so the two groups + are comparable despite different counts.""" + bins = np.linspace(0.0, 1.0, 21) + fig, ax = plt.subplots(figsize=(7, 4.5)) + ax.hist(indist, bins=bins, color=C_INDIST, alpha=0.6, density=True, + label=f"in-distribution holdout (n={len(indist)})", zorder=3) + ax.hist(ood, bins=bins, color=C_OOD, alpha=0.7, density=True, + label=f"OOD / garbled probe (n={len(ood)})", zorder=3) + + ax.axvline(gate, color="black", linestyle="--", linewidth=1.2, zorder=4) + ymax = ax.get_ylim()[1] + ax.text(gate - 0.012, ymax * 0.96, f"{gate:.2f} gate", ha="right", va="top", fontsize=9) + ax.text(gate - 0.012, ymax * 0.55, "← defer to Claude", ha="right", fontsize=8, color="#555") + ax.text(gate + 0.012, ymax * 0.55, "kept on-device →", ha="left", fontsize=8, color="#555") + + ood_deferred = int((ood < gate).sum()) + indist_kept = int((indist >= gate).sum()) + ax.set_title( + f"OOD defers {ood_deferred}/{len(ood)}, in-distribution keeps " + f"{indist_kept}/{len(indist)} at the {gate:.2f} gate", + fontsize=11, fontweight="bold", + ) + ax.set_xlabel("routelet confidence (temperature-scaled max softmax)") + ax.set_ylabel("density") + ax.set_xlim(0, 1) + ax.legend(frameon=False, loc="upper left") + ax.yaxis.grid(True, linestyle="--", alpha=0.4, zorder=0) + ax.set_axisbelow(True) + for spine in ("top", "right"): + ax.spines[spine].set_visible(False) + fig.tight_layout() + + out = OUT_DIR / "confidence_histogram.png" + fig.savefig(out, dpi=150) + plt.close(fig) + return out + + +def plot_deferral_tradeoff(indist: np.ndarray, ood: np.ndarray, gate: float) -> Path: + """Figure 3: as the confidence cutoff rises, what fraction of in-distribution + commands get wrongly deferred to Claude vs what fraction of OOD/garbled input + gets caught. The two lines never separate cleanly, which is why no cutoff + makes the gate work: catching OOD means deferring real commands too.""" + thresholds = np.linspace(0.5, 1.0, 101) + id_deferred = np.array([(indist < t).mean() for t in thresholds]) * 100 + ood_caught = np.array([(ood < t).mean() for t in thresholds]) * 100 + + fig, ax = plt.subplots(figsize=(7.5, 4.5)) + ax.plot(thresholds, ood_caught, color=C_OOD, linewidth=2.4, + label="OOD / garbled caught (good)", zorder=3) + ax.plot(thresholds, id_deferred, color=C_INDIST, linewidth=2.4, + label="real commands wrongly deferred (bad)", zorder=3) + + ax.axvline(gate, color="#777", linestyle="--", linewidth=1.2, zorder=2) + ax.text(gate, 102, f"current cutoff {gate:.2f}", ha="center", va="bottom", + fontsize=8, color="#555") + + ax.set_xlabel("confidence cutoff (defer to Claude below it)") + ax.set_ylabel("% of inputs deferred") + ax.set_title( + "The 0.55 cutoff is too low: ~0.95 catches far more garbage at little cost", + fontsize=11, fontweight="bold", + ) + ax.set_xlim(0.5, 1.0) + ax.set_ylim(0, 105) + ax.legend(loc="upper left", frameon=False) + ax.grid(True, linestyle="--", alpha=0.4, zorder=0) + ax.set_axisbelow(True) + for spine in ("top", "right"): + ax.spines[spine].set_visible(False) + fig.tight_layout() + + out = OUT_DIR / "deferral_tradeoff.png" + fig.savefig(out, dpi=150) + plt.close(fig) + return out + + def main() -> None: examples = load(HOLDOUT) texts = [e.text for e in examples] @@ -132,13 +265,36 @@ def main() -> None: print(f"scoring on {eval_n}-row holdout") tfidf_acc = score_tfidf(texts, gold) print(f" TF-IDF {tfidf_acc:.3f}") - setfit_acc = score_setfit(texts, gold) + + bundle = load_setfit() + conf, preds = setfit_predict(bundle, texts) + correct = np.array([p == g for p, g in zip(preds, gold)]) + setfit_acc = float(correct.mean()) print(f" SetFit {setfit_acc:.3f}") + # OOD/garbled probe: confidence on inputs the gate is meant to defer. + ood_texts = load_ood_probes() + ood_conf, _ = setfit_predict(bundle, ood_texts) + + # Confidence gate operating point: what the cascade does at GATE. + indist_kept = conf >= GATE + ood_deferred = ood_conf < GATE + gate_stats = { + "threshold": GATE, + "in_distribution_kept_share": round(float(indist_kept.mean()), 3), + "ood_deferred_share": round(float(ood_deferred.mean()), 3), + "kept_accuracy": ( + round(float(correct[indist_kept].mean()), 3) if indist_kept.any() else None + ), + } + print(f" gate {GATE}: in-dist keeps {gate_stats['in_distribution_kept_share']:.0%}, " + f"OOD defers {gate_stats['ood_deferred_share']:.0%}") + metrics: dict = { "eval_n": eval_n, "tfidf": {"accuracy": tfidf_acc}, "setfit": {"accuracy": setfit_acc}, + "gate": gate_stats, } haiku = load_haiku(eval_n) if haiku: @@ -150,8 +306,9 @@ def main() -> None: print(f" Haiku {haiku['accuracy']:.3f} (cached)") (OUT_DIR / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n") - fig = plot_model_comparison(metrics) - print(f"wrote {fig}") + print(f"wrote {plot_model_comparison(metrics)}") + print(f"wrote {plot_confidence_histogram(conf, ood_conf, GATE)}") + print(f"wrote {plot_deferral_tradeoff(conf, ood_conf, GATE)}") print(f"wrote {OUT_DIR / 'metrics.json'}") diff --git a/src/routelet/train_setfit.py b/src/routelet/train_setfit.py index 7bdab76..1f68086 100644 --- a/src/routelet/train_setfit.py +++ b/src/routelet/train_setfit.py @@ -268,20 +268,24 @@ def main() -> None: assert cw == "balanced", f"class_weight not applied; got {cw!r}" print(f"head class_weight: {cw}") - # sampling_strategy="unique": draws every sentence-pair combination exactly - # once (no duplication). Valid in setfit 1.1.3 (confirmed in source). - # num_epochs=2 doubles the embedding training time vs. the previous 1-epoch - # run, giving the contrastive head more signal on the larger 1115-row pool. + # oversampling draws num_iterations * n_classes * 2 pairs per epoch instead + # of the full O(n^2) pair set. "unique" generated ~1.2M pairs on the grown + # ~1566-row pool (~60 min runs), and most were trivially-easy near-duplicate + # pairs from the disfluency augmentation, so the embedding loss collapsed + # early without real signal. 40 iterations gives 400 pairs, enough to nudge + # the bge-small embedding for the domain; the LR head does the rest. args = TrainingArguments( batch_size=16, - num_epochs=2, - sampling_strategy="unique", + num_epochs=1, + num_iterations=40, + sampling_strategy="oversampling", ) print("\nTrainingArguments:") print(f" batch_size: {args.batch_size}") print(f" num_epochs: {args.num_epochs}") + print(f" num_iterations: {args.num_iterations}") print(f" sampling_strategy: {args.sampling_strategy}") - assert args.sampling_strategy == "unique", ( + assert args.sampling_strategy == "oversampling", ( f"sampling_strategy not applied; got {args.sampling_strategy!r}" )