ms-swift ChatML 渲染

register_template(
    QwenTemplateMeta(
        MLLMTemplateType.qwen3_5,
        template_cls=Qwen3_5Template,
        default_system=None,
        thinking_prefix='<think>\n',
        non_thinking_prefix='<think>\n\n</think>\n\n',
        agent_template='qwen3_5',
        is_thinking=True,
    )
)

@dataclass
class TemplateArguments:
    response_prefix: Optional[str] = None
    enable_thinking: Optional[bool] = None
    add_non_thinking_prefix: bool = True

class Template(ProcessorMixin):
    def __init__(
        self,
        processor: Processor,
        template_meta: 'TemplateMeta',
        default_system: Optional[str] = None,
        max_length: Optional[int] = None,
        *,
        # 省略与 thinking 无关的关键字参数
        response_prefix: Optional[str] = None,
        enable_thinking: Optional[bool] = None,
        preserve_thinking: Optional[bool] = None,
        add_non_thinking_prefix: bool = True,
    ) -> None: ...

template_meta = deepcopy(template_meta)

if enable_thinking is None:
    enable_thinking = (
        template_meta.is_thinking and not template_meta.non_thinking_prefix
    )

self.response_prefix = response_prefix
self.template_meta = template_meta
self.enable_thinking = enable_thinking
self.add_non_thinking_prefix = add_non_thinking_prefix

def _get_enable_thinking(self, inputs=None):
    enable_thinking = (
        None
        if inputs is None
        else inputs.chat_template_kwargs.get('enable_thinking')
    )
    if enable_thinking is None:
        enable_thinking = self.enable_thinking
    return enable_thinking


def _get_response_prefix(self, inputs=None):
    response_prefix = (
        None
        if inputs is None
        else inputs.chat_template_kwargs.get('response_prefix')
    )
    if response_prefix is None:
        response_prefix = self.response_prefix
    if response_prefix is not None:
        return response_prefix

    enable_thinking = self._get_enable_thinking(inputs)
    if enable_thinking:
        return self.template_meta.thinking_prefix
    else:
        return self.template_meta.non_thinking_prefix

response_prefix = self._get_response_prefix(inputs)

if response is not None:
    context_list.append('{{RESPONSE}}')
elif response_prefix:
    context_list.append(response_prefix)

outputs = model.generate(**inputs, return_dict_in_generate=True)
input_length = inputs.input_ids.shape[1]
generated_tokens = outputs.sequences[:, input_length:]
response = tokenizer.decode(generated_tokens[0])

for new_token_id in new_token_ids:
    self.token_ids.append(new_token_id)
    self.output_text += self.decode_next(new_token_id)

finished = finish_reason is not None
delta = self.output_kind == RequestOutputKind.DELTA
text = self.detokenizer.get_next_output_text(finished, delta)

def get_generate_ids(
    self, generate_ids: Union[torch.Tensor, List[int]], num_prompt_tokens: int
) -> Union[torch.Tensor, List[int]]:
    if self.skip_prompt:
        generate_ids = generate_ids[..., num_prompt_tokens:]
    return generate_ids


def decode_generate_ids(
    self,
    generate_ids: List[int],
    *,
    is_finished: bool = True,
    first_token=True,
    template_inputs=None,
    **kwargs,
) -> Any:
    if kwargs.get('spaces_between_special_tokens') is None:
        kwargs['spaces_between_special_tokens'] = False
    generate_ids = self.skip_stop_tokens(generate_ids, is_finished)
    response = self.tokenizer.decode(generate_ids, **kwargs)
    response_prefix = self._get_response_prefix(template_inputs)
    if first_token and response_prefix:
        response = response_prefix + response
    return response

for output in result.outputs:
    output.token_ids = list(output.token_ids)
    response = self.template.decode_generate_ids(
        output.token_ids,
        template_inputs=inputs['template_inputs'],
    )

elif role == 'assistant':
    stripped = content.strip()
    if '</think>' in stripped and '<think>' in stripped:
        before, _, after = stripped.partition('</think>')
        reasoning = (
            before.rstrip('\n').rsplit('<think>', 1)[-1].lstrip('\n').strip()
        )
        rest = after.lstrip('\n')
        message['content'] = f'<think>\n{reasoning}\n</think>\n\n{rest}'
    else:
        message['content'] = stripped

