mixle.task.calibrated_generator module

CalibratedGenerator – certified selective generation with 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 label set. Open-ended generation has no fixed label space: candidates are newly sampled and their identities change by prompt, so applying label-set conformal prediction to candidate slots would not provide classification coverage.

This module instead certifies one fixed measurable outcome: whether the top-scored generated candidate is correct when its score clears a threshold. Calibration is split into threshold proposal and independent certification halves. Candidate thresholds are proposed on the first half; the second half supplies simultaneous Bonferroni-corrected exact binomial upper bounds, allowing the most permissive threshold whose accepted-error upper bound is at most alpha to be selected without reusing its evidence. The guarantee is a confidence statement about selective risk under the usual i.i.d./exchangeability and stable-generation assumptions, not conformal coverage of changing strings.

An uncertified prompt returns ABSTAIN (None), the same sentinel as mixle.task.calibrate.ESCALATE, so cascades can escalate without special casing.

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

Bases: object

Draw k candidates and serve the best under a certified selective-risk gate.

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; calibration thresholds its ordering statistic directly.

  • alpha (float) – maximum certified error rate among accepted candidates.

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

  • qhat (float | None) – optional manually supplied acceptance threshold. This remains callable but has no risk certificate; only calibrate() populates risk_receipt.

  • confidence (float) – simultaneous confidence level for the held-out risk certificate.

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

calibrate(prompts, is_correct, *, seed=None, outcomes='per-prompt', sampling='constructed')[source]

Certify a held-out accepted-error threshold using an explicit correctness oracle.

The trailing half proposes score-margin thresholds; the independent leading half evaluates every proposal with simultaneous exact binomial bounds. Only the certification half consults is_correct, once per row in order, so a metered or side-effecting oracle is not spent on rows that cannot reach the bound (MXR-080-1849). If no nonempty accepted subset certifies risk <= alpha, the threshold is +inf and serving abstains everywhere.

Scope of the certificate: it covers exactly the SERVED stochastic policy – candidate draws seeded from the generator’s own seed and the prompt alone, the same derivation candidate_set()/serve() use (STAT-RR17-07: certifying under a per-row schedule measured 1/150 errors and bounded risk at 0.0366 while the served policy produced 1000/1000 errors). Assuming the calibration prompts and future queries are i.i.d. draws from the same distribution, with probability at least confidence over the calibration draw, the TRUE accepted-slice error at the deployed threshold is <= alpha (the per-threshold binomial bounds are Bonferroni-corrected across the proposal family, so the selection of the loosest passing threshold stays covered). Distribution shift voids the statement silently; re-certify on drifted traffic.

WHAT the bound is a bound ON depends on how duplicated prompts got into the stream, and only the caller knows that – sampling is that declaration (the same premise-declaration contract as outcomes):

  • sampling="constructed" (fail-safe default): the prompt list is a curated set, so under outcomes="per-prompt" duplicates are redundant copies of one Bernoulli and COLLAPSE to one certification row (300 copies of one always-correct prompt certify from n = 1, bound 0.975 – STAT-RR18-04). The certified estimand is then UNIFORM over the DISTINCT certification prompts. It is NOT traffic-weighted: measured on 40%-heavy i.i.d. traffic whose heavy prompt always errs, collapsing certified error_upper 0.080 while the served traffic risk was 0.37-0.41, every trial (the pass-19 blocker). If your serving mix repeats prompts, this default’s certificate does not cover it – declare the sampling truthfully instead.

  • sampling="iid-traffic": the caller’s recorded assertion that the calibration rows are an i.i.d. draw of the serving traffic itself. Row error indicators are then i.i.d. Bernoulli(traffic accepted-error rate) EVEN under per-prompt outcomes – the randomness is the prompt draw, and a duplicate’s multiplicity IS its traffic weight – so rows all count and the bound covers the TRAFFIC-WEIGHTED accepted-error rate serving experiences. On the measurement above, this declaration refuses to certify alpha = 0.15 (the heavy errors are ~40% of rows), which is the correct answer.

Disagreeing duplicate verdicts are refused under outcomes="per-prompt" in both sampling modes (a deterministic decision with a fixed per-prompt outcome cannot differ between copies); the receipt records both declarations and names the certified estimand.

Parameters:
Return type:

CalibratedGenerator

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

Return [best_candidate] when its calibrated statistic clears the risk gate, else [].

Parameters:

  • prompt (Any)

  • seed (int | None)

Return type:

list[Any]

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

Return the best candidate only when the certified risk gate accepts, else ABSTAIN.

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

smallest_certifiable_calibration_set(alpha, *, confidence=0.95, thresholds_tested=2, limit=100_000)[source]

Smallest calibrate(...) set size whose risk certificate can reach alpha at all.

CalibratedGenerator.calibrate() splits the set in half and certifies on the leading half with a one-sided Clopper-Pearson bound, Bonferroni-corrected across the thresholds it tests. That bound has a floor set by the certification count alone: with c certification rows and zero observed errors it is still 1 - (tail)**(1/c). Below the size returned here, alpha sits under that floor, no threshold can ever be eligible, qhat stays +inf and serving abstains on every input – for a reason that has nothing to do with the model. Sizing a calibration set from this avoids mistaking a structurally impossible target for a model that needs tuning.

thresholds_tested is the Bonferroni divisor; calibrate() tests one threshold per distinct proposal-split statistic plus -inf, so 2 is the floor and the realistic best case. Pass the number you expect if the statistic takes many distinct values, since more thresholds tighten the tail.

Parameters:
Return type:

int