diff --git a/src/tracer/fit/pipeline.py b/src/tracer/fit/pipeline.py index 4b0f056..f7e3a8b 100644 --- a/src/tracer/fit/pipeline.py +++ b/src/tracer/fit/pipeline.py @@ -233,7 +233,11 @@ def _calibrate_threshold(scores, preds, y_teacher, target_ta, alpha=0.1, min_acc def _predict(clf, X): probs = clf.predict_proba(X) - preds = probs.argmax(axis=1).astype(int) + idx = probs.argmax(axis=1) + # predict_proba columns follow clf.classes_. A stage trained on a subset of + # the label space (the RSB residual stage) must map them back to label ids. + classes = getattr(clf, "classes_", None) + preds = (np.asarray(classes)[idx] if classes is not None else idx).astype(int) return preds, probs diff --git a/tests/test_stage_labels.py b/tests/test_stage_labels.py new file mode 100644 index 0000000..4c2fa3d --- /dev/null +++ b/tests/test_stage_labels.py @@ -0,0 +1,92 @@ +"""A stage's predictions must be teacher label ids, not predict_proba columns. + +The RSB second stage is trained only on the rows the first stage rejected, so +its surrogate usually sees a subset of the label space (e.g. the easy class the +first stage fully absorbs is missing). For such a model, column ``i`` of +``predict_proba`` is ``clf.classes_[i]``, not label id ``i``. +""" +import json +import tempfile +from pathlib import Path + +import numpy as np +from sklearn.linear_model import LogisticRegression + +from tracer.fit.pipeline import _predict, apply_stage + +LINEAR_ONLY = ("gbt", "rf", "et", "xgb", "dt", "mlp_1h", "mlp_2h") + + +def test_predict_maps_proba_columns_to_class_labels(): + rng = np.random.RandomState(0) + # Trained on label ids {1, 3} only (0 and 2 absent from this subset). + X = np.vstack([rng.randn(50, 2) - 4, rng.randn(50, 2) + 4]) + y = np.r_[np.full(50, 1), np.full(50, 3)] + clf = LogisticRegression().fit(X, y) + + preds, probs = _predict(clf, X) + assert probs.shape == (100, 2) + assert set(np.unique(preds)) == {1, 3} + assert (preds == y).mean() > 0.99 + + stage = {"clf": clf, "accept_all": True} + stage_preds, _, _ = apply_stage(stage, X) + assert set(np.unique(stage_preds)) == {1, 3} + + +def _residual_traces(tmp, n=4000): + """Class 0 is trivially separable, so stage 1 accepts every class-0 row. + + Classes 1-3 come in two groups. The 'easy' group carries a per-class + marker; the 'hard' group encodes the class along a direction that the easy + group uses the opposite way, which a global linear surrogate cannot + resolve. Those rows are rejected by stage 1 and are exactly what the RSB + residual stage should pick up, with only labels {1, 2, 3} in its pool. + """ + rng = np.random.RandomState(0) + X, y = [], [] + for _ in range(n): + c = rng.randint(0, 4) + x = rng.randn(6) * 0.5 + if c == 0: + x[0] += 10 + elif rng.rand() > 0.3: + x[c] += 4 + x[5] -= (c - 2) * 8 + x[4] -= 3 + else: + x[4] += 3 + x[5] += (c - 2) * 2 + X.append(x) + y.append(c) + X = np.asarray(X, dtype=np.float32) + names = ["cls_0", "cls_1", "cls_2", "cls_3"] + path = Path(tmp) / "traces.jsonl" + with path.open("w") as f: + for i, c in enumerate(y): + f.write(json.dumps({"input": f"t{i}", "teacher": names[c]}) + "\n") + return path, X, np.asarray([names[c] for c in y]) + + +def test_rsb_residual_stage_routes_with_teacher_labels(): + from tracer.api import fit, load_router + from tracer.config import FitConfig + + with tempfile.TemporaryDirectory() as tmp: + path, X, teacher = _residual_traces(tmp) + out = Path(tmp) / ".tracer" + cfg = FitConfig(target_teacher_agreement=0.95, frontier_targets=(0.95,), + skip_candidates=LINEAR_ONLY, verbose=False) + result = fit(path, artifact_dir=out, embeddings=X, config=cfg) + + # The residual stage is certified, so RSB beats single-stage L2D. + assert result.manifest.selected_method == "rsb" + assert result.manifest.coverage_cal > 0.95 + + router = load_router(out) + batch = router.predict_batch(X) + stage2 = np.asarray(batch["stage_id"]) == 1 + assert stage2.sum() > 100 + labels = np.asarray(batch["labels"], dtype=object)[stage2] + # Stage-2 rows carry the class they belong to, not a shifted label. + assert (labels == teacher[stage2]).mean() >= 0.9