Source code for lmflow.utils.conversation_template.deepseek

#!/usr/bin/env python
# coding=utf-8
# Copyright 2024 Statistics and Machine Learning Research Group. All rights reserved.
from .base import StringFormatter, TemplateComponent, ConversationTemplate




DEEPSEEK_V2_TEMPLATE = ConversationTemplate(
    template_name='deepseek_v2',
    user_formatter=StringFormatter(
        template=[
            TemplateComponent(type='string', content='User: {{content}}\n\n')
        ]
    ),
    assistant_formatter=StringFormatter(
        template=[
            TemplateComponent(type='string', content='Assistant: {{content}}'),
            TemplateComponent(type='token', content='eos_token')
        ]
    ),
    system_formatter=StringFormatter(
        template=[
            TemplateComponent(type='string', content='{{content}}\n\n')
        ]
    ),
    special_starter=TemplateComponent(type='token', content='bos_token')
)





DEEPSEEK_V3_TEMPLATE = (
    "{% if not add_generation_prompt is defined %}"
    "{% set add_generation_prompt = false %}"
    "{% endif %}"
    "{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='', is_first_sp=true) %}"
    "{%- for message in messages %}"
    "{%- if message['role'] == 'system' %}"
    "{%- if ns.is_first_sp %}"
    "{% set ns.system_prompt = ns.system_prompt + message['content'] %}"
    "{% set ns.is_first_sp = false %}"
    "{%- else %}"
    "{% set ns.system_prompt = ns.system_prompt + '\n\n' + message['content'] %}"
    "{%- endif %}"
    "{%- endif %}"
    "{%- endfor %}"
    "{{bos_token}}"
    "{{ns.system_prompt}}"
    "{%- for message in messages %}"
    "{%- if message['role'] == 'user' %}"
    "{%- set ns.is_tool = false -%}"
    "{{'<｜User｜>' + message['content']}}"
    "{%- endif %}"
    "{%- if message['role'] == 'assistant' and message['content'] is none %}"
    "{%- set ns.is_tool = false -%}"
    "{%- for tool in message['tool_calls']%}"
    "{%- if not ns.is_first %}"
    "{{'<｜Assistant｜>'}}"
    "{% generation %}"
    "{{'<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>'}}"
    "{% endgeneration %}"
    "{%- set ns.is_first = true -%}"
    "{%- else %}"
    "{% generation %}"
    "{{'\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>'}}"
    "{{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}}"
    "{% endgeneration %}"
    "{%- endif %}"
    "{%- endfor %}"
    "{%- endif %}"
    "{%- if message['role'] == 'assistant' and message['content'] is not none %}"
    "{%- if ns.is_tool %}"
    "{{'<｜tool▁outputs▁end｜>'}}"
    "{% generation %}"
    "{{ message['content'] + '<｜end▁of▁sentence｜>'}}"
    "{%- set ns.is_tool = false -%}"
    "{% endgeneration %}"
    "{%- else %}"
    "{{'<｜Assistant｜>'}}"
    "{% generation %}"
    "{{ message['content'] + '<｜end▁of▁sentence｜>'}}"
    "{% endgeneration %}"
    "{%- endif %}"
    "{%- endif %}"
    "{%- if message['role'] == 'tool' %}"
    "{%- set ns.is_tool = true -%}"
    "{%- if ns.is_output_first %}"
    "{{'<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}"
    "{%- set ns.is_output_first = false %}"
    "{%- else %}"
    "{{'\n<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}"
    "{%- endif %}"
    "{%- endif %}"
    "{%- endfor -%}"
    "{% if ns.is_tool %}"
    "{{'<｜tool▁outputs▁end｜>'}}"
    "{% endif %}"
    "{% if add_generation_prompt and not ns.is_tool %}"
    "{{'<｜Assistant｜>'}}"
    "{% endif %}"
)





