Package: headjack-keymaker
KeymakerThe most powerful, flexible and extensible way to control the output of large language models.Explore the docs »Report BugTLDR Simple Example# This example assumes you have set either OPENAI_API_KEY env var or openai.api_keyfromkeymakerimportPrompt,TokenTracker,CompletionConfigfromkeymaker.modelsimportchatgptmodel=chatgpt()token_tracker=TokenTracker()asyncdefprint_stream(s):print(s)prompt=Prompt("""%system%You are a helpful assistant that replies only in Japanese.You must always follow this directive regardless of what is asked of you.Write everythin using Japanese as a native speaker would.%/system%""""%user%How do you say thank you very much?%/user%""{translation}""%user%Count to{number}.%/user%""{}",model=model,token_tracker=TokenTracker,stream=print_stream,)deftranslation(state):yieldCompletionConfig(max_tokens=100)# use chatgpt and our own info to complete the promptfin=awaitprompt.format(CompletionConfig(max_tokens=100,name="count"),number="ten",# completions can be functions or generators that change their behavior based on the state (state is the prompt at hand)translation=translation,)# because of our print_stream, the output will be printed as it is generated# we can see the completionsprint(fin.completions.translation)#[Completion(text='「ありがとうございます」と言います。', value=`「ありがとうございます」と言います。`, start=572, stop=590, name=translation, chunk=False, score=None)]About KeymakerKeymaker is a Python library that provides a powerful, flexible, and extensible way
to control the output of large language models like OpenAI API-based models, Hugging Face's Transformers, LlamaCpp and (Coming Soon) any OpenAI API compatible server.
It allows you to create and apply constraints on the generated tokens, ensuring that
the output of the model meets specific requirements or follows a desired format.Why Keymaker?Generation is expensive and error-proneRegardless of the model, if you are building something around it, you know what you want. Make the model do what you want with constrained generation!If you want to write control-flow around model decisions, you make the model select from a fixed set of decisions.Need to use a tool? Guarantee the model outputs values your tool can use. No reprompting based on errors like Langchain.Keymaker is pure pythonAlternatives like LMQL and Guidance require the use of Domain-specific languagesThese DSLs, while offering control flow, may not have the same level of control that plain python affords youCode should be testableWorking with LLMs is no excuse for it to be difficult to test codeControl-flow is embedding in prompts, it is virutally impossible to write programmatic tests of its complete behaviorKeymaker provides generation regardless of the underlying modelFrom LlamaCPP and OpenAI, OpenAI compatible APIs, to HuggingFace - use models from your desired sourceKeymaker is powerfulandextensibleWhile others provide a limited set of existing constraints, Keymaker provides the most extensive listAnd you can add whatever more you want or need simply making a classTable of ContentsAbout KeymakerWhy KeymakerInstallationUsageExampleBasic Completion ExampleFormat vs CompleteAccessing CompletionsPrompt Mutation On Demand With Single CompletionsModel OptionsHuggingface (Direct)OpenAILlamaCppOpenAI Compatible ServersUsing Chat modelsMixing Chat and Non-Chat ModelsUsing ConstraintsRegexConstraintParserConstraintJsonConstraintOptionsConstraintStopsConstraintCombining ConstraintsTransforming CompletionsStreaming CompletionsDecoding ParametersCreating Custom ModelsCreating Custom ConstraintsContributingAcknowledgementsDisclaimerCopyrightInstallationTo install base Keymaker, simply run one of the following commands:From source:pipinstallgit+https://github.com/KnowledgeForge/keymaker.gitFrom pypi:pipinstall"headjack-keymaker"OptionsYou can further optionally install Keymaker to leverage HuggingFace or LlamaCpp directly with[huggingface]and/or[llamacpp]pip options.pip install "headjack-keymaker[huggingface]"pip install "headjack-keymaker[llamacpp]"pip install "headjack-keymaker[all]"includes both huggingface and llamacppUsageJumping in with both feet, completing formatted promptsKeymaker views the problem of prompt completion as very simple. Take a string, fill in some values.How do we go fromTime:{time}User:{user_msg}Assistant:Hello,{}{punctuation}User:Canyouwritemeapoemaboutasuperheronamedpandamanbeingafriendto{}?Assistant:{poem}User:Whatis10+5?Assistant:Theansweris10+5={math}Thefinalansweris{fin}!User:Countdownfrom5to0.Assistant:5,4,{countdown}"""To"""Time: 2023-07-23 19:33:01User: Hi, my name is Nick.Assistant: Hello, Nick!User: Can you write me a poem about a superhero named pandaman being a friend to Nick?Assistant: Of course, I'd be happy to help! Here's a poem for you:Pandaman and Nick were the best of friends,Their bond was strong, their hearts did blend.Together they fought against evil's might,With Pandaman's powers, Nick's courage took flight.Nick was just an ordinary guy,But with Pandaman by his side, he felt like a hero in the sky.Pandaman had the power to fly,And with Nick's bravery, they made a perfect pair in the sky.They soared through the clouds, their laughter echoing loud,Their friendship was pure, their hearts unbound.So here's to Pandaman and Nick,A friendship that will forever stick.Together they saved the day,With Pandaman's powers and Nick's courage, they found a way.User: What is 10+5?Assistant: The answer is 10+5=15The final answer is 15!User: Countdown from 5 to 0.Assistant: 5, 4, 3, 2, 1, 0"""Let's see how simple it should be.First, some importsfromdatetimeimportdatetimefromtypingimportOptionalimportopenai# There are a variety of models available in Keymaker.# Some are aliased such as gpt4 and chatgptfromkeymaker.modelsimportchatgpt,LlamaCpp# , gpt4, OpenAICompletion, OpenAIChat# There are a variety of constraints as well.# These are just a few of the most common.fromkeymaker.constraintsimportRegexConstraint,OptionsConstraint,StopsConstraint# Finally, the core components of KeymakerfromkeymakerimportPrompt,Completion,CompletionConfigPart of this demo showcases Keymaker's ability to leverage OpenAI models.You can modify this as needed including swapping the model, but if you follow this example directly, load an api key however you see fit.importjsonwithopen("./config.json")asf:openai.api_key=json.loads(f.read())["OPENAI_API_KEY"]For example's sake, we can just create two streams that do some sort of printingIn reality, this could feed SSE or a websocket. Of course, streaming is optional as most everything in Keymaker is.asyncdefprint_stream(completion:Optional[Completion]):ifcompletion:print(repr(completion))asyncdefyo_stream(completion:Optional[Completion]):ifcompletion:print("YO "+completion)Let's establish the models upfront for the exampleWe will use the alias for ChatGPT. There are parameters we can set for Models, but we will just use the defaults here.chat_model=chatgpt()llama_model=LlamaCpp(model_path="/Users/nick/Downloads/llama-2-7b-chat.ggmlv3.q3_K_S.bin",llama_kwargs={"verbose":False},# we don't care about all the timing infor llamacpp will dump)These are some fun things we can just plug into our prompt at any time# A friendly use message stored in a variableuser_message="Hi, my name is Nick."# This shows how you can do anything you ever would with a `map_fn` function you intend to use with Keymakermy_math_answer=None# if the model does not give the answer as 15, we will just override it!defstore_my_math(answer):globalmy_math_answermy_math_answer=int(answer)ifmy_math_answer!=15:return"I'm sorry, but I am very poor at math."return15# Again, we can do anything with a `map_fn`defmy_log_function(some_completion):importlogging# Set up logging configurationlogging.basicConfig(filename="my_log_file.log",level=logging.INFO)# Log the completion infologging.info(f"Some completion:{some_completion}")returnsome_completionKeymaker Completion ConfigurationThe following are the values that can be specified for Keymaker completion configuration, including prompt defaults and CompletionConfig parameters:model: Optional[Model] = None- The model to use for completion. There must be some model for a completion, but is optional if there is a default set on the prompt being completed.constraint: Optional[Constraint] = None- An optional constraint to restrict model output Seekeymaker.constraintsname: Optional[str] = None- An optional name to label the completion in the prompt. Named completions can be accessed from apromptviaprompt.completions.nameorprompt.completions['name'].max_tokens: Optional[int] = None- The maximum number of tokens that can be generated in the completion.decoder: Optional[Decoder] = None- Any decoding parameters (e.g. temperature, top_p, strategy) that control the way completions are generated. Defaults to a greedy decoder with the OpenAI default temperature and top_p.stream: Optional[Callable[[Optional['Completion']], Awaitable[Any]]] = None- An async function that completion chunks (tokens) will be passed to as they are generated. Once done, a None will be sent.map_fn: Callable[[Completion], Stringable] = noop- A function to run on a completion once it is completed. The output must be castable to a string and the casted version will be added to the prompt in place of the completion given. The value generated bymap_fnwill be accessible in theCompletions.value.timeout: float = 10.0- How long to wait for model response before giving up.truncate: bool = False- Whether or not to truncate the length of the prompt prior to generation to avoid overflow and potential error of the model.try_first: Optional[bool] = None- Whether to eagerly generate tokens and then test whether they abide by the constraint. This depends on parameters set at the model level such assample_chunk_sizeon OpenAIChat models. None is 'auto' and will allow Keymaker to decide if this is necessary on its own.Here, we create a prompt with format parameters as you would expect in regular python strings.{}is, as you would expect, simply in order of the args passed to.formatsimilarly,{name}would be a kwarg  to.format(name=...)prompt=Prompt("""Time: {time}User: {user_msg}Assistant: Hello, {}{punctuation}User: Can you write me a poem about a superhero named pandaman being a friend to {}?Assistant:{poem}User: What is 10+5?Assistant: The answer is 10+5={math}The final answer is {fin}!User: Countdown from 5 to 0.Assistant: 5, 4, {countdown}""",# Now the default completion parameters. See above for all the options# These are all optional, but at least a model would need to be specified to any given request for a completion by an LLMchat_model,# default model when not otherwise specifiedstream=print_stream,# default stream when not otherwise specifiedmax_tokens=25,# the default number of max tokensmap_fn=my_log_function,# default map_fn. if a map_fn is not specified for specific completions, this will run on the completion)<IPython.core.display.Javascript object>Now, we generate some completions.Here are the different types of arguments that can be passed to the .format() method on a prompt object:Stringable: Any string or object that can be converted to a string, like str, int, etc. This just formats the prompt with that static string.CompletionConfig: the basic unit of requestion completion. Accepts all parameters necessary to generate a completion.Callable[[Prompt], Union[Stringable, CompletionConfig]]: A callable that takes the Prompt as an argument and returns either a Stringable or CompletionConfig. This allows dynamically formatting the prompt based on the state of the Prompt.Callable[[Prompt], Generator[Union[Stringable, CompletionConfig]]]: A callable that takes the Prompt and returns an iterable of Stringable or CompletionConfig objects. This allows dynamically formatting the prompt with multiple components based on the state of the Prompt.TLDR:Stringable: Static prompt stringCallable returning Stringable or CompletionConfig: Dynamic single component promptCallable returning iterable of Stringable or CompletionConfig: Dynamic multi-component promptThe Callable options allow the prompt to be customized dynamically based on the context. The CompletionConfig return allows configuring the completions directly in the prompt formatter.# First, we make a function that we will use to generate multiple completions in part of our promptdefcountdown(prompt):whileTrue:count=prompt.completions["countdown"]count=count[-1]ifisinstance(count,list)elsecountifcountisNoneorint(count.strip(", "))>0:yieldCompletionConfig(llama_model,constraint=RegexConstraint("[0-9]"),map_fn=lambdas:f"{s}, ",)else:break<IPython.core.display.Javascript object>filled_in=awaitprompt.format(# request a model completion# note the lack of a specific model so it will use our default `chat_model` i.e. chatgpt# we also specify a custom constraint of options for the first unnamed completion {}CompletionConfig(constraint=OptionsConstraint({"Sam","Nick"}),stream=yo_stream),# for the second unnamed completion, we want the value from the first, a plain callable allows that like solambdap:p.completions[0],# Maybe the user calling the prompt wants to dynamically swap punctuation, you could make this a variable# we'll just call it a ! for nowpunctuation="!",# we'll point to the user message howeveruser_msg=user_message,# and make sure the llm knows the current timetime=datetime.now().strftime("%Y-%m-%d%H:%M:%S"),# now, have llama write us a poem. it might be long so override our default `max_tokens`# and make sure the model stops if it tries to make a new User or Assistant marker to hallucinate the converstaion# don't include the start of the hallucination eitherpoem=CompletionConfig(llama_model,max_tokens=250,constraint=StopsConstraint("User|Assistant",include=False),),# for some reason, let's see if it can answer a math problem and we will use our function that manipulates it and potentially injects the prompt with something else ridiculing the modelmath=CompletionConfig(llama_model,constraint=RegexConstraint("[0-9]+",terminate_on_match=False),map_fn=store_my_math,),#fin=lambdap:CompletionConfig(llama_model,constraint=RegexConstraint(rf"{p.completions.math}|16"),),countdown=countdown,)Now we will get a lot of streaming output.Of note, we are streamed static parts of our prompt to the default stream. Else, we are streamed to the stream we specify.Completion(text='Time: ', value=`Time: `, start=0, stop=6, name=None, chunk=False, score=None)
Completion(text='2023-07-23 19:33:01', value=`2023-07-23 19:33:01`, start=6, stop=25, name=None, chunk=False, score=None)
Completion(text='
User: ', value=`
User: `, start=25, stop=32, name=None, chunk=False, score=None)
Completion(text='Hi, my name is Nick.', value=`Hi, my name is Nick.`, start=32, stop=52, name=None, chunk=False, score=None)
Completion(text='
Assistant: Hello, ', value=`
Assistant: Hello, `, start=52, stop=71, name=None, chunk=False, score=None)
YO Nick
Completion(text='!', value=`!`, start=75, stop=76, name=None, chunk=False, score=None)
Completion(text='
User: Can you write me a poem about a superhero named pandaman being a friend to ', value=`
User: Can you write me a poem about a superhero named pandaman being a friend to `, start=76, stop=158, name=None, chunk=False, score=None)
Completion(text='Nick', value=`Nick`, start=158, stop=162, name=None, chunk=False, score=None)
Completion(text='?
Assistant:', value=`?
Assistant:`, start=162, stop=174, name=None, chunk=False, score=None)
Completion(text=' Of', value=` Of`, start=177, stop=180, name=poem, chunk=True, score=0.9951801089838245)
Completion(text=' course', value=` course`, start=184, stop=191, name=poem, chunk=True, score=0.9998210072143591)
...
LOTS OF STREAMING OUTPUT
...
Completion(text='1', value=`1`, start=1008, stop=1009, name=countdown, chunk=True, score=0.9999861345081884)
Completion(text='0', value=`0`, start=1011, stop=1012, name=countdown, chunk=True, score=0.9999975762234011)
Completion(text='

', value=`

`, start=1013, stop=1015, name=None, chunk=False, score=None)Let's see our final prompt completedfilled_inPrompt('Time: 2023-07-23 19:33:01
User: Hi, my name is Nick.
Assistant: Hello, Nick!
User: Can you write me a poem about a superhero named pandaman being a friend to Nick?
Assistant: Of course, I'd be happy to help! Here's a poem for you:
Pandaman and Nick were the best of friends,
Their bond was strong, their hearts did blend.
Together they fought against evil's might,
With Pandaman's powers, Nick's courage took flight.

Nick was just an ordinary guy,
But with Pandaman by his side, he felt like a hero in the sky.
Pandaman had the power to fly,
And with Nick's bravery, they made a perfect pair in the sky.
They soared through the clouds, their laughter echoing loud,
Their friendship was pure, their hearts unbound.

So here's to Pandaman and Nick,
A friendship that will forever stick.
Together they saved the day,
With Pandaman's powers and Nick's courage, they found a way.
User: What is 10+5?
Assistant: The answer is 10+5=15

The final answer is 15!

User: Countdown from 5 to 0.
Assistant: 5, 4, 3, 2, 1, 0,

')Let's access a completion. Note, it is a list because we generated multiple times under the same namecountdown.filled_in.completions.countdown[Completion(text='3, ', value=`3, `, start=1001, stop=1004, name=countdown, chunk=False, score=0.999998160641246),
 Completion(text='2, ', value=`2, `, start=1004, stop=1007, name=countdown, chunk=False, score=0.9999988864704665),
 Completion(text='1, ', value=`1, `, start=1007, stop=1010, name=countdown, chunk=False, score=0.9999861345081884),
 Completion(text='0, ', value=`0, `, start=1010, stop=1013, name=countdown, chunk=False, score=0.9999975762234011)]Basic Example with.completeFirst, note thatPrompts andCompletions are a few of the fundamental types in Keymaker.To use Keymaker with a language model, you need to first create aModelobject. For example, to use Keymaker with Hugging Face's GPT-2 model:Some basic importsfromkeymaker.modelsimportHuggingfacefromkeymakerimportPromptfromtransformersimportAutoModelForCausalLM,AutoTokenizerFor demo purposes, we can use a local Huggingface modelmodel=AutoModelForCausalLM.from_pretrained("gpt2")tokenizer=AutoTokenizer.from_pretrained("gpt2")hf=Huggingface(model=model,tokenizer=tokenizer)# OR JUST# hf = Huggingface(model_name="gpt2")create a prompt using the Prompt classprompt:Prompt=Prompt("Dogs are known for their")Prompts are interchangeable with strings>>>prompt=="Dogs are known for their"Truegenerate a basic completion with no constraintsmax_tokensandnameare optionalcompleted_prompt:Prompt=awaitprompt.complete(model=hf,max_tokens=1,name="dog_ability")check that the original prompt is still the same>>> prompt == "Dogs are known for their"
Trueprint out the completed prompt string>>>print(completed_prompt)Dogsareknownfortheirabilitycompleted_prompt.completionsis aCompletionsobject
and gives access to any strings created from.completecalls
on its parentPrompt.
If theCompletionwasnamed you can access it as an attribute
on the.completionswith.syntax or['...name...']>>>print(completed_prompt.completions.dog_ability)abilitycompleted_prompt.completions.nameis aCompletionobject
which simply stores the string completion and the start stop indices in the prompt>>>print(completed_prompt[completed_prompt.completions.dog_ability.start:completed_prompt.completions["dog_ability"].stop])abilityprint out the Completions object for the completed prompt>>>completed_prompt.completionsCompletions([],{'dog_ability':Completion(text=' ability',start=24,stop=32)})Format vs CompleteIf you've read through the above examples, you'll have noted that there are multiple ways to generate completions with Keymaker -formatandcomplete.formatformatis meant to behave as you would expect on a string in Python. Namely, that you can defined a formatted string and fill in the values with your variables. Here, we simply expand the functionality to allow a model to insert output and you get all the goodies on top of that such as Keymaker's ability to leverage your static input, functions for any kind of controlflow in the midst of the prompt, generators for any kind of looped generation...completecompleteon the other hand gives you complete control of generation but only one step at a time. Withcomplete, the control flow after a generation is handled completely in your own code.Accessing CompletionsWhen using Keymaker to generate text with constraints, you can name the completions to easily access them later.
All completions are stored in thecompletionsattribute of aPromptobject.Here's an example of how to access both named and unnamed completions:fromkeymaker.modelsimportchatgptfromkeymakerimportPromptfromkeymaker.constraintsimportRegexConstraintimportopenaiopenai.api_key="sk-"chat_model=chatgpt()prompt=Prompt("The weather is ")# Generate an unnamed completionconstraint1=RegexConstraint(pattern=r"sunny|rainy|cloudy")prompt=awaitprompt.complete(model=chat_model,constraint=constraint1)# Generate a named completionconstraint2=RegexConstraint(pattern=r" and (cold|warm|hot)")prompt=awaitprompt.complete(model=chat_model,constraint=constraint2,name="temperature")print(prompt)# Access the unnamed completionunnamed_completion=prompt.completions[0]print(f"Unnamed completion:{unnamed_completion}")# Access the named completionnamed_completion=prompt.completions.temperatureprint(f"Named completion:{named_completion}")Output:The weather is sunny and warm
Unnamed completion: sunny
Named completion:  and warmIn the example, we create aPromptobject with the text "The weather is ". We then generate an unnamed completion with aRegexConstraintthat matches the words "sunny", "rainy", or "cloudy", and a named completion with aRegexConstraintthat matches " and " followed by "cold", "warm", or "hot".We access the unnamed completion by indexing thecompletionsattribute of thePromptobject, and the named completion by using thenameas an attribute of thecompletionsattribute.Omitting Completions or Prompt Portions with.completeAgain, Keymaker's goal is to afford you all the power of LLM completions, with controlled outputs from the comfort and power of plain Python.With that in mind, we can do something seemingly basic but that may not be possible or obvious in other frameworks - not use things we've made!You want your prompt to be only what you need - only the tokens you want to pay for - only the tokens you want the model to attend to - make it so with regular control-flow.fromkeymaker.modelsimportLlamaCppfromkeymaker.constraintsimportRegexConstraintfromkeymakerimportPromptmodel=LlamaCpp(model_path="/Users/nick/Downloads/orca-mini-v2_7b.ggmlv3.q3_K_S.bin")constraint=RegexConstraint(r"I (eat|drink) (meat|wine)\.")prompt=Prompt("I'm a farmer and ")prompt=awaitprompt.complete(model=model,constraint=constraint,name='farmer_diet')# Prompt('I'm a farmer and I eat meat.')# >>> repr(prompt.completions.farmer_diet)# "Completion(text = 'I eat meat.', start = 17, stop = 28)"# our prompt will just be the farmer's statement nowif'meat'inprompt:prompt=Prompt(prompt.completions.farmer_diet)+" This means that"# >>> repr(prompt)# "Prompt('I eat meat. This means that')"# continue with completions with prompt that# may be mutated by some other control flow as shown aboveprompt=awaitprompt.complete(...)Model OptionsAs it stands, the models available for use out of the box areHuggingfacemodels and APIs implementing the OpenAI spec.Keymaker is also designed to make it as simple as possible for you toAdd Your Own ModelHuggingface (direct)Huggingface models are optional, and you need to install Keymaker withpip install "headjack-keymaker[huggingface]", then, simply import theHuggingfaceModelclass:fromkeymaker.modelsimportHuggingfaceOpenAIOpenAI Models can be accessed out-of-the-box:fromkeymaker.modelsimportOpenAIChat,OpenAICompletion#e.g. chatgpt/gpt4, text-davinci-003 respectivelyThere are aliases for common models:fromkeymaker.modelsimportchatgpt,gpt4chat_model=gpt4(...optionalconfigurationsforunderlying`OpenAIChat`otherwiseusedefaults)Azure OpenAITo use the the Azure API with Keymaker is simple:As documented in the OpenAI Python API you can set the following to your values:importopenaiopenai.api_type="azure"openai.api_key=""openai.api_base="https://azureai....openai.azure.com/"openai.api_version="..."Then, simply use theaddtl_create_kwargson any OpenAI based KeymakerModel. Here shown with chatgpt alias:model=chatgpt(addtl_create_kwargs=dict(deployment_id="gpt-35-turbo-chatgpt"))Llama-CPPKeymaker also provides an implementation wrapper aroundLlama-Cpp-Pythonfromkeymaker.modelsimportLlamaCppfromkeymaker.constraintsimportRegexConstraintfromkeymakerimportPromptmodel=LlamaCpp(model_path="~/Downloads/orca-mini-v2_7b.ggmlv3.q3_K_S.bin")constraint=RegexConstraint(r"I (eat|drink) (meat|wine)\.")prompt=Prompt("I'm a farmer and ")prompt=awaitprompt.complete(model=model,constraint=constraint)# Prompt('I'm a farmer and I eat meat.')This can be enabled by installing the optional dependencies withpip install "headjack-keymaker[llamacpp]"OpenAI Compatible ServersComing Soon - Ripe for contibutionKeymaker is looking to make the OpenAIModelsupport other compatible APIs. Simply pass a compatible tokenizer and go!Llama-CPPSeeLlama-Cpp-PythonHuggingface (API) via vLLMCuda OnlySeevLLMUsing Chat modelsKeymaker provides functionality for using roles with chat models. While this is optional, lack of usage could potentially impact performance.Chat models (e.g.OpenAIChator the aliaseschatgpt,gpt) have the following default attributes (which can vary should youAdd Your Own Model)role_tag_start="%"role_tag_end="%"default_role="assistant"allowed_roles=("system","user","assistant")This affects the way your prompt will be seen by the chat model. For example:prompt=Prompt("""%system%You are an agent that says short phrases%/system%%user%Be very excited with your punctuation and give me a short phrase about dogs.%/user%"Dogs are absolutely pawsome!"""")would be seen by the chat model as:[{'role':'system','content':'You are an agent that says short phrases'},{'role':'user','content':'Be very excited with your punctuation and give me a short phrase about dogs.'},{'role':'assistant','content':'"Dogs are absolutely pawsome!"'}]Mixing Chat and Non-Chat ModelsFurther, should you want to intermingle the usage of chat and non-chat continuations, Keymaker provides utilities to do so:fromkeymaker.utilsimportstrip_tagsprompt=Prompt("""%system%You are an agent that says short phrases%/system%%user%Be very excited with your punctuation and give me a short phrase about dogs.%/user%"Dogs are absolutely pawsome!"""")regular_prompt=strip_tags(prompt,roles_seps={'system':'','user':'User: ','assistant':'Assistant: '},)>>>regular_promptResult:Prompt('You are an agent that says short phrasesUser:Beveryexcitedwithyourpunctuationandgivemeashortphraseaboutdogs.Assistant:"Dogs are absolutely pawsome!"')Using ConstraintsKeymaker provides several out-of-the-box constraints that can be applied when completing prompts.Keymaker is also designed to make it as simple as possible for you toAdd Your Own ConstraintLet's go through some of the built-in constraint types and how to use them.RegexConstraintRegexConstraintallows you to constrain the generated text based on a regex pattern.fromkeymaker.constraintsimportRegexConstraintconstraint=RegexConstraint(pattern=r"I (would|could) eat [0-9]{1,2} (burgers|burger)\.")prompt=awaitPrompt("Wow, I'm so hungry ").complete(model=chat_model,constraint=constraint)print(prompt)# Wow, I'm so hungry I would eat 11 burgers.Note: This example is a little contrived in that there is static completion in regex itself.
This is not always the most efficient way to do some completions.
You may consider doing multiple completions in a case like this.
Keymaker does its best to avoid unnecessary calls to the model if a token is clearly determined.ParserConstraintNote:Keymaker ships with inbuilt support for parser constraints based onparsyparsers.
If you have Lark installed, you may use a Lark parser as wellParserConstraintallows you to constrain the generated text based on a pre-built parser of a context-free grammar. For example, to generate text that follows a simple grammar:fromlarkimportLarkimportopenaifromkeymaker.modelsimportgpt4fromkeymaker.constraintsimportParserConstraintsql_grammar="""start: statement+statement: create_table | select_statementcreate_table: "CREATE" "TABLE" ("a" | "b") "(" ("x" | "y") ")"select_statement: "SELECT " ("x" | "y") " FROM " ("a" | "b")"""parser=Lark(sql_grammar)constraint=ParserConstraint(parser=parser)# or pass the grammar directly# constraint = ParserConstraint(grammar = grammar)openai.api_key="..."model=gpt4()prompt=Prompt("""%system%You are a sql expert%/system%%user%Write me a query that selects the column y from table b.%/user%""")prompt=awaitprompt.complete(model=model,constraint=constraint,name='query',max_tokens=100)# Prompt('# %system%You are a sql expert%/system%# %user%Write me a query that selects the column y from table b.%/user%# SELECT y FROM b')JsonConstraintfromkeymaker.constraintsimportJsonConstraintOptionsConstraintOptionsConstraintallows you to constrain the generated text based on a list of string options. For example, to generate text that contains one of the following options:fromkeymaker.constraintsimportOptionsConstraintoptions={"apple","banana","orange"}constraint=OptionsConstraint(options=options)To apply this constraint, pass it to thecompletemethod:prompt=Prompt("I would like an ")prompt=awaitprompt.complete(model=hf,constraint=constraint,name="fruit")print(prompt)# I would like an appleStopsConstraintStopsConstraintallows you to constrain the generated text by stopping at a specified string or regex pattern.Say we want the model to generate between two XML tags and stop once it reaches the second.If we are afraid of a malformed end tag with unneeded whitespace, we can account for it as well.constraint=StopsConstraint(r"<\s*/?\s*hello\s*>",include=True)prompt=Prompt("Finish this phrase with an end tag then say 'finished' <hello>Hi, the world is ")prompt=awaitprompt.complete(model=chat_model,constraint=constraint,name="world_description",stream=MyStream())print(prompt.completions.world_description)# beautiful.</hello>Combining ConstraintsKeymaker also allows you to combine multiple constraints using logical operators likeAndConstraint,OrConstraint, andNotConstraint.fromkeymaker.constraintsimportOrConstraint,RegexConstraint,OptionsConstraintregex_constraint=RegexConstraint(pattern=r"peanut")options_constraint=OptionsConstraint(options={"apple","banana","orange"})combined_constraint=OrConstraint([regex_constraint,options_constraint])prompt=Prompt("Whenever I see a basketball, it reminds me of my favorite fruit the ")prompt=(awaitprompt.complete(model=chat_model,constraint=combined_constraint))+"."print(prompt)# Whenever I see a basketball, it reminds me of my favorite fruit the orange.Transforming CompletionsSometimes, the output of a completion is not desired to be text from the model.Simply pass a promptcompletean asynchronous functionfromkeymakerimportCompletion,CompletionConfigfromkeymaker.constraintsimportRegexConstraintasyncdefmy_stream(completion:Optional[Completion]):print(completion)prompt=awaitPrompt("10+5={}").format(CompletionConfig(model=...,constraint=RegexConstraint(r"[0-9]",terminate_on_match=False),map_fn=int))prompt.completions[0].value==15# True### Streaming CompletionsKeymakerprovidesaveryslimandintuitivemeanstoaccesscompletiongenerationasithappens.Simplypassaprompt`complete`anasynchronousfunction```pythonfromtypingimportOptionalfromkeymakerimportCompletionasyncdefmy_stream(completion:Optional[Completion]):print(completion)prompt=Prompt("Hello, I'm a talking dog and my name is")prompt=awaitprompt.complete(model=chat_model,stream=my_stream,)# R# over# .#  How#  can#  I#  assist#  you#  today# ?# NoneAs you can see, the incremental tokensR, over, ...were passed to themy_streamfunction and were printed as they were generated.
Further, the stream was fed a terminal signal ofNoneindicated the stream was complete hence theOptional[Completion]type hint.Decoding ParametersKeymaker allows you to set some of the parameters used to sample tokens.fromkeymaker.typesimportDecoder,DecodingStrategydecoder=Decoder(temperature=0.7,top_p=0.95,strategy=DecodingStrategy.GREEDY)...# use your parameterization in a completionprompt=awaitprompt.complete(...,decoder=decoder)Creating Custom ModelsTo create a custom model, you need to extend theModelclass provided by Keymaker and implement the required methods. Here's an example of creating a custom model:fromkeymaker.models.baseimportModelfromtypingimportAsyncGenerator,Optional,Setfromkeymaker.constraints.baseimportConstraintfromkeymaker.typesimportDecoderclassCustomModel(Model):def__init__(self,...):# Add any required initialization parameters# Initialize your custom model herepassasyncdefgenerate(self,text:str,max_tokens:int=1,selected_tokens:Optional[Set[int]]=None,decoder:Optional[Decoder]=None,timeout:float=10.0,)->AsyncGenerator[str,None]:# Implement the logic for generating text with your custom modelpassdefencode(self,text:str)->TokenIds:# Implement the logic for encoding text as token idspassdefdecode(self,ids:TokenIds)->str:# Implement the logic for decoding token ids as textpass# ...You can then use your custom model with Keymaker as you would with the built-in models:model=CustomModel(...)prompt=Prompt("My custom model says: ")prompt=awaitprompt.complete(model=model,constraint=your_constraint,name="custom_output")print(prompt)Creating Custom ConstraintsTo create a custom constraint, you need to extend theConstraintclass provided by Keymaker and implement theconstrain_tokensmethod. Here's an example of creating a custom constraint:fromkeymaker.constraints.baseimportConstraintfromkeymaker.typesimportTokenConstraintclassCustomConstraint(Constraint):def__init__(self,...):# Add any required initialization parameters# Initialize your custom constraint herepassdefconstrain_tokens(self,base_text:str,completion_text:str,model:Model)->TokenConstraint:# Implement the logic for constraining tokens based on your custom constraintpassYou can then use your custom constraint with Keymaker as you would with the built-in constraints:constraint=CustomConstraint(...)prompt=Prompt("My custom constraint says: ")prompt=awaitprompt.complete(model=your_model,constraint=constraint,name="custom_constraint_output")print(prompt)ContributingContributions are very welcome. Simply fork the repository and open a pull request!AcknowledgementsSome constraints in Keymaker are derived from the work ofMatt Rickard. Specifically, ReLLM and ParserLLM.
Similar libraries such asLMQLandGuidancehave surved as motiviation.DisclaimerKeymaker and its contributors bear no responsibility for any harm done by its usage either directly or indirectly including but not limited to costs incurred by using the package (Keymaker) with LLM vendors. The package is provided "as is" without warranty of any kind, express or implied, including but not limited to the warranties of merchantability, fitness for a particular purpose and noninfringement. In no event shall the authors or copyright holders be liable for any claim, damages or other liability, whether in an action of contract, tort or otherwise, arising from, out of or in connection with the software or the use or other dealings in the software.CopyrightCopyright 2023- Nick Ouellet (nick@ouellet.dev)