mixle.enumeration.rescore module

Speculative enumeration: build the index with a cheap DRAFT model, score results with the TARGET.

Speculative decoding’s economics applied to enumeration. Building any autoregressive index costs one forward per live prefix – prohibitive when the model is a large transformer. But the ordering work (which sequences are near a rank/threshold) tolerates approximation, while the scores must be the real model’s. So: let a cheap draft (an n-gram, a distilled student, a quantized twin) pay for the tree or envelope build, and touch the target only for the sequences a query actually returns – one batched teacher-forcing forward for all of them (AutoregressiveEnumerable.score_sequences()).

Contract (honest): every returned log_prob is the target’s exact score. The order is draft-approximate, repaired locally by window reranking: top_k(k) / slice pull k + rerank_window draft-ordered candidates, rescore them all with the target in one batch, and sort by target score. That is exact whenever no unpulled sequence out-scores the returned ones – guaranteed if the draft-to-target log-prob gap is globally bounded by assumed_gap and the window edge clears it (the certified flag); without an assumed bound the observed gap diagnostic is reported and the certificate is honestly None.

class RescoredIndex(draft_index, target, *, rerank_window=64, assumed_gap=None)[source]

Bases: object

Draft-ordered, target-scored enumeration with window reranking.

Parameters:

  • draft_index (Any) – any index with unrank(i) -> (sequence, draft_log_prob) – a SeekIndex over a cheap AutoregressiveEnumerable, an AREnvelopeIndex, or anything equivalent.

  • target (Any) – the expensive model – an AutoregressiveEnumerable (its score_sequences() batch scorer is used) or a bare callable [seqs] -> log_probs.

  • rerank_window (int) – extra draft candidates pulled around a query and reranked by target score. Larger = more robust to draft/target disagreement, one batched forward either way.

  • assumed_gap (float | None) – optional global bound on |target_lp - draft_lp| (nats). When supplied, results carry a sound certified verdict; otherwise certified is None and the observed gap is reported as a diagnostic.

top_k(k)[source]

The k best sequences by TARGET score among the k + rerank_window draft head.

Returns {"items": [(seq, target_lp), ...], "certified": bool | None, "gap": float} – target-exact scores, draft+window-approximate completeness (see the class docstring).

Parameters:

k (int)

Return type:

dict[str, Any]

slice(start, k)[source]

Target-reranked [start, start + k) slice of the pulled start + k + rerank_window head.

Same semantics as top_k(): order within the pulled set is target-exact; the certificate covers whether an unpulled sequence could belong in (or before) the slice.

Parameters:
Return type:

dict[str, Any]

unrank(i)[source]

The draft’s rank-i sequence with the TARGET’s exact log-probability.

The rank coordinate is the draft’s (no reranking): the cheap random-access primitive. Use top_k() / slice() when local target-order matters.

Parameters:

i (int)

Return type:

tuple[tuple, float]