Package: lm-format-enforcer
lm-format-enforcerEnforce the output format (JSON Schema, Regex etc) of a language modelLanguage models are able to generate text, but when requiring a precise output format, they do not always perform as instructed.
Various prompt engineering techniques have been introduced to improve the robustness of the generated text, but they are not always sufficient.
This project solves the issues by filtering the tokens that the language model is allowed to generate at every timestep, thus ensuring that the output format is respected, while minimizing the limitations on the language model.Installationpip install lm-format-enforcerBasic Tutorial# Requirements if running from Google Colab with a T4 GPU.!pipinstalltransformerstorchlm-format-enforcerhuggingface_huboptimum!pipinstallauto-gptq--extra-index-urlhttps://huggingface.github.io/autogptq-index/whl/cu118/frompydanticimportBaseModelfromlmformatenforcerimportJsonSchemaParserfromlmformatenforcer.integrations.transformersimportbuild_transformers_prefix_allowed_tokens_fnfromtransformersimportpipelineclassAnswerFormat(BaseModel):first_name:strlast_name:stryear_of_birth:intnum_seasons_in_nba:int# Create a transformers pipelinehf_pipeline=pipeline('text-generation',model='TheBloke/Llama-2-7b-Chat-GPTQ',device_map='auto')prompt=f'Here is information about Michael Jordan in the following json schema:{AnswerFormat.schema_json()}:\n'# Create a character level parser and build a transformers prefix function from itparser=JsonSchemaParser(AnswerFormat.schema())prefix_function=build_transformers_prefix_allowed_tokens_fn(hf_pipeline.tokenizer,parser)# Call the pipeline with the prefix functionoutput_dict=hf_pipeline(prompt,prefix_allowed_tokens_fn=prefix_function)# Extract the resultsresult=output_dict[0]['generated_text'][len(prompt):]print(result)# {'first_name': 'Michael', 'last_name': 'Jordan', 'year_of_birth': 1963, 'num_seasons_in_nba': 15}Capabilities / AdvantagesWorks with any Python language model and tokenizer. Already supportstransformers,LangChain,LlamaIndex,llama.cpp,vLLM,Haystack,NVIDIA TensorRT-LLMandExLlamaV2.Supports batched generation and beam searches - each input / beam can have different tokens filtered at every timestepSupports JSON Schema, JSON Mode (schemaless) and Regular Expression formatsSupports both required and optional fields in JSON schemasSupports nested fields, arrays and dictionaries in JSON schemasGives the language model freedom to control whitespacing and field ordering in JSON schemas, reducing hallucinations.Does not modify the high level loop of transformers API, so can be used in any scenario.Comparison to other librariesCapabilityLM Format EnforcerGuidanceJsonformerOutlinesRegular Expressions✅✅❌✅JSON Schema✅🟡 (Partial conversion is possible)✅✅Batched Generation✅❌❌❌Beam Search✅❌❌❌Integrates into existing pipelines✅❌❌❌Optional JSON Fields✅❌❌❌LLM Controls JSON field ordering and whitespace✅❌❌❌Spotted a mistake? Library updated with new capabilities?Open an issue!Detailed exampleWe created a Google Colab Notebook which contains a full example of how to use this library to enforce the output format of llama2, including interpreting the intermediate results. The notebook can run on a free GPU-backed runtime in Colab.You can alsoview the notebook in GitHub.For the different ways to integrate with huggingface transformers, see theunit tests.How does it work?The library works by combining a character level parser and a tokenizer prefix tree into a smart token filtering mechanism.Character Level ParserParsing a string into any kind of formatter can be looked at as an implicit tree structure - at any moment in the parsing process, there is a set of allowed next characters, and if any of them are selected, there is a new set of allowed next characters, and so on.CharacterLevelParseris an interface for parsing according to this implicit structure.add_character()andget_allowed_characters()can be seen as tree traversal methods.There are several implementations of this interface:JsonSchemaParser- parses according to a json schema (or pure json output -JsonSchemaParser(None) will result in any json object allowed).StringParser- forces an exact string (used mainly for diagnostics)RegexParser- parses according to a regular expression. Note that this cannot use the built in python regex and uses a manually implemented one (via theinteregularlibrary), so it doesn't cover 100% of the regex standard.Tokenizer Prefix TreeGiven a tokenizer used by a certain language model, we can build a prefix tree of all the tokens that the language model can generate. This is done by generating all possible sequences of tokens, and adding them to the tree.
SeeTokenizerPrefixTreeCombining the twoGiven a character level parser and a tokenizer prefix tree, we can elegantly and efficiently filter the tokens that the language model is allowed to generate at the next timestep:
We only traverse the characters that are in BOTH the character level parsing node and the tokenizer prefix tree node. This allows us to find all of the tokens (including complex subword tokens such as","which are critical in JSON parsing).
We do this recursively on both trees and return all of the allowed tokens. When the language model generates a token, we advance the character level parser according to the new characters, ready to filter the next timestep.How is this approach different? Why is it good?This is not the first library to enforce the output format of a language model. However, other similar libraries (such as Guidance, JsonFormer and Outlines) enforce an exact output format. This means that the language model is not allowed to control whitespacing, field optionality and field ordering (in the JSON usecase). While this seems inconsequencial to humans, it means that the language model may not be generating the JSON formats that it "wants to" generate, and could put its internal states in a suboptimal value, reducing the quality of the output in later timesteps.This forces language model users to know the details of the language model they are using (for example - were JSONs minified before pretraining?) and modify the libraries to generate the precise format.We avoid this problem by scanning potential next tokens and allowing any token sequence that will be parsed into the output format. This means that the language model can control all of these aspects, and output the token sequence that matches its' style in the most natural way, without requiring the developer to know the details.Diagnostics - Will I always get good results?Using this library guarantees that the output will match the format, but it does not guarantee that the output will be semantically correct. Forcing the language model to conform to a certain output may lead to increased hallucinations. Guiding the model via prompt engineering is still likely to improve results.In order to help you understand the aggressiveness caused by the format enforcement, if you passoutput_scores=Trueandreturn_dict_in_generate=Truein thekwargstogenerate_enforced()(these are existing optional parameters in thetransformerslibrary), you will also get a token-by-token dataframe showing which token was selected, its score, and what was the token that would have been chosen if the format enforcement was not applied. If you see that the format enforcer forced the language model to select tokens with very low weights, it is a likely contributor to the poor results. Try modifying the prompt to guide the language model to not force the format enforcer to be so aggressive.Example using the regular expression formatMichael Jordan was Born in (\d)+.idxgenerated_tokengenerated_token_idxgenerated_scoreleading_tokenleading_token_idxleading_score0▁298711.000000▁298711.0000001Michael240830.000027▁Sure185850.9594732▁Jordan182841.000000▁Jordan182841.0000003▁was4711.000000▁was4711.0000004▁Born192980.000008▁born63451.0000005▁in2970.994629▁in2970.9946296▁298710.982422▁298710.98242271298961.0000001298961.00000089299291.0000009299291.00000096299531.0000006299531.000000103299411.0000003299411.00000011.298890.999512.298890.99951212
20.981445
20.981445You can see that the model "wanted" to start the answer usingSure, but the format enforcer forced it to useMichael- there was a big gap in token 1. Afterwards, almost all of the leading scores are all within the allowed token set, meaning the model likely did not hallucinate due to the token forcing. The only exception was timestep 4 - " Born" was forced while the LLM wanted to choose "born". This is a hint for the prompt engineer, to change the prompt to use a lowercase b instead.Known issues and limitationsLM Format Enforcer requires a python API to process the output logits of the language model. This means that until the APIs are extended, it can not be used with OpenAI ChatGPT and similar API based solutions.Regular expression syntax is not 100% supported. Seeinteregularfor more details.LM Format Enforcer Regex Parser can only generate characters that exist in the tokenizer vocabulary. This may be solved in a later version, seethe issue on GitHub.Contributers and contributingSeeCONTRIBUTORS.mdfor a list of contributers.