fix(factory.py): support llama3 instuct chat template

allows automatic templating for llama3 instruct requests
This commit is contained in:
Krrish Dholakia
2024-04-24 20:35:10 -07:00
parent ffc277c2bb
commit df7db2b870
2 changed files with 44 additions and 5 deletions
+27 -4
View File
@@ -232,7 +232,15 @@ known_tokenizer_config = {
"eos_token": "</s>",
},
"status": "success",
}
},
"meta-llama/Meta-Llama-3-8B-Instruct": {
"tokenizer": {
"chat_template": "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}",
"bos_token": "<|begin_of_text|>",
"eos_token": "",
},
"status": "success",
},
}
@@ -640,7 +648,7 @@ def convert_to_anthropic_tool_invoke_xml(tool_calls: list) -> str:
if get_attribute_or_key(tool, "type") != "function":
continue
tool_function = get_attribute_or_key(tool,"function")
tool_function = get_attribute_or_key(tool, "function")
tool_name = get_attribute_or_key(tool_function, "name")
tool_arguments = get_attribute_or_key(tool_function, "arguments")
parameters = "".join(
@@ -833,8 +841,14 @@ def convert_to_anthropic_tool_invoke(tool_calls: list) -> list:
{
"type": "tool_use",
"id": get_attribute_or_key(tool, "id"),
"name": get_attribute_or_key(get_attribute_or_key(tool, "function"), "name"),
"input": json.loads(get_attribute_or_key(get_attribute_or_key(tool, "function"), "arguments")),
"name": get_attribute_or_key(
get_attribute_or_key(tool, "function"), "name"
),
"input": json.loads(
get_attribute_or_key(
get_attribute_or_key(tool, "function"), "arguments"
)
),
}
for tool in tool_calls
if get_attribute_or_key(tool, "type") == "function"
@@ -1341,6 +1355,14 @@ def prompt_factory(
try:
if "meta-llama/llama-2" in model and "chat" in model:
return llama_2_chat_pt(messages=messages)
elif "meta-llama/llama-3" in model and "instruct" in model:
return hf_chat_template(
model=model,
messages=messages,
chat_template=known_tokenizer_config[ # type: ignore
"meta-llama/Meta-Llama-3-8B-Instruct"
]["tokenizer"]["chat_template"],
)
elif (
"tiiuae/falcon" in model
): # Note: for the instruct models, it's best to use a User: .., Assistant:.. approach in your prompt template.
@@ -1382,6 +1404,7 @@ def prompt_factory(
messages=messages
) # default that covers Bloom, T-5, any non-chat tuned model (e.g. base Llama2)
def get_attribute_or_key(tool_or_function, attribute, default=None):
if hasattr(tool_or_function, attribute):
return getattr(tool_or_function, attribute)
+17 -1
View File
@@ -14,9 +14,24 @@ from litellm.llms.prompt_templates.factory import (
anthropic_messages_pt,
claude_2_1_pt,
llama_2_chat_pt,
prompt_factory,
)
def test_llama_3_prompt():
messages = [
{"role": "system", "content": "You are a good bot"},
{"role": "user", "content": "Hey, how's it going?"},
]
received_prompt = prompt_factory(
model="meta-llama/Meta-Llama-3-8B-Instruct", messages=messages
)
print(f"received_prompt: {received_prompt}")
expected_prompt = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nYou are a good bot<|eot_id|><|start_header_id|>user<|end_header_id|>\n\nHey, how's it going?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"""
assert received_prompt == expected_prompt
def test_codellama_prompt_format():
messages = [
{"role": "system", "content": "You are a good bot"},
@@ -109,6 +124,7 @@ def test_anthropic_messages_pt():
messages = []
with pytest.raises(Exception) as err:
anthropic_messages_pt(messages)
assert("Invalid first message." in str(err.value))
assert "Invalid first message." in str(err.value)
# codellama_prompt_format()