mixle.task.calibrate module

CalibratedTaskModel wraps a task model in conformal answer sets.

A distilled student classifies by argmax over a softmax, but the softmax value alone is not a coverage guarantee. Conformal prediction adds the serving contract: 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 under the usual exchangeability assumption.

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 used by the cost model. Conformal coverage is marginal, and a softmax still cannot 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]

Return labels in the probability-vector order used by the adapter.

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]

Return the conformal label set for one input.

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]

Return local labels or ESCALATE for a batch of inputs.

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