# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Final, final
from typing_extensions import override
from fairseq2.data.tokenizers import (
TokenDecoder,
TokenEncoder,
Tokenizer,
VocabularyInfo,
)
from fairseq2.data.tokenizers.hg import (
HuggingFaceTokenDecoder,
HuggingFaceTokenEncoder,
HuggingFaceTokenModel,
load_hg_token_model,
)
from fairseq2.device import Device
from fairseq2.error import NotSupportedError
# Chat template with assistant mask support.
QWEN_HG_CHAT_TEMPLATE: Final = """{%- if tools %}
{{- '<|im_start|>system\\n' }}
{%- if messages[0].role == 'system' %}
{{- messages[0].content + '\\n\\n' }}
{%- endif %}
{{- "# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>" }}
{%- for tool in tools %}
{{- "\\n" }}
{{- tool | tojson }}
{%- endfor %}
{{- '\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\\n</tool_call><|im_end|>\\n' }}
{%- else %}
{%- if messages[0].role == 'system' %}
{{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}
{%- endif %}
{%- endif %}
{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
{%- for message in messages[::-1] %}
{%- set index = (messages|length - 1) - loop.index0 %}
{%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
{%- set ns.multi_step_tool = false %}
{%- set ns.last_query_index = index %}
{%- endif %}
{%- endfor %}
{%- for message in messages %}
{%- if message.content is string %}
{%- set content = message.content %}
{%- else %}
{%- set content = '' %}
{%- endif %}
{%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
{{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}
{%- elif message.role == "assistant" %}
{%- set reasoning_content = '' %}
{%- if message.reasoning_content is string %}
{%- set reasoning_content = message.reasoning_content %}
{%- else %}
{%- if '</think>' in content %}
{%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}
{%- set content = content.split('</think>')[-1].lstrip('\\n') %}
{%- endif %}
{%- endif %}
{{- '<|im_start|>' + message.role }}
{%- generation %}
{%- if loop.index0 > ns.last_query_index %}
{%- if loop.last or (not loop.last and reasoning_content) %}
{{- '<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}
{%- else %}
{{- content }}
{%- endif %}
{%- else %}
{{- content }}
{%- endif %}
{{- '<|im_end|>' }}
{%- endgeneration %}
{%- endif %}
{%- endfor %}
{%- if add_generation_prompt %}
{{- '<|im_start|>assistant\\n' }}
{%- if enable_thinking is defined and enable_thinking is false %}
{{- '<think>\\n\\n</think>\\n\\n' }}
{%- endif %}
{%- endif %}"""
[docs]
@final
class QwenTokenizer(Tokenizer):
def __init__(self, model: HuggingFaceTokenModel, eos_token: str) -> None:
self._model = model
self._eos_token = eos_token
[docs]
@override
def create_encoder(
self,
*,
task: str | None = None,
lang: str | None = None,
mode: str | None = None,
device: Device | None = None,
pin_memory: bool = False,
) -> TokenEncoder:
if task is not None:
raise ValueError(f"`task` must be `None`, but is '{task}' instead.")
if lang is not None:
raise ValueError(f"`lang` must be `None`, but is '{lang}' instead.")
match mode:
case None | "default":
suffix_tokens = [self._eos_token]
case "prompt":
# In prompt mode, we expect the generator to finish the sequence.
suffix_tokens = []
case "prompt_response":
suffix_tokens = [self._eos_token]
case "as_is":
suffix_tokens = []
case _:
raise NotSupportedError(
f"`mode` must be a supported mode, but is {mode} instead. Supported modes are prompt, prompt_response, as_is."
)
return HuggingFaceTokenEncoder(
self._model,
prefix_tokens=[],
suffix_tokens=suffix_tokens,
device=device,
pin_memory=pin_memory,
)
[docs]
@override
def create_raw_encoder(
self, *, device: Device | None = None, pin_memory: bool = False
) -> TokenEncoder:
return HuggingFaceTokenEncoder(
self._model, device=device, pin_memory=pin_memory
)
[docs]
@override
def create_decoder(self, *, skip_special_tokens: bool = False) -> TokenDecoder:
return HuggingFaceTokenDecoder(
self._model, skip_special_tokens=skip_special_tokens
)
@property
@override
def vocab_info(self) -> VocabularyInfo:
return self._model.vocab_info
[docs]
@dataclass(kw_only=True)
class QwenTokenizerConfig:
use_im_end: bool = False
def load_qwen_tokenizer(path: Path, config: QwenTokenizerConfig) -> Tokenizer:
eos_token = "<|im_end|>" if config.use_im_end else "<|endoftext|>"
model = load_hg_token_model(
path,
unk_token=None,
bos_token=None,
eos_token=eos_token,
pad_token="<|endoftext|>",
boh_token=None,
eoh_token=None,
)
model.overwrite_chat_template(QWEN_HG_CHAT_TEMPLATE)
return QwenTokenizer(model, eos_token)