olmo_tap.inference.poe_demo_no_kv

Implements the Spec-Decode PoE method detailed here: https://www.overleaf.com/read/cjnxtpqjbbcx#4f6a0e This provides a security guarantee that no harmful token is ever sampled provided there exists at least 1 honest head in the jury which assigns negligible probability mass to the harmful token.

TODO: This is for demonstration purposes only and does not use KV cache.

Functions

poe_generate_visual_diff(model, tokenizer, ...)

olmo_tap.inference.poe_demo_no_kv.poe_generate_visual_diff(model: HydraTransformer, tokenizer: PreTrainedTokenizerBase, prompt_text: str, n_heads: int, gamma: int = 4, beta: float = 1.0, max_new_tokens: int = 200)