def _add_non_thinking_prefix(self, inputs, thinking_prefix='<think>'):
    messages = inputs.messages
    non_thinking_prefix = self.template_meta.non_thinking_prefix

    if non_thinking_prefix:
        if (
            not self.is_training
            or self.loss_scale.base_strategy == 'last_round'
        ) and not self.template_meta.preserve_thinking:
            start_idx = get_last_user_round(messages)
        else:
            start_idx = -1

        for i, message in enumerate(messages):
            if not self._is_add_non_thinking_round(messages, i, start_idx):
                continue

            content = message['content']
            if isinstance(content, str):
                if not content.startswith(
                    (thinking_prefix, non_thinking_prefix)
                ):
                    message['content'] = non_thinking_prefix + content

def _set_loss_scale(self):
    if (not self.disable_ignore_empty_think
            and getattr(self, 'template_meta', None) is not None):
        template_meta = self.template_meta
        if template_meta.is_thinking and template_meta.non_thinking_prefix:
            if self.loss_scale and 'ignore_empty_think' not in self.loss_scale:
                self.loss_scale = self.loss_scale + '+ignore_empty_think'

{
  "^<think>\\s*</think>\\s*": [0.0]
}

for context, loss_weight in zip(
    context_list,
    loss_scale_list,
):
    token_list = self._tokenize(context)
    input_ids += token_list

    if loss_weight > 0.0:
        labels += token_list
    else:
        labels += [-100] * len(token_list)

--add_non_thinking_prefix true \
--loss_scale default+ignore_empty_think

def _preprocess_inputs(
    self, samples: List[OnPolicySample]
) -> List[OnPolicySample]:
    """Preprocess samples before inference"""
    samples = self._set_inputs_system(samples)
    samples = self._add_prompt_id_to_inputs(samples)
    for s in samples:
        remove_response(s.messages)
    return samples

def _generate_completions(
    self, samples: List[OnPolicySample]
) -> List[OnPolicySample]:
    samples = self._preprocess_inputs(samples)

    mode = 'train' if self.model.training else 'eval'
    if self.use_fast_infer:
        samples = self._fast_infer(samples)
    else:
        with (
            unwrap_model_for_generation(
                self.model_wrapped,
                self.accelerator,
                gather_deepspeed3_params=self.args.ds3_gather_for_generation,
            ),
            self.template.generate_context(),
            self.multi_turn_completion_length_context(),
        ):
            samples = self._infer_single_or_multi_turn(
                samples, self.request_config
            )
            if mode == 'train':
                self.model.train()

    return samples

def get_response_prefix_ids(
    template: Template,
    sample_enable_thinking: Optional[bool] = None,
    *,
    chat_template_kwargs: Optional[Dict[str, Any]] = None,
) -> Optional[List[int]]:
    kwargs = dict(chat_template_kwargs or {})
    if sample_enable_thinking is not None:
        kwargs['enable_thinking'] = sample_enable_thinking
    inputs = StdTemplateInputs(messages=[], chat_template_kwargs=kwargs)
    prefix_str = template._get_response_prefix(inputs)
    if prefix_str:
        return template.tokenizer.encode(prefix_str, add_special_tokens=False)
    return None

def encode_sample(
    sample: OnPolicySample,
    template: Template,
    *,
    encode_prompt_only: bool = False,
) -> Dict[str, Any]:
    data = sample.to_template_dict()
    if sample.response_token_ids:
        loss_mask = sample.response_loss_mask or None
        msgs = data.get('messages')
        if msgs is not None:
            msgs = [m.copy() for m in msgs]
        ctk = sample.extra.get('chat_template_kwargs') or {}
        prefix_ids = get_response_prefix_ids(template, chat_template_kwargs=ctk)
        data['messages'] = replace_assistant_response_with_ids(
            msgs,
            sample.response_token_ids,
            loss_mask,
            non_thinking_prefix_ids=prefix_ids,
        )

    if encode_prompt_only:
        messages = data.get('messages', [])
        if messages and messages[-1].get('role') == 'assistant':
            data = {
                **data,
                'messages': messages[:-1] + [{**messages[-1], 'content': None}],
            }

    encoded = template.encode(data, return_length=True)
    return encoded

