14
14
15
15
from ._utils import suppress_stdout_stderr , Singleton
16
16
17
+ ### Common Chat Templates and Special Tokens ###
18
+
19
+ # Source: https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/blob/main/tokenizer_config.json
20
+ CHATML_CHAT_TEMPLATE = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n ' + message['content'] + '<|im_end|>' + '\n '}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n ' }}{% endif %}"
21
+ CHATML_BOS_TOKEN = "<s>"
22
+ CHATML_EOS_TOKEN = "<|im_end|>"
23
+
24
+ # Source: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1/blob/main/tokenizer_config.json
25
+ MISTRAL_INSTRUCT_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 %}"
26
+ MISTRAL_INSTRUCT_BOS_TOKEN = "<s>"
27
+ MISTRAL_INSTRUCT_EOS_TOKEN = "</s>"
28
+
29
+
30
+ ### Chat Completion Handler ###
17
31
18
32
class LlamaChatCompletionHandler (Protocol ):
19
33
"""Base Protocol for a llama chat completion handler.
@@ -118,7 +132,6 @@ def decorator(f: LlamaChatCompletionHandler):
118
132
119
133
### Chat Formatter ###
120
134
121
-
122
135
@dataclasses .dataclass
123
136
class ChatFormatterResponse :
124
137
"""Dataclass that stores completion parameters for a given chat format and
@@ -440,7 +453,20 @@ def hf_tokenizer_config_to_chat_completion_handler(
440
453
return chat_formatter_to_chat_completion_handler (chat_formatter )
441
454
442
455
456
+ def guess_chat_format_from_gguf_metadata (metadata : Dict [str , str ]) -> Optional [str ]:
457
+ if "tokenizer.chat_template" not in metadata :
458
+ return None
459
+
460
+ if metadata ["tokenizer.chat_template" ] == CHATML_CHAT_TEMPLATE :
461
+ return "chatml"
462
+
463
+ if metadata ["tokenizer.chat_template" ] == MISTRAL_INSTRUCT_CHAT_TEMPLATE :
464
+ return "mistral-instruct"
465
+
466
+ return None
467
+
443
468
### Utility functions for formatting chat prompts ###
469
+ # TODO: Replace these with jinja2 templates
444
470
445
471
446
472
def _get_system_message (
@@ -929,7 +955,6 @@ def format_openchat(
929
955
_prompt = _format_chatml (system_message , _messages , _sep )
930
956
return ChatFormatterResponse (prompt = _prompt , stop = _sep )
931
957
932
-
933
958
# Chat format for Saiga models, see more details and available models:
934
959
# https://huggingface.co/collections/IlyaGusev/saiga2-saigamistral-6505d4ccc3d1e53166b636cd
935
960
@register_chat_format ("saiga" )
@@ -951,6 +976,7 @@ def format_saiga(
951
976
_prompt += "<s>bot"
952
977
return ChatFormatterResponse (prompt = _prompt .strip ())
953
978
979
+ # Tricky chat formats that require custom chat handlers
954
980
955
981
@register_chat_completion_handler ("functionary" )
956
982
def functionary_chat_handler (
0 commit comments