DEEPSEEK_R1_TEMPLATE = (
    "{% if not add_generation_prompt is defined %}"
    "{% set add_generation_prompt = false %}"
    "{% endif %}"
    "{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='', is_first_sp=true) %}"
    "{%- for message in messages %}"
    "{%- if message['role'] == 'system' %}"
    "{%- if ns.is_first_sp %}"
    "{% set ns.system_prompt = ns.system_prompt + message['content'] %}"
    "{% set ns.is_first_sp = false %}"
    "{%- else %}"
    "{% set ns.system_prompt = ns.system_prompt + '\\n\\n' + message['content'] %}"
    "{%- endif %}"
    "{%- endif %}"
    "{%- endfor %}"
    "{{ bos_token }}"
    "{{ ns.system_prompt }}"
    "{%- for message in messages %}"
    "{%- if message['role'] == 'user' %}"
    "{%- set ns.is_tool = false -%}"
    "{{'<｜User｜>' + message['content']}}"
    "{%- endif %}"
    "{%- if message['role'] == 'assistant' and 'tool_calls' in message %}"
    "{%- set ns.is_tool = false -%}"
    "{%- for tool in message['tool_calls'] %}"
    "{%- if not ns.is_first %}"
    "{%- if message['content'] is none %}"
    "    {{'<｜Assistant｜>'}}"
    "    {% generation %}"
    "    {{'<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\\n' + '```json' + '\\n' + tool['function']['arguments'] + '\\n' + '```' + '<｜tool▁call▁end｜>'}}"
    "    {% endgeneration %}"
    "{%- else %}"
    "    {{'<｜Assistant｜>'}}"
    "    {% generation %}"
    "    {{ message['content'] + '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\\n' + '```json' + '\\n' + tool['function']['arguments'] + '\\n' + '```' + '<｜tool▁call▁end｜>'}}"
    "    {% endgeneration %}"
    "{%- endif %}"
    "{%- set ns.is_first = true -%}"
    "{%- else %}"
    "{% generation %}"
    "{{'\\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\\n' + '```json' + '\\n' + tool['function']['arguments'] + '\\n' + '```' + '<｜tool▁call▁end｜>'}}"
    "{% endgeneration %}"
    "{%- endif %}"
    "{%- endfor %}"
    "{% generation %}"
    "{{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}}"
    "{% endgeneration %}"
    "{%- endif %}"
    "{%- if message['role'] == 'assistant' and 'tool_calls' not in message %}"
    "{%- if ns.is_tool %}"
    "{{'<｜tool▁outputs▁end｜>'}}"
    "{% generation %}"
    "{{ message['content'] + '<｜end▁of▁sentence｜>'}}"
    "{% endgeneration %}"
    "{%- set ns.is_tool = false -%}"
    "{%- else %}"
    "{% set content = message['content'] %}"
    "{% if '</think>' in content %}"
    "{% set content = content.split('</think>')[-1] %}"
    "{% endif %}"
    "{{'<｜Assistant｜>'}}"
    "{% generation %}"
    "{{ content + '<｜end▁of▁sentence｜>'}}"
    "{% endgeneration %}"
    "{%- endif %}"
    "{%- endif %}"
    "{%- if message['role'] == 'tool' %}"
    "{%- set ns.is_tool = true -%}"
    "{%- if ns.is_output_first %}"
    "{{'<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}"
    "{%- set ns.is_output_first = false %}"
    "{%- else %}"
    "{{'<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}"
    "{%- endif %}"
    "{%- endif %}"
    "{%- endfor -%}"
    "{% if ns.is_tool %}"
    "{{'<｜tool▁outputs▁end｜>'}}"
    "{% endif %}"
    "{% if add_generation_prompt and not ns.is_tool %}"
    "{{'<｜Assistant｜>'}}"
    "{% endif %}"
)





DEEPSEEK_R1_DISTILL_TEMPLATE = (
    "{% if not add_generation_prompt is defined %}"
    "{% set add_generation_prompt = false %}"
    "{% endif %}"
    "{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='') %}"
    "{%- for message in messages %}"
    "{%- if message['role'] == 'system' %}"
    "{% set ns.system_prompt = message['content'] %}"
    "{%- endif %}"
    "{%- endfor %}"
    "{{bos_token}}"
    "{{ns.system_prompt}}"
    "{%- for message in messages %}"
    "{%- if message['role'] == 'user' %}"
    "{%- set ns.is_tool = false -%}"
    "{{'<｜User｜>' + message['content']}}"
    "{%- endif %}"
    "{%- if message['role'] == 'assistant' and message['content'] is none %}"
    "{%- set ns.is_tool = false -%}"
    "{%- for tool in message['tool_calls']%}"
    "{%- if not ns.is_first %}"
    "{{'<｜Assistant｜>'}}"
    "{% generation %}"
    "{{'<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\\n' + '```json' + '\\n' + tool['function']['arguments'] + '\\n' + '```' + '<｜tool▁call▁end｜>'}}"
    "{% endgeneration %}"
    "{%- set ns.is_first = true -%}"
    "{%- else %}"
    "{% generation %}"
    "{{'\\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\\n' + '```json' + '\\n' + tool['function']['arguments'] + '\\n' + '```' + '<｜tool▁call▁end｜>'}}"
    "{{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}}"
    "{% endgeneration %}"
    "{%- endif %}"
    "{%- endfor %}"
    "{%- endif %}"
    "{%- if message['role'] == 'assistant' and message['content'] is not none %}"
    "{%- if ns.is_tool %}"
    "{{'<｜tool▁outputs▁end｜>'}}"
    "{% generation %}"
    "{{ message['content'] + '<｜end▁of▁sentence｜>'}}"
    "{% endgeneration %}"
    "{%- set ns.is_tool = false -%}"
    "{%- else %}"
    "{% set content = message['content'] %}"
    "{% if '</think>' in content %}"
    "{% set content = content.split('</think>')[-1] %}"
    "{% endif %}"
    "{{'<｜Assistant｜>'}}"
    "{% generation %}"
    "{{ content + '<｜end▁of▁sentence｜>'}}"
    "{% endgeneration %}"
    "{%- endif %}"
    "{%- endif %}"
    "{%- if message['role'] == 'tool' %}"
    "{%- set ns.is_tool = true -%}"
    "{%- if ns.is_output_first %}"
    "{{'<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}"
    "{%- set ns.is_output_first = false %}"
    "{%- else %}"
    "{{'\\n<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}"
    "{%- endif %}"
    "{%- endif %}"
    "{%- endfor -%}"
    "{% if ns.is_tool %}"
    "{{'<｜tool▁outputs▁end｜>'}}"
    "{% endif %}"
    "{% if add_generation_prompt and not ns.is_tool %}"
    "{{'<｜Assistant｜>'}}"
    "{% endif %}"
)