-
Notifications
You must be signed in to change notification settings - Fork 430
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
9 changed files
with
197 additions
and
16 deletions.
There are no files selected for viewing
Empty file.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,29 @@ | ||
{{ (messages|selectattr('role', 'equalto', 'system')|list|last).content|trim if (messages|selectattr('role', 'equalto', 'system')|list) else '' }} | ||
|
||
{% for message in messages %} | ||
{% if message['role'] == 'user' %} | ||
### Instruction: | ||
{{ message['content']|trim -}} | ||
{% if not loop.last %} | ||
|
||
|
||
{% endif %} | ||
{% elif message['role'] == 'assistant' %} | ||
### Response: | ||
{{ message['content']|trim -}} | ||
{% if not loop.last %} | ||
|
||
|
||
{% endif %} | ||
{% elif message['role'] == 'user_context' %} | ||
### Input: | ||
{{ message['content']|trim -}} | ||
{% if not loop.last %} | ||
|
||
|
||
{% endif %} | ||
{% endif %} | ||
{% endfor %} | ||
{% if add_generation_prompt and messages[-1]['role'] != 'assistant' %} | ||
### Response: | ||
{% endif %} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,2 @@ | ||
{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %} | ||
{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,7 @@ | ||
{% 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 %} | ||
{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,24 @@ | ||
{%- if messages[0]['role'] == 'system' %} | ||
{%- set system_message = messages[0]['content'] %} | ||
{%- set loop_messages = messages[1:] %} | ||
{%- else %} | ||
{%- set loop_messages = messages %} | ||
{%- endif %} | ||
|
||
{{- bos_token }} | ||
{%- for message in loop_messages %} | ||
{%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) %} | ||
{{- raise_exception('After the optional system message, conversation roles must alternate user/assistant/user/assistant/...') }} | ||
{%- endif %} | ||
{%- if message['role'] == 'user' %} | ||
{%- if loop.first and system_message is defined %} | ||
{{- ' [INST] ' + system_message + '\n\n' + message['content'] + ' [/INST]' }} | ||
{%- else %} | ||
{{- ' [INST] ' + message['content'] + ' [/INST]' }} | ||
{%- endif %} | ||
{%- elif message['role'] == 'assistant' %} | ||
{{- ' ' + message['content'] + eos_token}} | ||
{%- else %} | ||
{{- raise_exception('Only user and assistant roles are supported, with the exception of an initial optional system message!') }} | ||
{%- endif %} | ||
{%- endfor %} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,52 @@ | ||
from typing import Optional | ||
|
||
from common import run_and_check_merge | ||
from test_basic_merges import model_b | ||
from test_tokenizer import model_base | ||
from transformers import AutoTokenizer | ||
|
||
from mergekit.config import InputModelDefinition, MergeConfiguration | ||
|
||
|
||
def check_chat_template(model_path: str, needle: Optional[str] = None): | ||
tokenizer = AutoTokenizer.from_pretrained(model_path) | ||
if needle is None: | ||
assert not tokenizer.chat_template, "Expected no chat template" | ||
return | ||
assert ( | ||
tokenizer.chat_template and needle in tokenizer.chat_template | ||
), f"Expected chat template to contain {needle}" | ||
|
||
|
||
class TestChatTemplate: | ||
def test_template_chatml(self, model_base, model_b): | ||
config = MergeConfiguration( | ||
merge_method="linear", | ||
models=[ | ||
InputModelDefinition(model=model_base, parameters={"weight": 0.5}), | ||
InputModelDefinition(model=model_b, parameters={"weight": 0.5}), | ||
], | ||
base_model=model_base, | ||
dtype="bfloat16", | ||
chat_template="chatml", | ||
) | ||
run_and_check_merge( | ||
config, | ||
validate=lambda p: check_chat_template(p, "<|im_start|>"), | ||
) | ||
|
||
def test_template_literal_jinja(self, model_base, model_b): | ||
config = MergeConfiguration( | ||
merge_method="linear", | ||
models=[ | ||
InputModelDefinition(model=model_base, parameters={"weight": 0.5}), | ||
InputModelDefinition(model=model_b, parameters={"weight": 0.5}), | ||
], | ||
base_model=model_base, | ||
dtype="bfloat16", | ||
chat_template="{{messages[0]['content']}}", | ||
) | ||
run_and_check_merge( | ||
config, | ||
validate=lambda p: check_chat_template(p, "{{messages[0]['content']}}"), | ||
) |