mixle.task.calibrate module

CalibratedTaskModel – wrap a task model in conformal sets so its escalate-or-answer decision is honest.

A distilled student classifies by argmax over a softmax that is not a describable random process: the numbers sum to 1, but a confident-looking 0.97 carries no guarantee. Gating a cost-aware cascade on that number is fiction. Conformal prediction fixes it without pretending the softmax is generative: on a held-out calibration set it learns a score threshold (mixle.inference.conformal.conformal_label_threshold()) such that the prediction set covers the true label with probability >= 1 - alpha.

The decision rule the cascade and the cost model consume:

  • singleton set -> answer locally (covered at 1 - alpha);

  • empty or multi-label set -> escalate to the expensive teacher/frontier (genuinely ambiguous).

escalation_rate is the empirical p_escalate – the number that makes “expected $/request” real rather than a vibe. Conformal coverage is marginal, and a softmax still can’t see true OOD; a generative-density gate (mixle.task.density) covers that residual. Calibration persists in the artifact, so a loaded model decides identically in a fresh process.

class CalibratedTaskModel(task, *, alpha=0.1, qhat=None, density_gate=None)[source]

Bases: object

A TaskModel plus a conformal threshold: predicts label sets and decides answer-vs-escalate.

Parameters:

  • task (TaskModel)

  • alpha (float)

  • qhat (float | None)

  • density_gate (Any)

property labels: list[str]
calibrate(texts, teacher_labels)[source]

Set the conformal threshold from held-out (texts, teacher_labels) for 1 - alpha set coverage.

Parameters:
Return type:

CalibratedTaskModel

predict_sets(texts)[source]

Conformal label set per input (the classes whose score clears the calibrated threshold).

Parameters:

texts (Sequence[Any])

Return type:

list[list[str]]

predict_set(text)[source]
Parameters:

text (Any)

Return type:

list[str]

decide(text)[source]

Return the label if the input is a confident, in-distribution singleton, else ESCALATE (None).

Parameters:

text (Any)

Return type:

Any

batch_decide(texts)[source]
Parameters:

texts (Sequence[Any])

Return type:

list[Any]

escalation_rate(texts)[source]

Empirical p_escalate – the fraction of inputs escalated (ambiguous set or, if gated, OOD).

Parameters:

texts (Sequence[Any])

Return type:

float

save(path)[source]

Persist the underlying model, the calibration (alpha, qhat), and any density gate in the artifact.

qhat can legitimately be +inf (a small calibration set / tight alpha: too little data to admit any confident singleton, so every input escalates). That is a real, callable threshold, so it is persisted as the JSON-safe sentinel "inf" and reloads back to float('inf') – a loaded model stays callable instead of raising “call calibrate”.

Parameters:

path (str)

Return type:

str

classmethod load(path, *, device='cpu')[source]

Rebuild a calibrated model (with its density gate, if any) from an artifact; decisions match exactly.

Parameters:
Return type:

CalibratedTaskModel