Add grammar-based sampling (for webui, llamacpp, and koboldcpp)

[cpacker]

Closes #273

• Adds koboldcpp and llama.cpp backend support
• Adds grammar-based sampling support for koboldcpp, llama.cpp, and webui backends
• Adds basic grammar (json.gbnf) and MemGPT-specific grammar (json_func_calls_with_inner_thoughts.gbnf)
  • Thank you @Drake-AI!
• If you are using a backend that supports grammars, will run airoboros wrapper w/ grammar by default

---

Adds grammar-based sampling to prevent JSON-related parsing errors when using local LLMs, see llama.cpp.

Example

Starting the server (llama.cpp on MacOS):

llama.cpp % ./server -m ~/models/TheBloke/dolphin-2.2.1-mistral-7B-GGUF/dolphin-2.2.1-mistral-7b.Q6_K.gguf -c 8000

Starting the server (koboldcpp on MacOS):

koboldcpp % ./koboldcpp.py ~/models/TheBloke/dolphin-2.2.1-mistral-7B-GGUF/dolphin-2.2.1-mistral-7b.Q6_K.gguf --contextsize 8192

Specifying a wrapper with a grammar:

python main.py --model airoboros-l2-70b-2.1-grammar --first

[Drake-AI]

I'm working on a grammar to avoid errors, it is in progress but it looks like this, includes names of functions, names and types of params:

root ::= Function
Function ::= SendMessage | PauseHeartbeats
SendMessage ::= "{" ws ""function":" ws ""send_message"," ws ""params":" ws SendMessageParams "}"
PauseHeartbeats ::= "{" ws ""function":" ws ""pause_heartbeats"," ws ""params":" ws PauseHeartbeatsParams "}"
SendMessageParams ::= "{}" | "{" ws Message ("," ws Message)* ws "}"
PauseHeartbeatsParams ::= "{}" | "{" ws Minutes ("," ws Minutes)* ws "}"
Message ::= ""inner_thoughts":" ws string | ""message":" ws string | ""request_heartbeat":" ws boolean
Minutes ::= ""minutes":" ws number
string ::= """ ([^"]) """
boolean ::= "true" | "false"
ws ::= [ \t\n]
number ::= [0-9]+ "."? [0-9]*

With this should be possible to use any model without finetuning.

[Drake-AI]

It is working now with all the basic functions, easy to expand to extra functions. Feel free to use it or adapt to your needs.

