scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix(factory.py): add mistral config to known tokenizer configs

Browse source
This commit is contained in:
Krrish Dholakia 2024-04-18 17:11:38 -07:00
parent eb04a929e6
commit cf1cc71b87
1 changed files with 21 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
litellm/llms/prompt_templates/factory.py
View file

@ -218,6 +218,18 @@ def phind_codellama_pt(messages):
return prompt
known_tokenizer_config = {
"mistralai/Mistral-7B-Instruct-v0.1": {
"tokenizer": {
"chat_template": "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token + ' ' }}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}",
"bos_token": "<s>",
"eos_token": "</s>",
},
"status": "success",
}
}
def hf_chat_template(model: str, messages: list, chat_template: Optional[Any] = None):
# Define Jinja2 environment
env = ImmutableSandboxedEnvironment()
@ -246,20 +258,23 @@ def hf_chat_template(model: str, messages: list, chat_template: Optional[Any] =
else:
return {"status": "failure"}
tokenizer_config = _get_tokenizer_config(model)
if model in known_tokenizer_config:
tokenizer_config = known_tokenizer_config[model]
else:
tokenizer_config = _get_tokenizer_config(model)
if (
tokenizer_config["status"] == "failure"
or "chat_template" not in tokenizer_config["tokenizer"]
):
raise Exception("No chat template found")
## read the bos token, eos token and chat template from the json
tokenizer_config = tokenizer_config["tokenizer"]
bos_token = tokenizer_config["bos_token"]
eos_token = tokenizer_config["eos_token"]
chat_template = tokenizer_config["chat_template"]
tokenizer_config = tokenizer_config["tokenizer"] # type: ignore
bos_token = tokenizer_config["bos_token"] # type: ignore
eos_token = tokenizer_config["eos_token"] # type: ignore
chat_template = tokenizer_config["chat_template"] # type: ignore
try:
template = env.from_string(chat_template)
template = env.from_string(chat_template) # type: ignore
except Exception as e:
raise e