Constrained decoding

[viktor-ferenczi]

Changes:

• Added allowed_token_ids to SamplingParams (configuration)
• Enforced allowed_token_ids in Sampler.forward by squashing the logits of disallowed tokens

It allows the user to generate only specific tokens.

Please note, that it is the caller's responsibility to add the EOS and and additional stop tokens to the list of allowed tokens, but this is required only for open ended generations (not limited to 1 or a few tokens). It may be error prone, so we may want to add them automatically to the allowed tokens list.

Idea is to make a separate call for each segment of the generation which has different allowed_token_ids. The tokens known for sure can be efficiently "skipped" by appending them at the end of prompt for the next call (segment). The end of segments can be detected by adding them temporarily to stop_token_ids or by detecting them on the fly from a streaming generation. It gives the caller maximum control over the schema.

TODO:

• Write an example.
• Add support for regexp based constraint on the text returned (prompt+generated or generated only). It would alleviate the need for multiple generation calls in most cases, which is also more friendly with REST API calls. Regexp validation is the basis of the outlines library, study that library and figure out how can we implement this efficiently.
• Consider OpenAI API integration: Unsure whether we should do it or how to do it.

Constrained generation libraries we may want to provide adapters for:

• outlines
• guidance
• jsonformer

The adapters could be separate libraries or just examples. They should go into separate PRs (or repo).

Issue: #288

[jpeig]

Did you check out LMQL?

[viktor-ferenczi]

Did you check out LMQL?

Please give me a link or some hint where to find it.

[iiLaurens]

Did you check out LMQL?

Please give me a link or some hint where to find it.

https://lmql.ai/

[viktor-ferenczi]

LMQL sounds like a great library to support. I was thinking about regex, but that's very crude and difficult to work with.

[viktor-ferenczi]

See the related LMQL ticket

I guess best would be to work together with them to finish LMQL support.

[DarkLight1337]

Closing as superseded by #8252.