/swift/template/templates/deepseek.py

# Copyright (c) ModelScope Contributors. All rights reserved.
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass, field
from PIL import Image, ImageOps
from transformers.dynamic_module_utils import get_class_from_dynamic_module
from typing import Any, Dict, List, Optional

from swift.utils import get_env_args
from ..base import Template
from ..constant import LLMTemplateType, MLLMTemplateType
from ..register import TemplateMeta, register_template
from ..template_inputs import StdTemplateInputs
from ..utils import Prompt, findall


@dataclass
class DeepseekTemplateMeta(TemplateMeta):
    prefix: Prompt = field(default_factory=lambda: [['bos_token_id']])
    prompt: Prompt = field(default_factory=lambda: ['User: {{QUERY}}\n\nAssistant:'])
    chat_sep: Optional[Prompt] = field(default_factory=lambda: [['eos_token_id']])
    suffix: Prompt = field(default_factory=lambda: [['eos_token_id']])
    system_prefix: Optional[Prompt] = field(default_factory=lambda: [['bos_token_id'], '{{SYSTEM}}\n\n'])


register_template(DeepseekTemplateMeta(LLMTemplateType.deepseek, ))

register_template(
    TemplateMeta(
        LLMTemplateType.deepseek_coder,
        prefix=['{{SYSTEM}}'],
        prompt=['### Instruction:\n{{QUERY}}\n### Response:\n'],
        chat_sep=['\n<|EOT|>\n'],
        suffix=['\n<|EOT|>'],
        stop_words=['<|EOT|>'],
        default_system=('You are an AI programming assistant, utilizing the Deepseek Coder model, '
                        'developed by Deepseek Company, and you only answer questions related to computer science. '
                        'For politically sensitive questions, security and privacy issues, '
                        'and other non-computer science questions, you will refuse to answer\n')))