if non_thinking_prefix_ids:
    n_prefix = len(non_thinking_prefix_ids)
    last_ids = list(completion_ids[-1])

    if last_ids[:n_prefix] != list(non_thinking_prefix_ids):
        if loss_mask is None:
            loss_mask = [[1] * len(ids) for ids in completion_ids]
        completion_ids[-1] = list(non_thinking_prefix_ids) + last_ids
        loss_mask[-1] = [0] * n_prefix + list(loss_mask[-1])

{"chat_template_kwargs": {"enable_thinking": true}}
{"chat_template_kwargs": {"enable_thinking": false}}

@dataclass
class ChatmlTemplateMeta(TemplateMeta):
    prompt: Prompt = field(
        default_factory=lambda: [
            '<|im_start|>user\n{{QUERY}}<|im_end|>\n' '<|im_start|>assistant\n'
        ]
    )
    chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<|im_end|>\n'])
    suffix: Prompt = field(default_factory=lambda: ['<|im_end|>\n'])
    system_prefix: Optional[Prompt] = field(
        default_factory=lambda: ['<|im_start|>system\n{{SYSTEM}}<|im_end|>\n']
    )


class Qwen3OmniTemplate(Qwen2_5OmniTemplate):
    version = 'omni_v3'
    placeholder_tokens = [
        '<|image_pad|>',
        '<|audio_pad|>',
        '<|video_pad|>',
    ]

{
  "messages": [
    {"role": "system", "content": "你是个有用无害的助手"},
    {
      "role": "user",
      "content": "<image>图片中是什么，<video>视频中是什么，<audio>音频中是什么"
    },
    {
      "role": "assistant",
      "content": "图片中是一头大象，视频中是一只小狗在草地上奔跑，音频中有鸟鸣声"
    }
  ],
  "images": ["/xxx/x.jpg"],
  "videos": ["/xxx/x.mp4"],
  "audios": ["/xxx/x.wav"]
}

context_list, loss_scale_list = self._split_special_tokens(
    context_list, loss_scale_list
)
context_list, loss_scale_list = self._pre_tokenize(
    context_list, loss_scale_list, inputs
)

def replace_tag(self, media_type: Literal['image', 'video', 'audio'],
                index: int, inputs: StdTemplateInputs) -> List[Context]:
    from qwen_omni_utils import fetch_image, fetch_video
    kwargs = (
        {'image_patch_size': self.processor.image_processor.patch_size}
        if self.version == 'omni_v3' else {}
    )
    sampling_rate = inputs.chat_template_kwargs.get('sampling_rate')
    if sampling_rate is None:
        sampling_rate = self.sampling_rate
    if self.mode == 'vllm':
        inputs.mm_processor_kwargs['do_resize'] = False
    if media_type == 'image':
        inputs.images[index] = fetch_image(
            {'image': inputs.images[index], **inputs.chat_template_kwargs},
            **kwargs,
        )
        if self.version == 'omni_v2_5':
            return ['<|vision_bos|><|IMAGE|><|vision_eos|>']
        elif self.version == 'omni_v3':
            return ['<|vision_start|><|image_pad|><|vision_end|>']
    elif media_type == 'audio':
        if self.mode != 'vllm':
            inputs.audios[index] = load_audio(inputs.audios[index], sampling_rate)
        if self.version == 'omni_v2_5':
            return ['<|audio_bos|><|AUDIO|><|audio_eos|>']
        elif self.version == 'omni_v3':
            return ['<|audio_start|><|audio_pad|><|audio_end|>']

    elif media_type == 'video':
        video = inputs.videos[index]
        video_inputs = {'video': video, **inputs.chat_template_kwargs}
        if isinstance(video, list):
            from qwen_omni_utils import vision_process
            video_inputs['sample_fps'] = vision_process.FPS
        _video, sample_fps = fetch_video(
            video_inputs, return_video_sample_fps=True, **kwargs
        )
        inputs.mm_processor_kwargs.setdefault('fps', []).append(sample_fps)
        if isinstance(_video, torch.Tensor):
            _video = _video.to(torch.uint8)
        inputs.videos[index] = _video
        if self.version == 'omni_v2_5':
            return ['<|vision_bos|><|VIDEO|><|vision_eos|>']
        elif self.version == 'omni_v3':
            return ['<|vision_start|><|video_pad|><|vision_end|>']

