Add beam search

[cartazio]

Hello! I'm trying to understand the capabilities of the current interface, and it seems like kwargs aren't passed to the underlying models, are they? I'd like to do stuff like beam search so i can get useful structured queries that are "good quality"

[cartazio]

to be clear: i'm also happy to write a patch for it if thats feasible

[brandonwillard]

to be clear: i'm also happy to write a patch for it if thats feasible

Patches/PRs are always welcome! In cases of large or complicated changes, we do prefer to have high-level discussions first, though, so that we make the best use of contributors' time.

[mattkindy]

Did anything come of this discussion privately?

[brandonwillard]

Did anything come of this discussion privately?

We've been working on improving the performance of the basic sampling steps recently and will start addressing alternative approaches soon after.

[rlouf]

Now that #416 has been merged we can consider implementing Beam Search in Outlines.

The samplers currently return next token ids. We will now need them to return the next token ids, as well as the "ancestor" of these token ids, i.e. the sequences to which they need to be appended. The list of ancestors will allow to update:

1. The sequences by adding new token ids to the ancestors; `token_ids = update_token_ids(token_ids, next_token_ids, ancestors)
2. The cumulative log-probability of the sequences: `logprobs = update_logprobs(logprobs, token_ids, logits, ancestors)
3. The KV-cache kv_cache = update_kv_cache(kv_cache, ancestors)
4. The FSMs and FSM states