|
11 | 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | 12 | # See the License for the specific language governing permissions and |
13 | 13 | # limitations under the License. |
| 14 | +# flake8: noqa |
14 | 15 |
|
15 | 16 | """ |
16 | 17 | Various models and templates for Marin. |
17 | 18 | """ |
18 | 19 |
|
19 | | -marin_tokenizer = "stanford-crfm/marin-tokenizer" |
| 20 | +marin_tokenizer = "marin-community/marin-tokenizer" |
20 | 21 | """ |
21 | 22 | The HF Hub name for the Marin tokenizer. |
22 | 23 | The Marin tokenizer is (currently) just the Llama 3 tokenizer with a custom chat template (MARIN_CHAT_TEMPLATE). |
23 | 24 | """ |
24 | 25 |
|
25 | | -# to be clear this is the Olmo 2 template except we use llama3's special tokens |
| 26 | +# inspired by the smollm3 template and the Olmo 2 template, using llama3's special tokens |
26 | 27 | MARIN_CHAT_TEMPLATE = """ |
27 | 28 | {{ bos_token }} |
| 29 | +{%- if enable_thinking is defined -%} |
| 30 | + {%- if enable_thinking is sameas true -%} |
| 31 | + {%- set _reasoning_mode = "/think" -%} |
| 32 | + {%- elif enable_thinking is sameas false -%} |
| 33 | + {%- set _reasoning_mode = "/nothink" -%} |
| 34 | + {%- else -%} |
| 35 | + {%- set _reasoning_mode = enable_thinking -%} |
| 36 | + {%- endif -%} |
| 37 | +{%- else -%} |
| 38 | + {%- set _reasoning_mode = none -%} |
| 39 | +{%- endif -%} |
| 40 | +{%- set _custom_instructions = custom_instructions | default(None, true) -%} |
| 41 | +{%- set _xml_tools_list = xml_tools | default([], true) -%} |
| 42 | +{%- if tools is defined and tools -%} |
| 43 | + {%- set _xml_tools_list = tools -%} |
| 44 | +{%- endif -%} |
| 45 | +{%- set _python_tools = python_tools | default([], true) -%} |
| 46 | +{%- set _has_aux_header = (_reasoning_mode is not none) or _custom_instructions or (_xml_tools_list) or (_python_tools) -%} |
| 47 | +{%- if _has_aux_header -%} |
| 48 | +<|start_header_id|>system<|end_header_id|> |
| 49 | +{%- if _reasoning_mode is not none -%} |
| 50 | +Reasoning: {{ _reasoning_mode }} |
| 51 | +{%- endif %} |
| 52 | +{%- if _custom_instructions %} |
| 53 | +{{ _custom_instructions | trim }} |
| 54 | +{%- endif %} |
| 55 | +{% if _xml_tools_list or _python_tools %} |
| 56 | +{{ "\n### Tools\n" }} |
| 57 | +You may call one or more functions to assist with the user query. |
| 58 | +{% if _xml_tools_list %} |
| 59 | +You are provided with function signatures within <tools> </tools> tags: |
| 60 | +
|
| 61 | +<tools> |
| 62 | +{% for tool in _xml_tools_list %} |
| 63 | +{{ tool | string }}{% if not loop.last %} |
| 64 | +{% endif %} |
| 65 | +{% endfor %} |
| 66 | +</tools> |
| 67 | +
|
| 68 | +For each function call, pass a json object with function name and arguments within <tool_call> </tool_call> tags: |
| 69 | +<tool_call> |
| 70 | +{"name": <function-name>, "arguments": <args-json-object>} |
| 71 | +</tool_call> |
| 72 | +
|
| 73 | +{% endif %} |
| 74 | +{% if _python_tools %} |
| 75 | +When you send a message containing Python code between <|python_tag|> and <|eom_id|> tags, it will be executed in a stateful Jupyter notebook environment, and you will then be given the output. |
| 76 | +
|
| 77 | +You can use the following tools in your python code like regular functions: |
| 78 | +<tools> |
| 79 | +{% for tool in _python_tools %} |
| 80 | +{{ tool | string }}{% if not loop.last %} |
| 81 | +{% endif %} |
| 82 | +{% endfor %} |
| 83 | +</tools> |
| 84 | +{% endif %} |
| 85 | +{% endif %} |
| 86 | +<|eot_id|> |
| 87 | +{%- endif -%} |
28 | 88 | {%- for message in messages -%} |
29 | | -{%- if message['role'] == 'assistant' -%} |
30 | | - <|start_header_id|>{{ message['role'] }}<|end_header_id|> |
31 | | -{% generation %}{{- message['content'] | trim }}<|eot_id|>{% endgeneration %}\n |
| 89 | + {%- set has_tool_calls = message.get('tool_calls') is not none and message.get('tool_calls') -%} |
| 90 | + {%- if not (message.get('role') in ['tool', 'ipython'] or has_tool_calls) -%} |
| 91 | + {%- if message.get('role') == 'assistant' -%} |
| 92 | +<|start_header_id|>assistant<|end_header_id|> |
| 93 | +{% set content = message.get('content') %} |
| 94 | +{% if content is string %} |
| 95 | +{% generation %}{{- content | trim }}<|eot_id|>{% endgeneration %} |
| 96 | +{% elif content is mapping %} |
| 97 | +{% generation %}{{- content.get('text', '') | trim }}<|eot_id|>{% endgeneration %} |
| 98 | +{% elif content is iterable %} |
| 99 | +{% generation %} |
| 100 | +{%- for chunk in content -%} |
| 101 | + {%- if chunk.get('type') == 'text' -%} |
| 102 | + {{ chunk.get('text', '') | trim }} |
| 103 | + {%- endif -%} |
| 104 | +{%- endfor -%} |
| 105 | +<|eot_id|> |
| 106 | +{% endgeneration %} |
32 | 107 | {% else %} |
| 108 | +{% generation %}{% endgeneration %}<|eot_id|> |
| 109 | +{% endif %} |
| 110 | + {%- else -%} |
33 | 111 | <|start_header_id|>{{ message['role'] }}<|end_header_id|> |
34 | | -{{ message['content'] | trim }}<|eot_id|> |
| 112 | +{% set content = message.get('content') %} |
| 113 | +{% if content is string %} |
| 114 | +{{ content | trim }}<|eot_id|> |
| 115 | +{% elif content is mapping %} |
| 116 | +{{ content.get('text', '') | trim }}<|eot_id|> |
| 117 | +{% elif content is iterable %} |
| 118 | +{%- for chunk in content -%} |
| 119 | + {%- if chunk.get('type') == 'text' -%} |
| 120 | + {{ chunk.get('text', '') | trim }} |
| 121 | + {%- endif -%} |
| 122 | +{%- endfor -%}<|eot_id|> |
| 123 | +{% else %} |
| 124 | +<|eot_id|> |
| 125 | +{% endif %} |
| 126 | + {%- endif -%} |
| 127 | +
|
| 128 | + {%- elif message.get('role') == 'tool' -%} |
| 129 | + {%- set _tool_name = message.get('name') -%} |
| 130 | + {%- set _tool_id = message.get('tool_call_id') -%} |
| 131 | + {%- set _attr_name = ' name=\"' ~ _tool_name ~ '\"' if _tool_name else '' -%} |
| 132 | + {%- set _attr_id = ' id=\"' ~ _tool_id ~ '\"' if _tool_id else '' -%} |
| 133 | +<|start_header_id|>tool<|end_header_id|> |
| 134 | +<tool_response{{ _attr_name }}{{ _attr_id }}> |
| 135 | +{%- set tool_content = message.get('content') -%} |
| 136 | +{%- if tool_content is mapping or (tool_content is iterable and tool_content is not string) -%} |
| 137 | +{{- tool_content | tojson }} |
| 138 | +{%- else -%} |
| 139 | +{{- tool_content if tool_content is not none else '' }} |
| 140 | +{%- endif -%} |
| 141 | +</tool_response><|eot_id|> |
| 142 | +{{- "\n" -}} |
| 143 | + {%- elif message.get('role') == 'ipython' -%} |
| 144 | +<|start_header_id|>ipython<|end_header_id|> |
| 145 | +{% set ipy_content = message.get('content') %} |
| 146 | +{% if ipy_content is string %} |
| 147 | +{{- { "output": ipy_content } | tojson -}} |
| 148 | +{% elif ipy_content is iterable %} |
| 149 | +{%- for chunk in ipy_content -%} |
| 150 | + {%- if chunk.get('type') == 'text' -%} |
| 151 | + {{- { "output": chunk.get('text', '') } | tojson -}} |
| 152 | + {%- endif -%} |
| 153 | +{%- endfor -%} |
| 154 | +{% else %} |
| 155 | +{{- { "output": ipy_content } | tojson -}} |
35 | 156 | {% endif %} |
| 157 | +<|eot_id|> |
| 158 | +{% elif has_tool_calls -%} |
| 159 | + {%- if message.tool_calls|length != 1 -%} |
| 160 | + {{- raise_exception("This template expects exactly one tool call per assistant turn.") -}} |
| 161 | + {%- endif -%} |
| 162 | + {%- set tool_call = message.tool_calls[0].function -%} |
| 163 | +<|start_header_id|>assistant<|end_header_id|> |
| 164 | +{% generation %} |
| 165 | +{{- '{\"name\": \"' + tool_call.name + '\", \"arguments\": ' -}} |
| 166 | +{{- tool_call.arguments | tojson -}} |
| 167 | +{{- \"}\" -}}<|eot_id|> |
| 168 | +{% endgeneration %} |
| 169 | + {%- endif -%} |
36 | 170 | {%- endfor -%} |
37 | 171 | {%- if add_generation_prompt -%} |
38 | | -<|start_header_id|>assistant<|end_header_id|>\n{% endif -%} |
| 172 | +<|start_header_id|>assistant<|end_header_id|> |
| 173 | +{% endif -%} |
39 | 174 | """.strip() |
40 | 175 |
|
41 | 176 | """ |
|
0 commit comments