litellm/litellm/llms/prompt_templates/factory.py
2023-09-18 13:44:19 -07:00

119 lines
No EOL
5.1 KiB
Python

def default_pt(messages):
return " ".join(message["content"] for message in messages)
# Llama2 prompt template
llama_2_special_tokens = ["<s>", "</s>"]
def llama_2_chat_pt(messages):
prompt = custom_prompt(
role_dict={
"system": {
"pre_message": "[INST] <<SYS>>\n",
"post_message": "\n<</SYS>>\n [/INST]\n"
},
"user": { # follow this format https://github.com/facebookresearch/llama/blob/77062717054710e352a99add63d160274ce670c6/llama/generation.py#L348
"pre_message": "[INST] ",
"post_message": " [/INST]\n"
},
"assistant": {
"post_message": "\n" # follows this - https://replicate.com/blog/how-to-prompt-llama
}
},
messages=messages
)
return prompt
# Falcon prompt template - from https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py#L110
def falcon_instruct_pt(messages):
prompt = ""
for message in messages:
if message["role"] == "system":
prompt += message["content"]
else:
prompt += message['role']+":"+ message["content"].replace("\r\n", "\n").replace("\n\n", "\n")
prompt += "\n\n"
return prompt
def falcon_chat_pt(messages):
prompt = ""
for message in messages:
if message["role"] == "system":
prompt += "System: " + message["content"]
elif message["role"] == "assistant":
prompt += "Falcon: " + message["content"]
elif message["role"] == "user":
prompt += "User: " + message["content"]
return prompt
# MPT prompt template - from https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py#L110
def mpt_chat_pt(messages):
prompt = ""
for message in messages:
if message["role"] == "system":
prompt += "<|im_start|>system" + message["content"] + "<|im_end|>" + "\n"
elif message["role"] == "assistant":
prompt += "<|im_start|>assistant" + message["content"] + "<|im_end|>" + "\n"
elif message["role"] == "user":
prompt += "<|im_start|>user" + message["content"] + "<|im_end|>" + "\n"
return prompt
# WizardCoder prompt template - https://huggingface.co/WizardLM/WizardCoder-Python-34B-V1.0#prompt-format
def wizardcoder_pt(messages):
prompt = ""
for message in messages:
if message["role"] == "system":
prompt += message["content"] + "\n\n"
elif message["role"] == "user": # map to 'Instruction'
prompt += "### Instruction:\n" + message["content"] + "\n\n"
elif message["role"] == "assistant": # map to 'Response'
prompt += "### Response:\n" + message["content"] + "\n\n"
return prompt
# Phind-CodeLlama prompt template - https://huggingface.co/Phind/Phind-CodeLlama-34B-v2#how-to-prompt-the-model
def phind_codellama_pt(messages):
prompt = ""
for message in messages:
if message["role"] == "system":
prompt += "### System Prompt\n" + message["content"] + "\n\n"
elif message["role"] == "user":
prompt += "### User Message\n" + message["content"] + "\n\n"
elif message["role"] == "assistant":
prompt += "### Assistant\n" + message["content"] + "\n\n"
return prompt
# Custom prompt template
def custom_prompt(role_dict: dict, messages: list, initial_prompt_value: str="", final_prompt_value: str=""):
prompt = initial_prompt_value
for message in messages:
role = message["role"]
pre_message_str = role_dict[role]["pre_message"] if role in role_dict and "pre_message" in role_dict[role] else ""
post_message_str = role_dict[role]["post_message"] if role in role_dict and "post_message" in role_dict[role] else ""
prompt += pre_message_str + message["content"] + post_message_str
prompt += final_prompt_value
return prompt
def prompt_factory(model: str, messages: list):
model = model.lower()
if "meta-llama/Llama-2" in model:
if "chat" in model:
return llama_2_chat_pt(messages=messages)
elif "tiiuae/falcon" in model: # Note: for the instruct models, it's best to use a User: .., Assistant:.. approach in your prompt template.
if model == "tiiuae/falcon-180B-chat":
return falcon_chat_pt(messages=messages)
elif "instruct" in model:
return falcon_instruct_pt(messages=messages)
elif "mosaicml/mpt" in model:
if "chat" in model:
return mpt_chat_pt(messages=messages)
elif "codellama/codellama" in model:
if "instruct" in model:
return llama_2_chat_pt(messages=messages) # https://huggingface.co/blog/codellama#conversational-instructions
elif "wizardlm/wizardcoder" in model:
return wizardcoder_pt(messages=messages)
elif "phind/phind-codellama" in model:
return phind_codellama_pt(messages=messages)
elif "togethercomputer/llama-2" in model and ("instruct" in model or "chat" in model):
return llama_2_chat_pt(messages=messages)
return default_pt(messages=messages) # default that covers Bloom, T-5, any non-chat tuned model (e.g. base Llama2)