root ::= Function
Function ::= SendMessage | PauseHeartbeats | CoreMemoryAppend | CoreMemoryReplace | ConversationSearch | ConversationSearchDate | ArchivalMemoryInsert | ArchivalMemorySearch
SendMessage ::= "{"   ws   "\"function\":"   ws   "\"send_message\","   ws   "\"params\":"   ws   SendMessageParams   "}"
PauseHeartbeats ::= "{"   ws   "\"function\":"   ws   "\"pause_heartbeats\","   ws   "\"params\":"   ws   PauseHeartbeatsParams   "}"
CoreMemoryAppend ::= "{"   ws   "\"function\":"   ws   "\"core_memory_append\","   ws   "\"params\":"   ws   CoreMemoryAppendParams   "}"
CoreMemoryReplace ::= "{"   ws   "\"function\":"   ws   "\"core_memory_replace\","   ws   "\"params\":"   ws   CoreMemoryReplaceParams   "}"
ConversationSearch  ::= "{"   ws   "\"function\":"   ws   "\"conversation_search\","   ws   "\"params\":"   ws   ConversationSearchParams   "}"
ConversationSearchDate  ::= "{"   ws   "\"function\":"   ws   "\"conversation_search_date\","   ws   "\"params\":"   ws   ConversationSearchDateParams   "}"
ArchivalMemoryInsert  ::= "{"   ws   "\"function\":"   ws   "\"archival_memory_insert\","   ws   "\"params\":"   ws   ArchivalMemoryInsertParams   "}"
ArchivalMemorySearch  ::= "{"   ws   "\"function\":"   ws   "\"archival_memory_search\","   ws   "\"params\":"   ws   ArchivalMemorySearchParams   "}"
SendMessageParams ::= "{"   ws   InnerThoughtsParam   ","   ws   "\"message\":"   ws   string   ws   "}"
PauseHeartbeatsParams ::= "{"   ws   InnerThoughtsParam   ","      ws   "\"minutes\":"   ws   number   ws   "}"
CoreMemoryAppendParams ::= "{"   ws   InnerThoughtsParam   ","      ws   "\"name\":"   ws   namestring   ","   ws   "\"content\":"   ws   string   ws   ","   ws   RequestHeartbeatParam   ws   "}"
CoreMemoryReplaceParams ::= "{"   ws   InnerThoughtsParam   ","      ws   "\"name\":"   ws   namestring   ","   ws   "\"old_content\":"   ws   string   ","   ws   "\"new_content\":"   ws   string   ws   ","   ws   RequestHeartbeatParam   ws   "}"
ConversationSearchParams ::= "{"   ws   InnerThoughtsParam   ","      ws   "\"query\":"   ws   string   ws   ","   ws   "\"page\":"   ws   number   ws   ","   ws   RequestHeartbeatParam   ws   "}"
ConversationSearchDateParams ::= "{"   ws   InnerThoughtsParam   ","      ws   "\"start_date\":"   ws   string   ws   ","      ws   "\"end_date\":"   ws   string   ws   ","   ws   "\"page\":"   ws   number   ws   ","   ws   RequestHeartbeatParam   ws   "}"
ArchivalMemoryInsertParams ::= "{"   ws   InnerThoughtsParam    ","   ws   "\"content\":"   ws   string   ws   ","   ws   RequestHeartbeatParam   ws   "}"
ArchivalMemorySearchParams ::= "{"   ws   InnerThoughtsParam   ","  ws   "\"query\":"   ws   string   ws   ","   ws   "\"page\":"   ws   number   ws   ","   ws   RequestHeartbeatParam   ws   "}"
InnerThoughtsParam ::= "\"inner_thoughts\":"   ws   string
RequestHeartbeatParam ::= "\"request_heartbeat\":"   ws   boolean
namestring ::= "\"human\"" | "\"persona\""
string ::= "\""   ([^"\[\]{}]*)   "\""
boolean ::= "true" | "false"
ws ::= [ \\t\\n]*
number ::= [0-9]+

[cpacker]

This is really awesome, thank you @Drake-AI ! I'll merge it into this PR as an additional grammar file.

Out of curiosity @Drake-AI while testing this what backend are you using? llama.cpp? web UI?

[Drake-AI]

This is really awesome, thank you @Drake-AI ! I'll merge it into this PR as an additional grammar file.

Out of curiosity @Drake-AI while testing this what backend are you using? llama.cpp? web UI?

koboldcpp which is a fork of llamacpp, should work with both.

[cpacker]

This is really awesome, thank you @Drake-AI ! I'll merge it into this PR as an additional grammar file.
Out of curiosity @Drake-AI while testing this what backend are you using? llama.cpp? web UI?

koboldcpp which is a fork of llamacpp, should work with both.

awesome, I guess I should probably also add "official" support for kobold with a catch for BACKEND_TYPE=kobold, though if it's a fork of llama.cpp I wouldn't be surprised if BACKEND_TYPE=llamacpp works for it

[Drake-AI]

No, it needs this variables:

"stop_sequence": [
    "\nUSER:",
    "\nASSISTANT:",
    "\nFUNCTION RETURN:",
    # '\n' +
    # '</s>',
    # '<|',
    # '\n#',
    # '\n\n\n',
],
"max_content_length": 4096,
"max_length": 512,

And you can send all parameters like temperature and all of that via api.

[Drake-AI]

I found a bug in my grammar, please change it before merge to this new version:

root ::= Function
Function ::= SendMessage | PauseHeartbeats | CoreMemoryAppend | CoreMemoryReplace | ConversationSearch | ConversationSearchDate | ArchivalMemoryInsert | ArchivalMemorySearch
SendMessage ::= "{"   ws   "\"function\":"   ws   "\"send_message\","   ws   "\"params\":"   ws   SendMessageParams   "}"
PauseHeartbeats ::= "{"   ws   "\"function\":"   ws   "\"pause_heartbeats\","   ws   "\"params\":"   ws   PauseHeartbeatsParams   "}"
CoreMemoryAppend ::= "{"   ws   "\"function\":"   ws   "\"core_memory_append\","   ws   "\"params\":"   ws   CoreMemoryAppendParams   "}"
CoreMemoryReplace ::= "{"   ws   "\"function\":"   ws   "\"core_memory_replace\","   ws   "\"params\":"   ws   CoreMemoryReplaceParams   "}"
ConversationSearch  ::= "{"   ws   "\"function\":"   ws   "\"conversation_search\","   ws   "\"params\":"   ws   ConversationSearchParams   "}"
ConversationSearchDate  ::= "{"   ws   "\"function\":"   ws   "\"conversation_search_date\","   ws   "\"params\":"   ws   ConversationSearchDateParams   "}"
ArchivalMemoryInsert  ::= "{"   ws   "\"function\":"   ws   "\"archival_memory_insert\","   ws   "\"params\":"   ws   ArchivalMemoryInsertParams   "}"
ArchivalMemorySearch  ::= "{"   ws   "\"function\":"   ws   "\"archival_memory_search\","   ws   "\"params\":"   ws   ArchivalMemorySearchParams   "}"
SendMessageParams ::= "{"   ws   InnerThoughtsParam   ","   ws   "\"message\":"   ws   string   ws   "}"
PauseHeartbeatsParams ::= "{"   ws   InnerThoughtsParam   ","      ws   "\"minutes\":"   ws   number   ws   "}"
CoreMemoryAppendParams ::= "{"   ws   InnerThoughtsParam   ","      ws   "\"name\":"   ws   namestring   ","   ws   "\"content\":"   ws   string   ws   ","   ws   RequestHeartbeatParam   ws   "}"
CoreMemoryReplaceParams ::= "{"   ws   InnerThoughtsParam   ","      ws   "\"name\":"   ws   namestring   ","   ws   "\"old_content\":"   ws   string   ","   ws   "\"new_content\":"   ws   string   ws   ","   ws   RequestHeartbeatParam   ws   "}"
ConversationSearchParams ::= "{"   ws   InnerThoughtsParam   ","      ws   "\"query\":"   ws   string   ws   ","   ws   "\"page\":"   ws   number   ws   ","   ws   RequestHeartbeatParam   ws   "}"
ConversationSearchDateParams ::= "{"   ws   InnerThoughtsParam   ","      ws   "\"start_date\":"   ws   string   ws   ","      ws   "\"end_date\":"   ws   string   ws   ","   ws   "\"page\":"   ws   number   ws   ","   ws   RequestHeartbeatParam   ws   "}"
ArchivalMemoryInsertParams ::= "{"   ws   InnerThoughtsParam    ","   ws   "\"content\":"   ws   string   ws   ","   ws   RequestHeartbeatParam   ws   "}"
ArchivalMemorySearchParams ::= "{"   ws   InnerThoughtsParam   ","  ws   "\"query\":"   ws   string   ws   ","   ws   "\"page\":"   ws   number   ws   ","   ws   RequestHeartbeatParam   ws   "}"
InnerThoughtsParam ::= "\"inner_thoughts\":"   ws   string
RequestHeartbeatParam ::= "\"request_heartbeat\":"   ws   boolean
namestring ::= "\"human\"" | "\"persona\""
string ::= "\""   ([^"\[\]{}]*)   "\""
boolean ::= "true" | "false"
ws ::= ""
number ::= [0-9]+

[cpacker]

I found a bug in my grammar, please change it before merge to this new version:

root ::= Function
Function ::= SendMessage | PauseHeartbeats | CoreMemoryAppend | CoreMemoryReplace | ConversationSearch | ConversationSearchDate | ArchivalMemoryInsert | ArchivalMemorySearch
SendMessage ::= "{"   ws   "\"function\":"   ws   "\"send_message\","   ws   "\"params\":"   ws   SendMessageParams   "}"
PauseHeartbeats ::= "{"   ws   "\"function\":"   ws   "\"pause_heartbeats\","   ws   "\"params\":"   ws   PauseHeartbeatsParams   "}"
CoreMemoryAppend ::= "{"   ws   "\"function\":"   ws   "\"core_memory_append\","   ws   "\"params\":"   ws   CoreMemoryAppendParams   "}"
CoreMemoryReplace ::= "{"   ws   "\"function\":"   ws   "\"core_memory_replace\","   ws   "\"params\":"   ws   CoreMemoryReplaceParams   "}"
ConversationSearch  ::= "{"   ws   "\"function\":"   ws   "\"conversation_search\","   ws   "\"params\":"   ws   ConversationSearchParams   "}"
ConversationSearchDate  ::= "{"   ws   "\"function\":"   ws   "\"conversation_search_date\","   ws   "\"params\":"   ws   ConversationSearchDateParams   "}"
ArchivalMemoryInsert  ::= "{"   ws   "\"function\":"   ws   "\"archival_memory_insert\","   ws   "\"params\":"   ws   ArchivalMemoryInsertParams   "}"
ArchivalMemorySearch  ::= "{"   ws   "\"function\":"   ws   "\"archival_memory_search\","   ws   "\"params\":"   ws   ArchivalMemorySearchParams   "}"
SendMessageParams ::= "{"   ws   InnerThoughtsParam   ","   ws   "\"message\":"   ws   string   ws   "}"
PauseHeartbeatsParams ::= "{"   ws   InnerThoughtsParam   ","      ws   "\"minutes\":"   ws   number   ws   "}"
CoreMemoryAppendParams ::= "{"   ws   InnerThoughtsParam   ","      ws   "\"name\":"   ws   namestring   ","   ws   "\"content\":"   ws   string   ws   ","   ws   RequestHeartbeatParam   ws   "}"
CoreMemoryReplaceParams ::= "{"   ws   InnerThoughtsParam   ","      ws   "\"name\":"   ws   namestring   ","   ws   "\"old_content\":"   ws   string   ","   ws   "\"new_content\":"   ws   string   ws   ","   ws   RequestHeartbeatParam   ws   "}"
ConversationSearchParams ::= "{"   ws   InnerThoughtsParam   ","      ws   "\"query\":"   ws   string   ws   ","   ws   "\"page\":"   ws   number   ws   ","   ws   RequestHeartbeatParam   ws   "}"
ConversationSearchDateParams ::= "{"   ws   InnerThoughtsParam   ","      ws   "\"start_date\":"   ws   string   ws   ","      ws   "\"end_date\":"   ws   string   ws   ","   ws   "\"page\":"   ws   number   ws   ","   ws   RequestHeartbeatParam   ws   "}"
ArchivalMemoryInsertParams ::= "{"   ws   InnerThoughtsParam    ","   ws   "\"content\":"   ws   string   ws   ","   ws   RequestHeartbeatParam   ws   "}"
ArchivalMemorySearchParams ::= "{"   ws   InnerThoughtsParam   ","  ws   "\"query\":"   ws   string   ws   ","   ws   "\"page\":"   ws   number   ws   ","   ws   RequestHeartbeatParam   ws   "}"
InnerThoughtsParam ::= "\"inner_thoughts\":"   ws   string
RequestHeartbeatParam ::= "\"request_heartbeat\":"   ws   boolean
namestring ::= "\"human\"" | "\"persona\""
string ::= "\""   ([^"\[\]{}]*)   "\""
boolean ::= "true" | "false"
ws ::= ""
number ::= [0-9]+

updated, thank you for the patch! also I added kobold support, though i haven't tested it yet

[cpacker]

Tested on kobold.cpp, seems to be working:

MemGPT % python main.py --model airoboros-l2-70b-2.1-grammar --first

[cpacker]

Pending extra review from @vivi will merge - @Drake-AI let me know if you catch any bugs in the grammar in the meantime.

[vivi]

LGTM 😈

[vivi]

Thanks for your help @Drake-AI! Will merge this in and add you as a co-author.