encoded = Template._encode(self, inputs)
inputs.audios = self._trim_omni_v3_audios(inputs.audios)

media_inputs = self.processor(
    text='',
    audio=inputs.audios or None,
    images=inputs.images or None,
    videos=inputs.videos or None,
    do_resize=False,
    return_tensors='pt',
)
media_inputs.pop('input_ids')
media_inputs.pop('attention_mask')

input_ids = encoded['input_ids']
labels = encoded['labels']
loss_scale = encoded.get('loss_scale')
config = self.config.thinker_config

def _get_new_tokens(i):
    token_len = media_grid_thw[i].prod() // (merge_size**2)
    return token_id * token_len

def _get_feat_extract_output_lengths(self, input_lengths):
    if self.version == 'omni_v2_5':
        return ((input_lengths - 1) // 2 + 1 - 2) // 2 + 1
    elif self.version == 'omni_v3':
        input_lengths_leave = input_lengths % 100
        feat_lengths = (input_lengths_leave - 1) // 2 + 1
        return (
            ((feat_lengths - 1) // 2 + 1 - 1) // 2
            + 1
            + (input_lengths // 100) * 13
        )

merge_size = 2

image_grid_thw = torch.tensor([[1, 32, 48]])
video_grid_thw = torch.tensor([[8, 24, 40]])
feature_attention_mask = torch.ones((1, 300), dtype=torch.long)

image_placeholder_len = int(image_grid_thw[0].prod() // (merge_size**2))
video_placeholder_len = int(video_grid_thw[0].prod() // (merge_size**2))
audio_placeholder_len = int(
    self._get_feat_extract_output_lengths(feature_attention_mask.sum(dim=1))[0]
)

assert image_placeholder_len == 384
assert video_placeholder_len == 1920
assert audio_placeholder_len == 39

vision_start_id = self._tokenize('<|vision_start|>')[0]
vision_end_id = self._tokenize('<|vision_end|>')[0]
audio_start_id = self._tokenize('<|audio_start|>')[0]
audio_end_id = self._tokenize('<|audio_end|>')[0]

image_segment = [
    vision_start_id,
    config.image_token_id,
    vision_end_id,
]
video_segment = [
    vision_start_id,
    config.video_token_id,
    vision_end_id,
]
audio_segment = [
    audio_start_id,
    config.audio_token_id,
    audio_end_id,
]

image_segment = (
    [vision_start_id]
    + [config.image_token_id] * image_placeholder_len
    + [vision_end_id]
)
video_segment = (
    [vision_start_id]
    + [config.video_token_id] * video_placeholder_len
    + [vision_end_id]
)
audio_segment = (
    [audio_start_id]
    + [config.audio_token_id] * audio_placeholder_len
    + [audio_end_id]
)

assert len(image_segment) == 386
assert len(video_segment) == 1922
assert len(audio_segment) == 41

new_tokens = get_new_tokens(i)
token_len = len(new_tokens)

input_ids = (
    input_ids[: idx + added_tokens_len]
    + new_tokens
    + input_ids[added_tokens_len + idx + 1 :]
)
labels = (
    labels[: idx + added_tokens_len]
    + [-100] * token_len
    + labels[added_tokens_len + idx + 1 :]
)

encoded['input_ids'] = input_ids
encoded['labels'] = labels
encoded['loss_scale'] = loss_scale
encoded.update(media_inputs)
return encoded

seq_lens = [len(seq) for seq in res['input_ids']]
res['attention_mask'] = [
    torch.ones(seq_len, dtype=torch.int64) for seq_len in seq_lens
]

res['input_ids'] = self._pad_sequence(
    res['input_ids'], self.tokenizer.pad_token_id
)
res['attention_mask'] = self._pad_sequence(res['attention_mask'], 0)

pixel_values = [
    b['pixel_values'] for b in batch if b.get('pixel_values') is not None
]
res['pixel_values'] = torch.concat(pixel_values)

pixel_values_videos = [
    b['pixel_values_videos']
    for b in batch
    if b.get('pixel_values_videos') is not None
]
res['pixel_values_videos'] = torch.concat(pixel_values_videos)

res['image_grid_thw'] = self.concat_tensor(batch, 'image_grid_thw', 0)
res['video_grid_thw'] = self.concat_tensor(batch, 'video_grid_thw', 0)

max_length = max(x.shape[-1] for x in input_features)
for i, input_feature in enumerate(input_features):
    mask = feature_attention_mask[i]
    input_features[i] = F.pad(
        input_feature,
        (0, max_length - input_feature.shape[-1]),
    )
    feature_attention_mask[i] = F.pad(
        mask,
        (0, max_length - mask.shape[-1]),
    )

res['input_features'] = torch.concat(input_features)
res['feature_attention_mask'] = torch.concat(feature_attention_mask)

feature_attention_mask = inputs.get('feature_attention_mask')
audio_feature_lengths = torch.sum(feature_attention_mask, dim=1)

position_ids, _ = self._get_get_rope_index()(
    input_ids,
    inputs.get('image_grid_thw'),
    inputs.get('video_grid_thw'),
    attention_mask,
    self.use_audio_in_video,
    audio_feature_lengths,
    video_second_per_grid,
)

text_position_ids = torch.arange(seq_len, device=position_ids.device).expand(
    1, *position_ids.shape[1:]
)

position_ids = torch.concat([text_position_ids, position_ids], dim=0)

class Qwen3OmniTemplate(Qwen2_5OmniTemplate):
    version = 'omni_v3'
    norm_bbox = 'norm1000'
    placeholder_tokens = [
        '<|image_pad|>',
        '<|audio_pad|>',
        '<|video_pad|>',
    ]

    def _post_encode(self, model, inputs):
        return inputs

if inputs_embeds is None:
    inputs_embeds = self.get_input_embeddings()(input_ids)

audio_features = self.get_audio_features(
    input_features,
    feature_attention_mask,
    audio_feature_lengths,
    return_dict=True,
).last_hidden_state

_, _, audio_mask = self.get_placeholder_mask(
    input_ids,
    inputs_embeds=inputs_embeds,
)

inputs_embeds = inputs_embeds.masked_scatter(
    audio_mask,
    audio_features,
)

image_outputs = self.get_image_features(
    pixel_values,
    image_grid_thw,
    return_dict=True,
)
image_embeds = image_outputs.pooler_output
image_mask, _, _ = self.get_placeholder_mask(
    input_ids,
    inputs_embeds=inputs_embeds,
    image_features=image_embeds,
)
inputs_embeds = inputs_embeds.masked_scatter(
    image_mask,
    image_embeds,
)

video_outputs = self.get_video_features(
    pixel_values_videos,
    video_grid_thw,
    return_dict=True,
)
video_embeds = video_outputs.pooler_output
_, video_mask, _ = self.get_placeholder_mask(
    input_ids,
    inputs_embeds=inputs_embeds,
    video_features=video_embeds,
)
inputs_embeds = inputs_embeds.masked_scatter(
    video_mask,
    video_embeds,
)

special_image_mask = input_ids == self.config.image_token_id
special_video_mask = input_ids == self.config.video_token_id
special_audio_mask = input_ids == self.config.audio_token_id

n_image_tokens = special_image_mask.sum()
n_video_tokens = special_video_mask.sum()

special_image_mask = special_image_mask.unsqueeze(-1)
special_video_mask = special_video_mask.unsqueeze(-1)
special_audio_mask = special_audio_mask.unsqueeze(-1)

torch_compilable_check(
    n_image_tokens * inputs_embeds.shape[-1] == image_features.numel(),
    "Image features and image tokens do not match",
)

visual_pos_masks = video_mask | image_mask

if deepstack_visual_embeds is not None and layer_idx in range(
    len(deepstack_visual_embeds)
):
    hidden_states = self._deepstack_process(
        hidden_states,
        visual_pos_masks,
        deepstack_visual_embeds[layer_idx],
    )