class DeepseekVLTemplate(Template):
    image_placeholder = ['<image_placeholder>']
    skip_prompt = False
    use_model = True
    placeholder_tokens = ['<image_placeholder>']

    image_token_num_per_image: int = 576

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        is_janus = getattr(self, 'is_janus', False)

        encoded = super()._encode(inputs)
        images = inputs.images
        processor = self.processor
        input_ids, labels = encoded['input_ids'], encoded['labels']

        if not inputs.generate_mode:  # understanding task
            idx_list = findall(input_ids, processor.image_id)  # '<image_placeholder>'
            new_input_ids, new_labels = [], []
            lo = 0
            for hi in idx_list:
                new_input_ids += input_ids[lo:hi]
                if labels is not None:
                    new_labels += labels[lo:hi]
                image_tokens = [processor.image_id] * processor.num_image_tokens
                if is_janus:
                    image_tokens = [processor.image_start_id] + image_tokens + [processor.image_end_id]
                new_input_ids += image_tokens
                new_labels += [-100] * len(image_tokens)
                lo = hi + 1
            new_input_ids += input_ids[lo:]
            if labels is not None:
                new_labels += labels[lo:]
            else:
                new_labels = None
            if is_janus:
                from janus.models.processing_vlm import VLChatProcessorOutput
            else:
                from deepseek_vl.models.processing_vlm import VLChatProcessorOutput

            images_outputs = processor.image_processor(images, return_tensors='pt')
            output = VLChatProcessorOutput(
                sft_format=None,
                input_ids=torch.tensor(new_input_ids),
                pixel_values=images_outputs.pixel_values,
                num_image_tokens=torch.tensor([processor.num_image_tokens] * len(idx_list)))
            encoded = {'output': output, 'input_ids': new_input_ids, 'labels': new_labels}
            return encoded

        else:  # image generation task
            if self.is_training:
                raise NotImplementedError('Only support the inference of generation of Janus series models.')
            sft_format = self.tokenizer.decode(input_ids)
            prompt = sft_format + processor.image_start_tag
            input_ids = processor.tokenizer.encode(prompt)
            input_ids = torch.LongTensor(input_ids)

            encoded = {'input_ids': input_ids, 'labels': labels, 'generate_mode': inputs.generate_mode}
            return encoded

    def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]:
        if not inputs.get('generate_mode'):
            inputs['pixel_values'] = inputs['pixel_values'].to(dtype=self.model_info.torch_dtype)
            inputs_embeds = model.prepare_inputs_embeds(**inputs)
            return {'inputs_embeds': inputs_embeds}
        else:
            return inputs

    def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
        gene_img_list = [b.get('generate_mode') for b in batch]
        if all(gene_img_list):
            generate_mode = True
        elif not any(gene_img_list):
            generate_mode = False
        else:
            raise NotImplementedError('Do not support understanding and image generation tasks in one batch.')

        if not generate_mode:
            output = self.fetch_inputs(batch, ['output'])['output']
            batched_output = dict(self.processor.batchify(output))
            res = super()._data_collator(batch, padding_to=padding_to)
            return {**batched_output, **res}
        else:
            res = super()._data_collator(batch, padding_to=padding_to)
            res['generate_mode'] = generate_mode
            return res

    def generate(self, model, *args, **kwargs):
        if not kwargs.get('generate_mode'):
            return super().generate(model, *args, **kwargs)

        else:
            # generate how many number of images for each prompt, it is named parallel_size in the author's code
            parallel_size = kwargs['generation_config'].num_return_sequences
            temperature = kwargs['generation_config'].temperature
            cfg_weight = get_env_args('cfg_weight', float, 5.0)

            input_ids = kwargs['input_ids']  # [bsz, max_input_token_num]
            bsz, max_input_token_num = input_ids.shape
            tokens = torch.zeros((bsz, parallel_size * 2, max_input_token_num),
                                 dtype=torch.int).cuda()  # [bsz, parallel_size*2, max_input_token_num]
            for i in range(parallel_size * 2):
                tokens[:, i, :] = input_ids
                if i % 2 != 0:
                    tokens[:, i, 1:-1] = self.processor.pad_id

            inputs_embeds = model.language_model.get_input_embeddings()(
                tokens)  # [bsz, parallel_size*2, max_input_token_num, 2048]

            generated_tokens = torch.zeros(
                (bsz, parallel_size, self.image_token_num_per_image),
                dtype=torch.int).cuda()  # [bsz, 16, image_token_num_per_image] placeholder for the generated tokens

            # set the first two dimensions into one dimension for batch size
            inputs_embeds = inputs_embeds.reshape(bsz * parallel_size * 2, max_input_token_num, -1)
            generated_tokens = generated_tokens.reshape(bsz * parallel_size, self.image_token_num_per_image)

            for i in range(self.image_token_num_per_image):  # generate the tokens of image in a auto-regression way
                outputs = model.language_model.model(
                    inputs_embeds=inputs_embeds,
                    use_cache=True,
                    past_key_values=outputs.past_key_values if i != 0 else None)
                hidden_states = outputs.last_hidden_state

                logits = self.model.gen_head(hidden_states[:, -1, :])
                logit_cond = logits[0::2, :]
                logit_uncond = logits[1::2, :]

                logits = logit_uncond + cfg_weight * (logit_cond - logit_uncond)
                probs = torch.softmax(logits / temperature, dim=-1)

                next_token = torch.multinomial(probs, num_samples=1)
                generated_tokens[:, i] = next_token.squeeze(dim=-1)  # [parallel_size, self.image_token_num_per_image]

                next_token = torch.cat([next_token.unsqueeze(dim=1), next_token.unsqueeze(dim=1)], dim=1).view(-1)
                img_embeds = model.prepare_gen_img_embeds(next_token)  # [parallel_size * 2, 2048]
                inputs_embeds = img_embeds.unsqueeze(dim=1)  # [parallel_size * 2, 1, 2048]

            # no need to reset the original first two dimensions, waiting for the update of the upper layer
            # inputs_embeds = inputs_embeds.reshape(bsz, parallel_size*2, -1)
            # generated_tokens = generated_tokens.reshape(bsz, parallel_size, self.image_token_num_per_image)

            return {'sequences': generated_tokens}

    def decode_generate_ids(self, generate_ids: List[int], **kwargs) -> Any:
        if 'template_inputs' not in kwargs or not kwargs['template_inputs'].generate_mode:
            return super().decode_generate_ids(generate_ids, **kwargs)
        else:
            img_size = get_env_args('img_size', int, 384)
            patch_size = 16

            num_to_decode = 1  # for now, generate_ids is a 1D list

            generate_ids = torch.tensor(generate_ids).unsqueeze(0)  # [num_to_decode=1, self.image_token_num_per_image]

            dec = self.model.gen_vision_model.decode_code(
                generate_ids.to(dtype=torch.int),
                shape=[num_to_decode, 8, img_size // patch_size, img_size // patch_size])
            dec = dec.to(torch.float32).cpu().numpy().transpose(0, 2, 3, 1)  # [num_to_decode, H, W, ch=3]

            dec = np.clip((dec + 1) / 2 * 255, 0, 255)

            visual_img = np.zeros((num_to_decode, img_size, img_size, 3), dtype=np.uint8)
            visual_img[:, :, :] = dec

            img_list = []
            for i in range(num_to_decode):
                cur_img = Image.fromarray(visual_img[i])
                img_list.append({'type': 'image', 'image': cur_img})
            return img_list


@dataclass
class DeepseekVLTemplateMeta(DeepseekTemplateMeta):
    default_system: Optional[str] = ('You are a helpful language and vision assistant. '
                                     'You are able to understand the visual content that the user provides, '
                                     'and assist the user with a variety of tasks using natural language.')


register_template(DeepseekVLTemplateMeta(
    MLLMTemplateType.deepseek_vl,
    template_cls=DeepseekVLTemplate,
))


class DeepseekJanus(DeepseekVLTemplate):
    is_janus = True
    image_placeholder = ['<image_placeholder>\n']


register_template(DeepseekVLTemplateMeta(MLLMTemplateType.deepseek_janus, template_cls=DeepseekJanus))


class DeepseekOCR(Template):
    version = 'v1'
    image_placeholder = ['<image>\n']

    def init_env_args(self):
        # Delay loading dynamic modules that require specific transformers versions
        # These will be loaded lazily in _preprocess_image when actually needed
        # This avoids triggering transformers version compatibility issues for vllm backend
        super().init_env_args()
        self._BasicImageTransform = None
        self._dynamic_preprocess = None
        self.crop_mode = get_env_args('crop_mode', bool, True)
        self.base_size = get_env_args('base_size', int, 1024)
        # image_size will be set after detecting version (v1: 640, v2: 768)
        self._image_size_override = get_env_args('image_size', int, None)

    @property
    def image_size(self):
        if self._image_size_override is not None:
            return self._image_size_override
        return 768 if self.version == 'v2' else 640

    @property
    def crop_threshold(self):
        # v1: 640, v2: 768
        return 768 if self.version == 'v2' else 640

    def _load_dynamic_modules(self):
        """Lazily load dynamic modules from model repository."""
        if self._BasicImageTransform is None:
            model_dir = self.model_info.model_dir
            model_type_name = 'deepseekocr2' if self.version == 'v2' else 'deepseekocr'
            self._BasicImageTransform = get_class_from_dynamic_module(f'modeling_{model_type_name}.BasicImageTransform',
                                                                      model_dir)
            self._dynamic_preprocess = get_class_from_dynamic_module(f'modeling_{model_type_name}.dynamic_preprocess',
                                                                     model_dir)

    @property
    def BasicImageTransform(self):
        self._load_dynamic_modules()
        return self._BasicImageTransform

    @property
    def dynamic_preprocess(self):
        self._load_dynamic_modules()
        return self._dynamic_preprocess

    def _preprocess_image(self, images, image_token_id):
        # Code borrowed from
        # https://modelscope.cn/models/deepseek-ai/DeepSeek-OCR/file/view/master/modeling_deepseekocr.py?status=1
        # https://modelscope.cn/models/deepseek-ai/DeepSeek-OCR-2/file/view/master/modeling_deepseekocr2.py?status=1
        crop_mode = self.crop_mode
        patch_size = 16
        downsample_ratio = 4
        valid_img_tokens = 0
        w, h = images[0].size
        ratio = 1 - ((max(w, h) - min(w, h)) / (max(w, h)))
        crop_threshold = self.crop_threshold
        image_size = self.image_size

        image_transform = self.BasicImageTransform(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), normalize=True)
        images_list, images_crop_list = [], []
        tokenized_str = []
        images_spatial_crop = []
        for image in images:
            if crop_mode:
                if image.size[0] <= crop_threshold and image.size[1] <= crop_threshold:
                    crop_ratio = [1, 1]
                else:
                    if crop_mode:
                        images_crop_raw, crop_ratio = self.dynamic_preprocess(image)
                    else:
                        crop_ratio = [1, 1]
                """process the global view"""
                global_view = ImageOps.pad(
                    image, (self.base_size, self.base_size), color=tuple(int(x * 255) for x in image_transform.mean))

                if self.base_size == 1024:
                    valid_img_tokens += int(256 * ratio)
                elif self.base_size == 1280:
                    valid_img_tokens += int(400 * ratio)

                images_list.append(image_transform(global_view).to(torch.bfloat16))
                width_crop_num, height_crop_num = crop_ratio

                images_spatial_crop.append([width_crop_num, height_crop_num])

                if width_crop_num > 1 or height_crop_num > 1:
                    """process the local views"""
                    for i in range(len(images_crop_raw)):
                        images_crop_list.append(image_transform(images_crop_raw[i]).to(torch.bfloat16))

                if image_size == 640:
                    valid_img_tokens += len(images_crop_list) * 100
                elif image_size == 768:
                    valid_img_tokens += len(images_crop_list) * 144

                num_queries = math.ceil((image_size // patch_size) / downsample_ratio)
                num_queries_base = math.ceil((self.base_size // patch_size) / downsample_ratio)
                """add image tokens"""
                # v1: adds newline token after each row, v2: no newline tokens in rows
                if self.version == 'v2':
                    tokenized_image = ([image_token_id] * num_queries_base) * num_queries_base
                    tokenized_image += [image_token_id]
                    if width_crop_num > 1 or height_crop_num > 1:
                        tokenized_image += ([image_token_id] * (num_queries * width_crop_num)) * (
                            num_queries * height_crop_num)
                else:
                    tokenized_image = ([image_token_id] * num_queries_base + [image_token_id]) * num_queries_base
                    tokenized_image += [image_token_id]
                    if width_crop_num > 1 or height_crop_num > 1:
                        tokenized_image += ([image_token_id] * (num_queries * width_crop_num) + [image_token_id]) * (
                            num_queries * height_crop_num)
                tokenized_str.append(tokenized_image)
            else:
                """process the global view"""
                if image_size <= crop_threshold:
                    image = image.resize((image_size, image_size))
                global_view = ImageOps.pad(
                    image, (image_size, image_size), color=tuple(int(x * 255) for x in image_transform.mean))
                images_list.append(image_transform(global_view).to(torch.bfloat16))

                if self.base_size == 1024:
                    valid_img_tokens += int(256 * ratio)
                elif self.base_size == 1280:
                    valid_img_tokens += int(400 * ratio)
                elif self.base_size == 640:
                    valid_img_tokens += int(100 * 1)
                elif self.base_size == 512:
                    valid_img_tokens += int(64 * 1)
                elif self.base_size == 768:
                    valid_img_tokens += int(144 * 1)

                width_crop_num, height_crop_num = 1, 1

                images_spatial_crop.append([width_crop_num, height_crop_num])
                """add image tokens"""
                num_queries = math.ceil((image_size // patch_size) / downsample_ratio)

                # v1: adds newline token after each row, v2: no newline tokens in rows
                if self.version == 'v2':
                    tokenized_image = ([image_token_id] * num_queries) * num_queries
                    tokenized_image += [image_token_id]
                else:
                    tokenized_image = ([image_token_id] * num_queries + [image_token_id]) * num_queries
                    tokenized_image += [image_token_id]
                tokenized_str.append(tokenized_image)
        if len(images_list) == 0:
            images_ori = torch.zeros((1, 3, self.image_size, self.image_size))
            images_spatial_crop = torch.zeros((1, 2), dtype=torch.long)
            images_crop = torch.zeros((1, 3, self.base_size, self.base_size))

        else:
            images_ori = torch.stack(images_list, dim=0)
            images_spatial_crop = torch.tensor(images_spatial_crop, dtype=torch.long)
            if images_crop_list:
                images_crop = torch.stack(images_crop_list, dim=0)
            else:
                images_crop = torch.zeros((1, 3, self.base_size, self.base_size))
        return tokenized_str, images_ori, images_crop, images_spatial_crop

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        encoded = super()._encode(inputs)
        input_ids = encoded['input_ids']
        labels = encoded['labels']
        loss_scale = encoded.get('loss_scale', None)
        image_token = self._tokenize('<image>')
        idx_list = findall(input_ids, image_token)
        if idx_list:
            tokenized_str, images_ori, images_crop, images_spatial_crop = self._preprocess_image(
                inputs.images, image_token[0])
            input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list,
                                                                lambda i: tokenized_str[i])
            encoded['input_ids'] = input_ids
            encoded['labels'] = labels
            encoded['loss_scale'] = loss_scale
            encoded['images'] = [(images_crop, images_ori)]
            encoded['images_seq_mask'] = (torch.tensor(input_ids) == image_token[0])[None]
            encoded['images_spatial_crop'] = images_spatial_crop
        return encoded

    def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]:
        res = super()._data_collator_mm_data(batch)
        images = self.gather_list(batch, 'images')
        if images:
            res['images'] = images
        images_seq_mask = [x['images_seq_mask'] for x in batch if x.get('images_seq_mask') is not None]
        images_spatial_crop = self.concat_tensor(batch, 'images_spatial_crop', 0)
        padding_side = self.padding_side if self.is_training else 'left'
        if images_seq_mask:
            max_len = max([x.shape[1] for x in images_seq_mask])
            res['images_seq_mask'] = torch.concat([
                F.pad(x, (0, max_len - x.shape[1]) if padding_side == 'right' else (max_len - x.shape[1], 0))
                for x in images_seq_mask
            ])
        if images_spatial_crop is not None:
            res['images_spatial_crop'] = images_spatial_crop
        return res


register_template(
    TemplateMeta(
        MLLMTemplateType.deepseek_ocr,
        prefix=['<｜begin▁of▁sentence｜>'],
        prompt=['{{QUERY}}'],
        chat_sep=None,
        template_cls=DeepseekOCR))


class DeepseekOCR2(DeepseekOCR):
    version = 'v2'


register_template(
    TemplateMeta(
        MLLMTemplateType.deepseek_ocr2,
        prefix=['<｜begin▁of▁sentence｜>'],
        prompt=['{{QUERY}}'],
        chat_sep=None,
        template_cls=DeepseekOCR2))


class UnlimitedOCR(DeepseekOCR):
    image_placeholder = ['<image>']  # Remove trailing newline; override the parent class default

    def init_env_args(self):
        super().init_env_args()
        self._device_fixed = False  # Instance variable; avoid sharing state across multiple instances.

    def _fix_device(self):
        if not self._device_fixed and self.model is not None:
            try:
                vision_device = next(self.model.model.vision_model.parameters()).device
                self.model.model.image_newline.data = self.model.model.image_newline.data.to(vision_device)
                self.model.model.view_seperator.data = self.model.model.view_seperator.data.to(vision_device)
                self._device_fixed = True
            except Exception:
                pass

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        self._fix_device()
        return super()._encode(inputs)

    def _load_dynamic_modules(self):
        if self._BasicImageTransform is None:
            model_dir = self.model_info.model_dir
            self._BasicImageTransform = get_class_from_dynamic_module('modeling_unlimitedocr.BasicImageTransform',
                                                                      model_dir)
            self._dynamic_preprocess = get_class_from_dynamic_module('modeling_unlimitedocr.dynamic_preprocess',
                                                                     model_dir)


register_template(
    TemplateMeta(
        MLLMTemplateType.unlimited_ocr,
        prefix=[['bos_token_id']],
        prompt=['{{QUERY}}'],
        chat_sep=None,
        template_cls=UnlimitedOCR,
    ))


@dataclass
class DeepseekV2_5TemplateMeta(TemplateMeta):
    prefix: Prompt = field(default_factory=lambda: ['<｜begin▁of▁sentence｜>{{SYSTEM}}'])
    prompt: Prompt = field(default_factory=lambda: ['<｜User｜>{{QUERY}}<｜Assistant｜>'])
    chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<｜end▁of▁sentence｜>'])
    suffix: Prompt = field(default_factory=lambda: ['<｜end▁of▁sentence｜>'])


register_template(DeepseekV2_5TemplateMeta(LLMTemplateType.deepseek_v2_5))

register_template(DeepseekV2_5TemplateMeta(LLMTemplateType.deepseek_r1, is_thinking=True, thinking_prefix='<think>\n'))


class DeepseekV3_1Template(Template):
    jinja_enable_thinking_key = 'thinking'
    non_thinking_prefix_only_after_user = True


register_template(
    DeepseekV2_5TemplateMeta(
        LLMTemplateType.deepseek_v3_1,
        agent_template='deepseek_v3_1',
        is_thinking=True,
        template_cls=DeepseekV3_1Template,
        thinking_prefix='<think>',
        non_thinking_prefix='</think>',
        history_thinking_prefix='</think>'))

REASONING_EFFORT_MAX = (
    'Reasoning Effort: Absolute maximum with no shortcuts permitted.\n'
    'You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve '
    'the root cause, rigorously stress-testing your logic against all potential paths, edge cases, '
    'and adversarial scenarios.\n'
    'Explicitly write out your entire deliberation process, documenting every intermediate step, '
    'considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n')


class DeepseekV4Template(DeepseekV3_1Template):

    def init_env_args(self):
        super().init_env_args()
        # reasoning_effort: "max", "high", or None
        self.reasoning_effort = get_env_args('reasoning_effort', str, None)
        if self.reasoning_effort is None:
            self.reasoning_effort = 'high' if self.enable_thinking else None
        self.enable_thinking = self.reasoning_effort in ('max', 'high')
        self.chat_template_kwargs['reasoning_effort'] = self.reasoning_effort

    def _get_enable_thinking(self, inputs=None):
        reasoning_effort = None if inputs is None else inputs.chat_template_kwargs.get('reasoning_effort')
        if reasoning_effort is not None:
            return reasoning_effort in ('max', 'high')
        return super()._get_enable_thinking(inputs)

    def _get_system(self, inputs):
        system = super()._get_system(inputs)
        reasoning_effort = inputs.chat_template_kwargs.get('reasoning_effort')
        if reasoning_effort is None:
            reasoning_effort = self.reasoning_effort
        if reasoning_effort == 'max' and self._get_enable_thinking(inputs):
            system = REASONING_EFFORT_MAX + (system or '')
        return system


register_template(
    DeepseekV2_5TemplateMeta(
        LLMTemplateType.deepseek_v4,
        agent_template='deepseek_v4',
        is_thinking=True,
        template_cls=DeepseekV4Template,
        thinking_prefix='<think>',
        non_thinking_prefix='</think>',
        history_thinking_prefix='</think>'))


class DeepseekVL2Template(DeepseekVLTemplate):
    image_placeholder = ['<image>\n']
    placeholder_tokens = ['<image>']

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        from deepseek_vl2.models.processing_deepseek_vl_v2 import VLChatProcessorOutput
        encoded = Template._encode(self, inputs)
        images = inputs.images
        processor = self.processor
        input_ids, labels = encoded['input_ids'], encoded['labels']
        images_seq_mask = [False] * len(input_ids)
        idx_list = findall(input_ids, processor.image_token_id)  # '<image>'
        _, images_list, _, images_spatial_crop, num_image_tokens = processor.tokenize_with_images(
            '<image>' * len(images), images, cropping=len(images) <= 2)
        new_num_tokens = 0
        for idx, n_image_tokens in zip(idx_list, num_image_tokens):
            image_tokens = [processor.image_token_id] * n_image_tokens
            input_ids = input_ids[:idx] + image_tokens + input_ids[idx + 1:]
            if labels is not None:
                labels = labels[:idx] + [-100] * n_image_tokens + labels[idx + 1:]
            images_seq_mask = images_seq_mask[:idx] + [True] * n_image_tokens + images_seq_mask[idx + 1:]
            new_num_tokens += n_image_tokens - 1

        output = VLChatProcessorOutput(
            sft_format=None,
            input_ids=torch.tensor(input_ids),
            target_ids=torch.tensor(input_ids),
            images=torch.stack(images_list) if images_list else torch.zeros((0, 3, 384, 384)),
            images_seq_mask=torch.tensor(images_seq_mask),
            images_spatial_crop=torch.tensor(images_spatial_crop),
            num_image_tokens=num_image_tokens)
        output.images = output.images.to(dtype=self.model_info.torch_dtype)
        encoded = {'output': output, 'input_ids': input_ids, 'labels': labels}
        return encoded

    def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]:
        inputs['images_seq_mask'] = inputs['images_seq_mask'].to(torch.bool)
        inputs['images_spatial_crop'] = inputs['images_spatial_crop'].to(torch.long)
        inputs_embeds = model.prepare_inputs_embeds(**inputs)
        return {'inputs_embeds': inputs_embeds}


register_template(
    DeepseekV2_5TemplateMeta(
        MLLMTemplateType.deepseek_vl2,
        prompt=['<|User|>: {{QUERY}}\n\n<|Assistant|>:'],
        template_cls=DeepseekVL2Template,
    ))

register_template(
    DeepseekVLTemplateMeta(
        MLLMTemplateType.deepseek_janus_pro,
        prompt=['<|User|>: {{QUERY}}\n\n<|Assistant|>:'],
        template_cls=DeepseekJanus))