Constrained decoding

[lucasavila00]

The public interface is the same used by TGI.

A grammar parameter that can have type "regex" or "yacc". (TGI has regex and json-schema)

Example 1 - Token healing

curl http://localhost:1234/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer EMPTY" \
-d '{
"model": "",
"grammar": {
    "type": "regex",
    "value": "Sure! (?s:.)*"
},
"messages": [
{
    "role": "user",
    "content": "Tell me a joke."
}
]
}'

Example 2 - Markdown List

curl http://localhost:1234/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer EMPTY" \
-d '{
"model": "",
"grammar": {
    "type": "regex",
    "value": "(- [^\n]*\n)+(- [^\n]*)(\n\n)?"
},
"messages": [
{
    "role": "user",
    "content": "Write a list of jokes. Return a markdown list where each item is a joke."
}
]
}'

[github-actions]

Code Metrics Report

  ───────────────────────────────────────────────────────────────────────────────
Language                 Files     Lines   Blanks  Comments     Code Complexity
───────────────────────────────────────────────────────────────────────────────
Rust                        47     16623     1176       706    14741        809
───────────────────────────────────────────────────────────────────────────────
Total                       47     16623     1176       706    14741        809
───────────────────────────────────────────────────────────────────────────────
Estimated Cost to Develop 55,563
Estimated Schedule Effort 10.201465 months
Estimated People Required 3.967364
───────────────────────────────────────────────────────────────────────────────
Processed 566260 bytes, 0.566 megabytes (SI)
───────────────────────────────────────────────────────────────────────────────
  

[lucasavila00]

@EricLBuehler What do you think about this approach?

The Windows build fails because the linker complains about missing functions that exist just for WASM. I'm not sure how to deal with this...

Also, I'm not a fan of using a library from Git. Should we vendor the code or something?

I haven't tried optimizing it yet, but on my system, this creates no overhead.

But it works 😃

[EricLBuehler]

This looks really great, thank you so much for adding it! However, I am concerned about the build failures on Windows, which I do want to support. An AICI issue: microsoft/aici#42 says that they do not support native Windows, which is a problem.

Perhaps we could support using aici_abi on non-Windows platforms, but use kalosm_sample's RegexParser otherwise or on all platforms for consistency? This is really great, and I wish that Windows support was available. Do you think you could modify it to use RegexParser? For example, they run the parser here and modify the logits.

[lucasavila00]

This looks really great, thank you so much for adding it! However, I am concerned about the build failures on Windows, which I do want to support. An AICI issue: microsoft/aici#42 says that they do not support native Windows, which is a problem.

I tried vendoring the code. It's 2k lines and by vendoring I was able to remove ~20 transient dependencies too.

Perhaps we could support using aici_abi on non-Windows platforms, but use kalosm_sample's RegexParser otherwise or on all platforms for consistency? This is really great, and I wish that Windows support was available. Do you think you could modify it to use RegexParser? For example, they run the parser here and modify the logits.

The parser is in "string-space" and we need to operate on "token-space". This would require a lot of glue code to do this conversion. I hope to make AICI code work for us, if it doesn't then I look into RegexParser

[EricLBuehler]

Thank you for pointing that out. Looks like the Windows build is passing; I think we are probably close to a merge once we do some testing, and you mark it as ready!

[lucasavila00]

@EricLBuehler I think this is ready for review and merge.

There is work left to do on optimization, and I haven't tested the yacc implementation yet.

But since the MR is already big, it might be a good point to merge.

[EricLBuehler]

I wrote some comments regarding error handling, other than that this looks excellent. Could you please add the example requests above into the examples/ directory?

[lucasavila00]

I added a Yacc example and tested that it works as expected.

[EricLBuehler]

Thank you!