mixle.task.calibrated_generator module

CalibratedGenerator – conformal generation with honest abstention.

The generation-side sibling of CalibratedTaskModel. That class gates classification: it turns an uncalibrated softmax into conformal label sets and escalates on an ambiguous (empty or multi-label) set. This module gates generation: draw k candidates from any generator (a CallableLLM, a sampler, a beam), score each candidate with any mixle-scoreable model, and calibrate a conformal threshold on a held-out set so that serving the best-scored candidate carries the same finite-sample coverage guarantee – instead of always emitting whatever scored highest, regardless of whether the score means anything.

The trick is treating the k candidates the way CalibratedTaskModel treats classes: raw candidate scores are softmax-normalized per prompt into selection probabilities, and mixle.inference.conformal.conformal_label_threshold() / conformal_label_sets() calibrate + apply the exact LAC threshold the classification sibling uses. A singleton conformal set -> serve that candidate (covered at 1 - alpha); an empty or multi-candidate set -> honest “I’m not sure” (ABSTAIN) rather than a silent guess. ABSTAIN is None, the same sentinel value as mixle.task.calibrate.ESCALATE, so a Cascade built on a CalibratedGenerator escalates on abstention exactly the way it escalates on ESCALATE – no special-casing needed on the cascade side.

class CalibratedGenerator(generate, score, alpha=0.1, *, k=8, qhat=None, seed=0)[source]

Bases: object

Draw k scored candidates and serve the best one under a conformal accept-or-abstain guarantee.

Parameters:

  • generate (Callable[..., Sequence[Any]]) – generate(prompt, k) -> Sequence[candidate] (an rng keyword is passed if the callable accepts one; falls back to the two-argument form otherwise). Any generator that can draw k candidates for a prompt works: a wrapped CallableLLM sampled k times, a beam, a stochastic sampler.

  • score (Callable[[Any], float]) – score(candidate) -> float, any mixle-scoreable model. Higher is better; the score need not be a calibrated probability – that is exactly what conformal calibration fixes.

  • alpha (float) – target miscoverage rate.

  • k (int) – number of candidates to draw per prompt.

  • seed (int) – base seed for candidate draws; combined with the prompt (see _derive_seed()) so different prompts get different, but reproducible, draws.

  • qhat (float | None)

calibrate(prompts, is_correct, *, seed=None)[source]

Fit the conformal threshold on a held-out set of prompts, given a correctness oracle.

For each held-out prompt, k candidates are drawn and scored; the scores are softmax-normalized into per-prompt selection probabilities. The calibration score for that prompt is the probability mass landing on whichever candidate(s) is_correct(prompt, candidate) accepts (0 if none of the k draws is correct – an honest miss that the calibration prices in, the same way a held-out example whose true class never appears in a small candidate set prices into a wider conformal set). mixle.inference.conformal.conformal_label_threshold() on those scores gives the same LAC threshold CalibratedTaskModel calibrates for label sets.

Parameters:
Return type:

CalibratedGenerator

candidate_set(prompt, *, seed=None)[source]

The conformal candidate set for prompt – candidates whose selection probability clears the calibrated threshold. A singleton set is what serve() accepts; empty/multi abstains.

Parameters:

  • prompt (Any)

  • seed (int | None)

Return type:

list[Any]

serve(prompt, *, seed=None)[source]

Draw k candidates and return the best one if it conformally clears the threshold, else ABSTAIN. ABSTAIN is returned for both an empty set (nothing confident enough) and a multi-candidate set (genuinely ambiguous) – the same honest-uncertainty split CalibratedTaskModel uses for classification.

Parameters:

  • prompt (Any)

  • seed (int | None)

Return type:

Any

decide(prompt, *, seed=None)[source]

Alias for serve() with the same name as CalibratedTaskModel.decide(), so a CalibratedGenerator drops into Cascade unmodified.

Parameters:

  • prompt (Any)

  • seed (int | None)

Return type:

Any

abstention_rate(prompts, *, seed=None)[source]

Empirical fraction of prompts that would abstain – the generation analogue of CalibratedTaskModel.escalation_rate().

Parameters:
Return type:

float