V2PE-256K/modeling_internvl_chat.py

# --------------------------------------------------------
# InternVL
# Copyright (c) 2024 OpenGVLab
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
import warnings
from typing import Any, List, Optional, Tuple, Union

import torch.distributed as dist
import torch.utils.checkpoint
import transformers
from internvl.conversation import get_conv_template
from internvl.model.internlm2.modeling_internlm2 import InternLM2ForCausalLM
from internvl.model.phi3.modeling_phi3 import Phi3ForCausalLM
from peft import LoraConfig, get_peft_model
from torch import nn
from torch.nn import CrossEntropyLoss
from transformers import (AutoModel, GenerationConfig, LlamaForCausalLM,
                          LlamaTokenizer, Qwen2ForCausalLM)
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import ModelOutput, logging

from .configuration_internvl_chat import InternVLChatConfig
from .modeling_intern_vit import InternVisionModel

logger = logging.get_logger(__name__)
from transformers import AutoTokenizer
import json
tokenizer_path="/mnt/petrelfs/share_data/chenziyi/InternVL2-2B"
global_tokenizer = AutoTokenizer.from_pretrained(
        tokenizer_path, add_eos_token=False, trust_remote_code=True, use_fast=False)
import random


def version_cmp(v1, v2, op='eq'):
    import operator

    from packaging import version
    op_func = getattr(operator, op)
    return op_func(version.parse(v1), version.parse(v2))
def extract_local(value, rank, world_size, dim=1):
    value_chunks = value.chunk(2 * world_size, dim=dim)
    local_value = torch.cat(
        [value_chunks[rank], value_chunks[2 * world_size - rank - 1]], dim=dim
    )
    return local_value.to(value.device)
def extract_local2(value, rank, world_size,  dim=1):
    dimension_size = value.shape[dim]
    sub_seq_length = dimension_size // world_size

    sub_seq_start = rank * sub_seq_length
    sub_seq_end = (rank + 1) * sub_seq_length
    local_value = value[:, sub_seq_start:sub_seq_end]

    return local_value.to(value.device)
class GatherLayer(torch.autograd.Function):
    """Gather tensors from all process, supporting backward propagation."""

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        output = [torch.zeros_like(input) for _ in range(dist.get_world_size(local_group))]
        dist.all_gather(output, input, group=local_group)
        return torch.stack(output, 0)

    @staticmethod
    def backward(ctx, grads):
        (input,) = ctx.saved_tensors
        dist.all_reduce(grads, group=local_group)
        grad_out = torch.zeros_like(input)
        grad_out[:] = grads[dist.get_rank(local_group)]
        return grad_out
class InternVLChatModel(PreTrainedModel):
    config_class = InternVLChatConfig
    main_input_name = 'pixel_values'
    _no_split_modules = ['InternVisionModel', 'LlamaDecoderLayer', 'InternLM2DecoderLayer',
                         'Phi3DecoderLayer', 'Qwen2DecoderLayer']

    def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None):
        super().__init__(config)

        assert version_cmp(transformers.__version__, '4.37.0', 'ge')
        image_size = config.force_image_size or config.vision_config.image_size
        patch_size = config.vision_config.patch_size
        self.patch_size = patch_size
        self.select_layer = config.select_layer
        self.template = config.template
        
        # batch_size: 批处理大小
        # patch_size: 图片分块大小
        # downsample_ratio: 缩放比例，将高分辨率图像转换为低分辨率图像
        # self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2))
        self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2))
        self.downsample_ratio = config.downsample_ratio
        self.ps_version = config.ps_version
        self.compress_seq = config.compress_seq
        self.attn_type = config.attn_type
        self.posid_type = config.posid_type
        if self.posid_type is None:
            self.posid_type='default'
        assert self.posid_type in ['default','None', 'qkvLearnable', 'qkLearnable', '1dROPE', '2dROPE']
        self.group_list = config.group_list
        self.chunk_num = config.chunk_num
        self.interaction = config.interaction


        logger.info(f'num_image_token: {self.num_image_token}')
        logger.info(f'ps_version: {self.ps_version}')
        config.llm_config.posid_type = self.posid_type
        config.llm_config.rope_pos_id_version=config.rope_pos_id_version
        if vision_model is not None:
            self.vision_model = vision_model
        else:
            self.vision_model = InternVisionModel(config.vision_config)
        if language_model is not None:
            self.language_model = language_model
        else:
            if config.llm_config.architectures[0] == 'LlamaForCausalLM':
                self.language_model = LlamaForCausalLM(config.llm_config)
            elif config.llm_config.architectures[0] == 'InternLM2ForCausalLM':
                self.language_model = InternLM2ForCausalLM(config.llm_config)
            elif config.llm_config.architectures[0] == 'Phi3ForCausalLM':
                self.language_model = Phi3ForCausalLM(config.llm_config)
            elif config.llm_config.architectures[0] == 'Qwen2ForCausalLM':
                self.language_model = Qwen2ForCausalLM(config.llm_config)
            else:
                raise NotImplementedError(f'{config.llm_config.architectures[0]} is not implemented.')

        vit_hidden_size = config.vision_config.hidden_size
        llm_hidden_size = config.llm_config.hidden_size

        self.mlp1 = nn.Sequential(
            nn.LayerNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2),
            nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size),
            nn.GELU(),
            nn.Linear(llm_hidden_size, llm_hidden_size)
        )
        
        if self.posid_type in ['qkvLearnable']:
            self.local_posid = nn.Embedding(self.num_image_token,llm_hidden_size)
            
        self.img_context_token_id = None
        self.conv_template = get_conv_template(self.template)
        self.system_message = self.conv_template.system_message
        self.num_samples = 0

        if config.use_backbone_lora:
            self.wrap_backbone_lora(r=config.use_backbone_lora, lora_alpha=2 * config.use_backbone_lora)

        if config.use_llm_lora:
            self.wrap_llm_lora(r=config.use_llm_lora, lora_alpha=2 * config.use_llm_lora)
    def init_embed(self):
        if hasattr(self,'local_posid'):
            nn.init.normal_(self.local_posid.weight, mean=0.0, std=0.02)
    def wrap_backbone_lora(self, r=128, lora_alpha=256, lora_dropout=0.05):
        lora_config = LoraConfig(
            r=r,
            target_modules=['attn.qkv', 'attn.proj', 'mlp.fc1', 'mlp.fc2'],
            lora_alpha=lora_alpha,
            lora_dropout=lora_dropout,
        )
        self.vision_model = get_peft_model(self.vision_model, lora_config)
        self.vision_model.print_trainable_parameters()

    def wrap_llm_lora(self, r=128, lora_alpha=256, lora_dropout=0.05):
        lora_config = LoraConfig(
            r=r,
            target_modules=['self_attn.q_proj', 'self_attn.k_proj', 'self_attn.v_proj', 'self_attn.o_proj',
                            'mlp.gate_proj', 'mlp.down_proj', 'mlp.up_proj'],
            lora_alpha=lora_alpha,
            lora_dropout=lora_dropout,
            task_type='CAUSAL_LM'
        )
        self.language_model = get_peft_model(self.language_model, lora_config)
        self.language_model.enable_input_require_grads()
        self.language_model.print_trainable_parameters()
    
    def forward(
            self,
            pixel_values: torch.FloatTensor,
            input_ids: torch.LongTensor = None,
            attention_mask: Optional[torch.Tensor] = None,
            position_ids: Optional[torch.Tensor] = None,
            image_flags: Optional[torch.LongTensor] = None,
            past_key_values: Optional[List[torch.FloatTensor]] = None,
            labels: Optional[torch.LongTensor] = None,
            use_cache: Optional[bool] = None,
            output_attentions: Optional[bool] = None,
            output_hidden_states: Optional[bool] = None,
            return_dict: Optional[bool] = None,
            statistics: Optional[torch.LongTensor] = None,
            loss_weight: Optional[List] = None,
            loss_reduction_all_gather: Optional[bool] = False,
            origin_cu_seq_lens: Optional[torch.Tensor] = None,
            rope_pos_id: Optional[torch.Tensor] = None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        # import ipdb
        # ipdb.set_trace()
        if isinstance(position_ids,list):
            position_ids=torch.tensor(position_ids).to(input_ids.device)
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        # print("Printing decoded input ids")
        # decoded_texts = [global_tokenizer.decode(ids, skip_special_tokens=True) for ids in input_ids]
        # for i, text in enumerate(decoded_texts):
        #     print(f"Sample {i+1}: {text}")
        global local_group
        if self.group_list is not None:
            for group_idx,group in enumerate(self.group_list):
                if type(group)==torch.distributed.distributed_c10d.ProcessGroup:
                    # assert type(group)==torch.distributed.distributed_c10d.ProcessGroup
                    break        # print("Printing decoded input ids")
            local_group=group
        else:
            group=None
            local_group=None
        image_flags = image_flags.squeeze(-1)
        input_embeds = self.language_model.get_input_embeddings()(input_ids).clone()
        if self.attn_type:
            if self.attn_type=='ring':
                group_size = dist.get_world_size(group)
                img_num_dim = 0
                pad_num=0
                if pixel_values.shape[img_num_dim] > group_size:
                    if pixel_values.shape[img_num_dim] % group_size!=0:
                        pad_num = group_size - pixel_values.shape[img_num_dim] % group_size
                        if pad_num < group_size:  # 仅在需要填充时进行
                            # 创建填充的张量，与 pixel_values 的形状匹配
                            pad_shape = list(pixel_values.shape)
                            pad_shape[img_num_dim] = pad_num  # 在目标维度上设置填充值
                            pad_pixel = torch.zeros(pad_shape, dtype=pixel_values.dtype, device=pixel_values.device)

                            # 在指定维度上拼接原始张量和填充张量
                            pixel_values = torch.cat([pixel_values, pad_pixel], dim=img_num_dim)

                    chunked_pixel=torch.chunk(pixel_values, group_size, dim=img_num_dim)
                    local_pixel=chunked_pixel[dist.get_rank(group)]
                    local_vit_embeds=self.extract_feature(local_pixel)
                    vit_embeds=GatherLayer.apply(local_vit_embeds)
                    vit_embeds=vit_embeds.view(-1,vit_embeds.shape[-2],vit_embeds.shape[-1])
                    if pad_num>0:
                        vit_embeds=vit_embeds[:-pad_num]
                else:
                    vit_embeds = self.extract_feature(pixel_values)
            else:
                vit_embeds = self.extract_feature(pixel_values)
        else:
            vit_embeds = self.extract_feature(pixel_values)
        
        if self.posid_type=='qkvLearnable':
            # added_embeds = self.local_posid(torch.arange(self.num_image_token).to(pixel_values.device))
            # vit_embeds = vit_embeds + added_embeds
            vit_embeds=vit_embeds+self.local_posid(torch.arange(self.num_image_token).to(pixel_values.device))


        vit_embeds = vit_embeds[image_flags == 1]
        vit_batch_size = pixel_values.shape[0]
        # print("Printing pixiel shape", pixel_values.shape)
        B, N, C = input_embeds.shape
        input_embeds = input_embeds.reshape(B * N, C)

        if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:
            print(f'dynamic ViT batch size: {vit_batch_size}, images per sample: {vit_batch_size / B}, dynamic token length: {N}')
            if statistics is not None:
                num_samples, num_padding_tokens, num_padding_images = statistics.tolist()
                self.num_samples += num_samples
                print(f'total_samples={self.num_samples}, {num_samples=}, {num_padding_tokens=}, {num_padding_images=}')
        input_ids = input_ids.reshape(B * N)
        selected = (input_ids == self.img_context_token_id)
        try:
            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)
            ignore_flag = False
        except Exception as e:
            vit_embeds = vit_embeds.reshape(-1, C)
            print(f'warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, '
                  f'vit_embeds.shape={vit_embeds.shape}')
            n_token = selected.sum()
            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds[:n_token]
            # ignore_flag = True
            ignore_flag = False

        input_embeds = input_embeds.reshape(B, N, C)
        if self.attn_type:
            if self.attn_type=='ulysses':
                input_embeds=extract_local2(input_embeds,dist.get_rank(group),dist.get_world_size(group))
                position_ids=extract_local2(position_ids,dist.get_rank(group),dist.get_world_size(group))
                labels=extract_local2(labels,dist.get_rank(group),dist.get_world_size(group))
                loss_weight=extract_local2(torch.tensor(loss_weight),dist.get_rank(group),dist.get_world_size(group))
                loss_weight=list(loss_weight.numpy())
                attention_mask=attention_mask//dist.get_world_size(group)
            elif self.attn_type=='ring':
                input_embeds=extract_local(input_embeds,dist.get_rank(group),dist.get_world_size(group))
                position_ids=extract_local(position_ids,dist.get_rank(group),dist.get_world_size(group))
                labels=extract_local(labels,dist.get_rank(group),dist.get_world_size(group))
                if loss_weight:
                    loss_weight=extract_local(torch.tensor(loss_weight),dist.get_rank(group),dist.get_world_size(group))
                    loss_weight=list(loss_weight.numpy())
                attention_mask=attention_mask//dist.get_world_size(group)
        outputs = self.language_model(
            inputs_embeds=input_embeds,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            compress_seq=self.compress_seq,
            group_list=self.group_list,
            chunk_num=self.chunk_num,
            origin_cu_seq_lens=origin_cu_seq_lens,
            interaction=self.interaction,
            selected=selected
        )
        logits = outputs.logits
        
        loss = None
        if labels is not None and loss_weight is not None:
            # decoded_labels = global_tokenizer.decode(labels[0][labels[0]!=-100], skip_special_tokens=True)
            loss_weight = torch.tensor(loss_weight, dtype=torch.float32, device=labels.device)
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            shift_weights = loss_weight[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = CrossEntropyLoss(reduction='none')
            shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            shift_weights = shift_weights.view(-1)
            # Enable model parallelism
            shift_labels = shift_labels.to(shift_logits.device)
            shift_weights = shift_weights.to(shift_logits.device)
            loss = loss_fct(shift_logits, shift_labels)

            shift_weights_sum = shift_weights.sum()
                
            if loss_reduction_all_gather:
                dist.all_reduce(shift_weights_sum, op=dist.ReduceOp.AVG)

            loss = loss * shift_weights
            loss = loss.sum() / shift_weights_sum
            if ignore_flag:
                loss = loss * 0.0
        elif labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = CrossEntropyLoss()
            shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            # Enable model parallelism
            shift_labels = shift_labels.to(shift_logits.device)
            loss = loss_fct(shift_logits, shift_labels)
            if ignore_flag:
                loss = loss * 0.0
        params=dict(self.named_parameters())
        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        # self.update_log(log_dict)
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    def pixel_shuffle(self, x, scale_factor=0.5):
        n, w, h, c = x.size()
        # N, W, H, C --> N, W, H * scale, C // scale
        x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))
        # N, W, H * scale, C // scale --> N, H * scale, W, C // scale
        x = x.permute(0, 2, 1, 3).contiguous()
        # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)
        x = x.view(n, int(h * scale_factor), int(w * scale_factor),
                   int(c / (scale_factor * scale_factor)))
        if self.ps_version == 'v1':
            warnings.warn("In ps_version 'v1', the height and width have not been swapped back, "
                          'which results in a transposed image.')
        else:
            x = x.permute(0, 2, 1, 3).contiguous()
        return x

    def extract_feature(self, pixel_values):
        # 选择视觉模型特定层的输出作为图片特征
        if self.select_layer == -1:
            vit_embeds = self.vision_model(
                pixel_values=pixel_values,
                output_hidden_states=False,
                return_dict=True).last_hidden_state
        else:
            vit_embeds = self.vision_model(
                pixel_values=pixel_values,
                output_hidden_states=True,
                return_dict=True).hidden_states[self.select_layer]
        # [batch_size, num_patches, vit_hidden_size]
        # 去除第一个标记
        vit_embeds = vit_embeds[:, 1:, :]

        # [batch_size, num_patches, vit_hidden_size] -> [batch_size, h, w, vit_hidden_size]
        h = w = int(vit_embeds.shape[1] ** 0.5)
        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
        # 像素混洗，降低分辨率，减少 num_patches
        vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)
        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
        # 线性层，vit_hidden_size -> llm_hidden_size
        vit_embeds = self.mlp1(vit_embeds)
        return vit_embeds

    def batch_chat(self, tokenizer, pixel_values, questions, generation_config, num_patches_list=None,
                   history=None, return_history=False, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>',
                   IMG_CONTEXT_TOKEN='<IMG_CONTEXT>', verbose=False, image_counts=None):
        if history is not None or return_history:
            print('Now multi-turn chat is not supported in batch_chat.')
            raise NotImplementedError

        if image_counts is not None:
            num_patches_list = image_counts
            print('Warning: `image_counts` is deprecated. Please use `num_patches_list` instead.')
        
        img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
        self.img_context_token_id = img_context_token_id

        if verbose and pixel_values is not None:
            image_bs = pixel_values.shape[0]
            print(f'dynamic ViT batch size: {image_bs}')

        queries = []
        for idx, num_patches in enumerate(num_patches_list):
            question = questions[idx]
            if pixel_values is not None and '<image>' not in question:
                question = '<image>\n' + question
            template = get_conv_template(self.template)
            template.append_message(template.roles[0], question)
            template.append_message(template.roles[1], None)
            query = template.get_prompt()

            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
            query = query.replace('<image>', image_tokens, 1)
            queries.append(query)

        # tokenizer.padding_side = 'left'
        model_inputs = tokenizer(queries, return_tensors='pt', padding=False)
        input_ids = model_inputs['input_ids'].cuda()
        attention_mask = model_inputs['attention_mask'].cuda()
        eos_token_id = tokenizer.convert_tokens_to_ids(template.sep)
        generation_config['eos_token_id'] = eos_token_id
        generation_output = self.generate(
            pixel_values=pixel_values,
            input_ids=input_ids,
            attention_mask=attention_mask,
            **generation_config
        )
        responses = tokenizer.batch_decode(generation_output, skip_special_tokens=True)
        responses = [response.split(template.sep)[0].strip() for response in responses]
        return responses

    def chat(self, tokenizer, pixel_values, question, generation_config, history=None, return_history=False,
             num_patches_list=None, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>',
             verbose=False,**kwargs):
        if history is None and pixel_values is not None and '<image>' not in question:
            question = '<image>\n' + question

        # num_patches_list 用法：
        if num_patches_list is None:
            num_patches_list = [pixel_values.shape[0]] if pixel_values is not None else []
        assert pixel_values is None or len(pixel_values) == sum(num_patches_list)

        # 设置图片上下文的 token id
        img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
        self.img_context_token_id = img_context_token_id

        # 获取 Chat 模板
        template = get_conv_template(self.template)
        # 设置系统消息
        template.system_message = self.system_message
        # 设置分隔符 End Of Sentence
        eos_token_id = tokenizer.convert_tokens_to_ids(template.sep)

        # 将历史对话添加到模板中
        history = [] if history is None else history
        for (old_question, old_answer) in history:
            template.append_message(template.roles[0], old_question)
            template.append_message(template.roles[1], old_answer)
        template.append_message(template.roles[0], question)
        template.append_message(template.roles[1], None)
        # 生成查询
        query = template.get_prompt()

        # verbose: 是否打印调试信息
        if verbose and pixel_values is not None:
            # pixel_values 形状: [batch_size, channels, height, width]
            # 其中 batch_size 即图片数量
            # 打印批处理大小信息
            image_bs = pixel_values.shape[0]
            print(f'dynamic ViT batch size: {image_bs}')

        # 将图片 token 插入到查询中，图片用占位符 IMG_CONTEXT_TOKEN 代替
        for num_patches in num_patches_list:
            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
            query = query.replace('<image>', image_tokens, 1)

        # 用分词器将查询转换为模型输入
        model_inputs = tokenizer(query, return_tensors='pt')
        # 文本对应的 token id，转换为 cuda 张量
        # ID 长度就是 Token 长度，形状为 [1, sequence_length]
        input_ids = model_inputs['input_ids'].cuda()
        # print(f'Token length: {input_ids.shape[1]}')
        # 实际输入掩码为 1，填充部分掩码为 0
        attention_mask = model_inputs['attention_mask'].cuda()
        # 分隔符 End Of Sentence
        generation_config['eos_token_id'] = eos_token_id
        if 'rope_pos_id_version' in kwargs:
            self.language_model.rope_pos_id_version=kwargs['rope_pos_id_version']
            pos_ids=[]
            ret={'input_ids':input_ids,'attention_mask':attention_mask}
            for i in range(input_ids.shape[0]):
                # cur_position_ids = ret['attention_mask'][i].long().cumsum(-1) - 1
                # cur_position_ids.masked_fill_(ret['attention_mask'][i] == 0, 1)

                if kwargs['rope_pos_id_version'] == 'default':
                    cur_dtype = torch.long
                    # bf16 -> long 会产生截断
                else:
                    cur_dtype = torch.float32

                if 'rope_pos_id_stride' in kwargs:
                    rope_pos_id_stride = kwargs['rope_pos_id_stride']
                else:
                    rope_pos_id_stride = None

                pos_ids.append(torch.tensor(get_rope_pos_id(ret, num_tiles=kwargs['num_tiles'][i], dtype=cur_dtype,
                                               rope_pos_id_version=kwargs['rope_pos_id_version'],
                                               position_id=torch.arange(0,input_ids.shape[1]),
                                               # position_id=cur_position_ids,
                                               boxes=kwargs['all_boxes'][i],
                                               orig_size=None,
                                               images=kwargs['image_list'][i],
                                               IMG_START_TOKEN=IMG_START_TOKEN,
                                               IMG_END_TOKEN=IMG_END_TOKEN, rope_pos_id_stride=rope_pos_id_stride)).cuda())
                
            pos_ids=torch.stack(pos_ids)
            if self.attn_type=='ulysses' or self.attn_type=='ring':
                if input_ids.shape[1]%(2*dist.get_world_size())!=0:
                    num_padding = 2*dist.get_world_size()-input_ids.shape[1]%(2*dist.get_world_size())
                    # 创建需要的 padding，input_ids 和 labels 填充值为 -100
                    padding_shape = (input_ids.shape[0], num_padding)
                    input_padding = torch.full(padding_shape, 1, dtype=input_ids.dtype, device=input_ids.device)
                    attn_mask_padding = torch.full(padding_shape, 1, dtype=attention_mask.dtype, device=attention_mask.device)
                    # 对 input_ids 和 labels 进行 padding
                    input_ids = torch.cat([input_ids, input_padding], dim=1)
                    attention_mask=torch.cat([attention_mask,attn_mask_padding],dim=1)
                                # position_ids 添加正确的递增填充
                    max_pos_id = pos_ids.max() + 1  # 找到当前最大 position_id
                    pos_padding = torch.arange(max_pos_id, max_pos_id + num_padding, device=input_ids.device)
                    pos_padding = pos_padding.unsqueeze(0).expand(input_ids.shape[0], -1)
                    pos_ids = torch.cat([pos_ids, pos_padding], dim=1)
            generation_output = self.generate(
                pixel_values=pixel_values,
                input_ids=input_ids,
                attention_mask=attention_mask,
                position_ids=pos_ids,
                **generation_config,
            )
        else:
            self.language_model.rope_pos_id_version='default'
            if self.attn_type=='ulysses' or self.attn_type=='ring':
                if input_ids.shape[1]%(2*dist.get_world_size())!=0:
                    num_padding = 2*dist.get_world_size()-input_ids.shape[1]%(2*dist.get_world_size())
                    # 创建需要的 padding，input_ids 和 labels 填充值为 -100
                    padding_shape = (input_ids.shape[0], num_padding)
                    input_padding = torch.full(padding_shape, 1, dtype=input_ids.dtype, device=input_ids.device)
                    attn_mask_padding = torch.full(padding_shape, 0, dtype=attention_mask.dtype, device=attention_mask.device)
                    # 对 input_ids 和 labels 进行 padding
                    input_ids = torch.cat([input_ids, input_padding], dim=1)
                    attention_mask=torch.cat([attention_mask,attn_mask_padding],dim=1)
            generation_output = self.generate(
                pixel_values=pixel_values,
                input_ids=input_ids,
                attention_mask=attention_mask,
                **generation_config,
            )
        # 解码生成的输出，跳过特殊 token
        response = tokenizer.batch_decode(generation_output, skip_special_tokens=True)[0]
        # 根据分隔符分段
        response = response.split(template.sep)[0].strip()
        # 将结果写入历史
        history.append((question, response))
        if return_history:
            return response, history
        else:
            query_to_print = query.replace(IMG_CONTEXT_TOKEN, '')
            query_to_print = query_to_print.replace(f'{IMG_START_TOKEN}{IMG_END_TOKEN}', '<image>')
            if verbose:
                print(query_to_print, response)
            return response

    @torch.no_grad()
    def generate(
            self,
            pixel_values: Optional[torch.FloatTensor] = None,
            input_ids: Optional[torch.FloatTensor] = None,
            attention_mask: Optional[torch.LongTensor] = None,
            visual_features: Optional[torch.FloatTensor] = None,
            generation_config: Optional[GenerationConfig] = None,
            output_hidden_states: Optional[bool] = None,
            return_dict: Optional[bool] = None,
            **generate_kwargs,
    ) -> torch.LongTensor:
        assert self.img_context_token_id is not None
        if pixel_values is not None:
            # 提取图片 embedding
            # [batch_size, channels, height, width] -> [batch_size, 每张图片的 patch 数, embedding_dim]
            if visual_features is not None:
                vit_embeds = visual_features
            else:
                vit_embeds = self.extract_feature(pixel_values)
            if self.posid_type=='qkvLearnable':
                added_embeds = self.local_posid(torch.arange(self.num_image_token).to(pixel_values.device))
                vit_embeds = vit_embeds + added_embeds
                # vit_embeds=vit_embeds+self.local_posid(torch.arange(self.num_image_token).to(pixel_values.device))
            # 通过嵌入层将 token id 转化为嵌入向量
            # 其中图片用占位符 IMG_CONTEXT_TOKEN 的 embedding 代替
            input_embeds = self.language_model.get_input_embeddings()(input_ids)
            # [1, sequence_length, embedding_dim] -> [sequence_length, embedding_dim]
            B, N, C = input_embeds.shape
            input_embeds = input_embeds.reshape(B * N, C)

            # [1, sequence_length] -> [sequence_length]
            input_ids = input_ids.reshape(B * N)

            selected = (input_ids == self.img_context_token_id)
            assert selected.sum() != 0
            
            # 图片 embedding: [总 Patch 数, embedding_dim]
            # 每个 patch 与一个占位符对应，对应一列 embedding
            input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)

            input_embeds = input_embeds.reshape(B, N, C)
        else:
            # 通过嵌入层将 token id 转化为嵌入向量
            # 例如 one hot 编码、Word2Vec、GloVe、FastText等
            # 嵌入层是一张查找表
            # [1, sequence_length] -> [1, sequence_length, embedding_dim]
            input_embeds = self.language_model.get_input_embeddings()(input_ids)
                    # 找到图片占位符的位置
        if 'position_ids' in generate_kwargs:
            pos_id=generate_kwargs['position_ids']
            if self.attn_type:
                if self.attn_type=='ulysses':
                    input_embeds=extract_local2(input_embeds,dist.get_rank(),dist.get_world_size())
                    attention_mask=extract_local2(attention_mask,dist.get_rank(),dist.get_world_size())
                    pos_id=extract_local2(pos_id,dist.get_rank(),dist.get_world_size())
                elif self.attn_type=='ring':
                    former_shape = input_embeds.shape
                    input_embeds=extract_local(input_embeds,dist.get_rank(),dist.get_world_size())
                    attention_mask=extract_local(attention_mask,dist.get_rank(),dist.get_world_size())
                    pos_id=extract_local(pos_id,dist.get_rank(),dist.get_world_size())
                generate_kwargs['position_ids']=pos_id
                
        else:
            if self.attn_type:
                if self.attn_type=='ulysses':
                    input_embeds=extract_local2(input_embeds,dist.get_rank(),dist.get_world_size())
                    attention_mask=extract_local2(attention_mask,dist.get_rank(),dist.get_world_size())
                elif self.attn_type=='ring':
                    former_shape = input_embeds.shape
                    input_embeds=extract_local(input_embeds,dist.get_rank(),dist.get_world_size())
                    attention_mask=extract_local(attention_mask,dist.get_rank(),dist.get_world_size())

        outputs = self.language_model.generate(
            inputs_embeds=input_embeds,
            attention_mask=attention_mask,
            generation_config=generation_config,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            use_cache=True,
            **generate_kwargs,
        )

        return outputs
    def update_log(self, new_log_dict):
        if not hasattr(self, 'log_dict'):
            self.log_dict = {}
        for key, value in new_log_dict.items():
            if 'loss' in key:
                if key not in self.log_dict:
                    self.log_dict[key] = value
                else:
                    self.log_dict[key] += value
            else:
                # just copy it
                self.log_dict[key] = value

def get_rope_pos_id(ret, num_tiles, dtype, rope_pos_id_version='default', position_id=None,boxes=None, orig_size=None,images=None,IMG_START_TOKEN='<img>',IMG_END_TOKEN='</img>',rope_pos_id_stride=None):
        image_start_token_id = global_tokenizer.convert_tokens_to_ids(IMG_START_TOKEN)
        image_end_token_id = global_tokenizer.convert_tokens_to_ids(IMG_END_TOKEN)
        num_image_token=256
        rope_pos_id_list = []

        input_ids_0 = ret['input_ids'][0]
        attention_mask_0 = ret['attention_mask'][0]
        image_start_token_id_idxs = torch.where(input_ids_0 == image_start_token_id)[0]
        image_end_token_id_idxs = torch.where(input_ids_0 == image_end_token_id)[0]

        last_record_pos_id = -1
        start_index = 0
        for i in range(len(image_start_token_id_idxs)):
            # 根据序列中的 IMG_START_TOKEN 出现的位置，锁定需要处理的图像 id 序列
            # 注：这里的 IMG_START_TOKEN 和 IMG_END_TOKEN 应当与文本的处理方式相同
            box = boxes[i]
            image = images[i]
            
            rope_pos_id_pre = attention_mask_0[start_index:image_start_token_id_idxs[i] + 1].long().cumsum(-1) - 1 + (last_record_pos_id + 1) # 从处理好的序列的最后一个 global id 开始 count
            rope_pos_id_pre.masked_fill_(attention_mask_0[start_index:image_start_token_id_idxs[i] + 1] == 0, 1)
            rope_pos_id_list.append(rope_pos_id_pre)

            last_record_pos_id = rope_pos_id_pre[-1].long()

            num_tile = num_tiles[i]
            num_sub_imgs = num_tile - 1
            is_last = (i == len(image_start_token_id_idxs) - 1)

            if rope_pos_id_version == 'v0':
                # 子图为小数，且不管多少个子图，其分配的总 global id 跨度为1；缩略图单独分配完整的，跨度为 1的 global id. Example:
                # start_id = 100; 100 - 101 (分给 4 * 256)，子图数目为4; 101 - 102 (分给 256) 缩略图
                if num_sub_imgs > 0:
                    split_img_id_idxs = torch.linspace(last_record_pos_id, last_record_pos_id + 1, (num_tile - 1) * num_image_token + 1)[1:].to(dtype=dtype)  # 小数数值的 tensor 作为变换的数据取值
                    origin_split_img_id_idxs = split_img_id_idxs
                    ############################## 进行位置变换 ##############################
                    # 先计算第一个子图对应 index
                    rearange_idx_list = []
                    rearange_idx_list_list = []
                    base_index_list = []
                    num_img_token_in_length = int(num_image_token ** 0.5)
                    num_patch_width = int(box[-1][2] // box[0][2])
                    num_patch_height = int(box[-1][3] // box[0][2])
                    assert num_patch_width * num_patch_height == len(box)

                    num_total_patch_width_token = num_patch_width * num_img_token_in_length
                    num_total_patch_height_token = num_patch_height * num_img_token_in_length
                    assert num_total_patch_width_token * num_total_patch_height_token == num_sub_imgs * num_image_token, (num_total_patch_width_token * num_total_patch_height_token, num_sub_imgs * num_image_token)

                    for k in range(num_image_token):
                        map_idx = (k // num_img_token_in_length) * num_total_patch_width_token + (k % num_img_token_in_length)
                        base_index_list.append(map_idx)

                    # 计算其他子图对应第一个子图的 offset

                    for k in range(num_sub_imgs):
                        patch_row = k // num_patch_width
                        patch_col = k % num_patch_width
                        offset = patch_row * (num_image_token * num_patch_width) + patch_col * num_img_token_in_length
                        # print(f'{k=}, {offset=}')
                        dst_index_list = [base_index + offset for base_index in base_index_list]
                        rearange_idx_list.extend(dst_index_list)
                        rearange_idx_list_list.append(dst_index_list)

                    ############################## plot 验证 ##############################
                    # img_boxes = [(deepcopy(img), cur_box, cur_posid) for img, cur_box, cur_posid in
                    #              zip(image[:-1], box, rearange_idx_list_list)]
                    # self.eval_posid_by_plot(img_boxes, rope_pos_id_version, None)
                    # img_boxes = [(deepcopy(img), cur_box, cur_posid) for img, cur_box, cur_posid in
                    #              zip(image[:-1], box, rearange_idx_list_list)]
                    # self.eval_posid_by_plot(img_boxes, rope_pos_id_version, split_img_id_idxs)

                    ############################## rearrange ##############################
                    split_img_id_idxs = split_img_id_idxs[rearange_idx_list]

                    rope_pos_id_list.append(split_img_id_idxs)
                    thumbnail_id_idxs = origin_split_img_id_idxs.reshape([num_image_token, -1]).to(dtype=dtype).mean(dim=1).view(-1)
                    rope_pos_id_list.append(thumbnail_id_idxs)
                    last_record_pos_id = origin_split_img_id_idxs[-1].long()
                else:
                    thumbnail_id_idxs = torch.linspace(last_record_pos_id, last_record_pos_id + 1,
                                                    num_image_token + 1)[1:].to(dtype=dtype)  # 缩略图
                    rope_pos_id_list.append(thumbnail_id_idxs)
                    last_record_pos_id = (last_record_pos_id + 1).long()

                # 验证是否能够恢复为等差数列
                if num_tile > 1:
                    gt_pos_id = torch.linspace(last_record_pos_id - 2, last_record_pos_id - 1, (num_tile - 1) * num_image_token + 1)[1:].to(dtype=dtype)
                    # self.eval_posid_by_rearange(box, rope_pos_id_list, gt_pos_id, num_tile, dtype, is_last)

            elif rope_pos_id_version == 'v1':
                # 子图为小数，若有 N 个子图，其分配的总 global id 跨度为 N；缩略图单独分配完整的，跨度为 1的 global id. Example:
                # start_id = 100; 100 - 104 (分给 4 * 256)，子图数目为4; 104 - 105 (分给 256) 缩略图
                if num_sub_imgs > 0:
                    split_img_id_idxs = torch.linspace(last_record_pos_id, last_record_pos_id + num_tile - 1, (num_tile - 1) * num_image_token + 1)[1:].to(dtype=dtype)  # 小数数值的 tensor 作为变换的数据取值
                    origin_split_img_id_idxs = split_img_id_idxs
                    ############################## 进行位置变换 ##############################
                    # 先计算第一个子图对应 index
                    rearange_idx_list = []
                    rearange_idx_list_list = []
                    base_index_list = []
                    # rearange_split_img_id_idxs_list = []
                    num_img_token_in_length = int(num_image_token ** 0.5)
                    num_patch_width = int(box[-1][2] // box[0][2])
                    num_patch_height = int(box[-1][3] // box[0][2])
                    assert num_patch_width * num_patch_height == len(box)

                    num_total_patch_width_token = num_patch_width * num_img_token_in_length
                    num_total_patch_height_token = num_patch_height * num_img_token_in_length
                    assert num_total_patch_width_token * num_total_patch_height_token == num_sub_imgs * num_image_token, (
                    num_total_patch_width_token * num_total_patch_height_token, num_sub_imgs * num_image_token)

                    for k in range(num_image_token):
                        map_idx = (k // num_img_token_in_length) * num_total_patch_width_token + (
                                    k % num_img_token_in_length)
                        base_index_list.append(map_idx)

                    # 计算其他子图对应第一个子图的 offset

                    for k in range(num_sub_imgs):
                        patch_row = k // num_patch_width
                        patch_col = k % num_patch_width
                        offset = patch_row * (
                                    num_image_token * num_patch_width) + patch_col * num_img_token_in_length
                        # print(f'{k=}, {offset=}')
                        dst_index_list = [base_index + offset for base_index in base_index_list]
                        rearange_idx_list.extend(dst_index_list)
                        rearange_idx_list_list.append(dst_index_list)
                        # rearange_split_img_id_idxs_list.append(split_img_id_idxs[dst_index_list])

                    ############################## plot 验证 ##############################
                    # img_boxes = [(deepcopy(img), cur_box, cur_posid) for img, cur_box, cur_posid in zip(image[:-1], box, rearange_idx_list_list)]
                    # self.eval_posid_by_plot(img_boxes, rope_pos_id_version, None)
                    # img_boxes = [(deepcopy(img), cur_box, cur_posid) for img, cur_box, cur_posid in zip(image[:-1], box, rearange_idx_list_list)]
                    # self.eval_posid_by_plot(img_boxes, rope_pos_id_version, split_img_id_idxs)

                    ############################## rearrange ##############################
                    split_img_id_idxs = split_img_id_idxs[rearange_idx_list]

                    rope_pos_id_list.append(split_img_id_idxs)
                    # thumbnail_id_idxs = torch.linspace(last_record_pos_id + 1, last_record_pos_id + 2, num_image_token + 1)[1:].to(dtype=dtype)  # 缩略图
                    thumbnail_id_idxs = origin_split_img_id_idxs.reshape([num_image_token, -1]).to(dtype=dtype).mean(dim=1).view(-1)
                    rope_pos_id_list.append(thumbnail_id_idxs)
                    last_record_pos_id = origin_split_img_id_idxs[-1].long()
                else:
                    thumbnail_id_idxs = torch.linspace(last_record_pos_id, last_record_pos_id + 1, num_image_token + 1)[1:].to(dtype=dtype)  # 缩略图
                    rope_pos_id_list.append(thumbnail_id_idxs)
                    last_record_pos_id = (last_record_pos_id + 1).long()

                # 验证是否能够恢复为等差数列
                if num_tile > 1:
                    gt_pos_id = torch.linspace(last_record_pos_id - 1 - (num_tile - 1), last_record_pos_id - 1, (num_tile - 1) * num_image_token + 1)[1:].to(dtype=dtype)
                    # self.eval_posid_by_rearange(box, rope_pos_id_list, gt_pos_id, num_tile, dtype)

            elif rope_pos_id_version == 'v2':
                # 子图处理方式同文本（N 个子图分配 N * 256 个 global id）；一个缩略图分配 256 * N 个的 global id.
                # 子图处理同 v0, v1，也对 global id 根据空间关系做 arrange
                if num_sub_imgs > 0:
                    split_img_id_idxs = torch.linspace(last_record_pos_id, last_record_pos_id + num_sub_imgs * num_image_token, num_sub_imgs * num_image_token + 1)[1:].long()  # long 数值的 tensor 作为变换的数据取值
                    last_id_for_split_img = last_record_pos_id + num_sub_imgs * num_image_token
                    origin_split_img_id_idxs = split_img_id_idxs
                    ############################## 进行位置变换 ##############################
                    # 先计算第一个子图对应 index
                    rearange_idx_list = []
                    rearange_idx_list_list = []
                    base_index_list = []
                    # rearange_split_img_id_idxs_list = []
                    num_img_token_in_length = int(num_image_token ** 0.5)
                    num_patch_width = int(box[-1][2] // box[0][2])
                    num_patch_height = int(box[-1][3] // box[0][2])
                    assert num_patch_width * num_patch_height == len(box)

                    num_total_patch_width_token = num_patch_width * num_img_token_in_length
                    num_total_patch_height_token = num_patch_height * num_img_token_in_length
                    assert num_total_patch_width_token * num_total_patch_height_token == num_sub_imgs * num_image_token, (
                        num_total_patch_width_token * num_total_patch_height_token, num_sub_imgs * num_image_token)

                    for k in range(num_image_token):
                        map_idx = (k // num_img_token_in_length) * num_total_patch_width_token + (
                                k % num_img_token_in_length)
                        base_index_list.append(map_idx)

                    # 计算其他子图对应第一个子图的 offset

                    for k in range(num_sub_imgs):
                        patch_row = k // num_patch_width
                        patch_col = k % num_patch_width
                        offset = patch_row * (
                                num_image_token * num_patch_width) + patch_col * num_img_token_in_length
                        # print(f'{k=}, {offset=}')
                        dst_index_list = [base_index + offset for base_index in base_index_list]
                        rearange_idx_list.extend(dst_index_list)
                        rearange_idx_list_list.append(dst_index_list)
                        # rearange_split_img_id_idxs_list.append(split_img_id_idxs[dst_index_list])

                    ############################## plot 验证 ##############################

                    # img_boxes = [(deepcopy(img), cur_box, cur_posid) for img, cur_box, cur_posid in
                    #              zip(image[:-1], box, rearange_idx_list_list)]
                    # self.eval_posid_by_plot(img_boxes, rope_pos_id_version, None)
                    # img_boxes = [(deepcopy(img), cur_box, cur_posid) for img, cur_box, cur_posid in
                    #              zip(image[:-1], box, rearange_idx_list_list)]
                    # self.eval_posid_by_plot(img_boxes, rope_pos_id_version, split_img_id_idxs)

                    ############################## rearrange ##############################
                    split_img_id_idxs = split_img_id_idxs[rearange_idx_list]

                    rope_pos_id_list.append(split_img_id_idxs)
                    thumbnail_id_idxs = origin_split_img_id_idxs.reshape([num_image_token, -1]).to(dtype=dtype).mean(dim=1).view(-1)
                    rope_pos_id_list.append(thumbnail_id_idxs)
                    last_record_pos_id = origin_split_img_id_idxs[-1].long()
                else:
                    thumbnail_id_idxs = torch.linspace(last_record_pos_id, last_record_pos_id + num_image_token, num_image_token + 1)[1:].long()  # 缩略图，和 default 处理一致
                    rope_pos_id_list.append(thumbnail_id_idxs)
                    last_record_pos_id = thumbnail_id_idxs[-1].long()

                # 验证是否能够恢复为等差数列
                if num_tile > 1:
                    gt_pos_id = torch.linspace(last_id_for_split_img - num_image_token * num_sub_imgs,
                                               last_id_for_split_img,
                                               num_sub_imgs * num_image_token + 1)[1:].long()
                    # self.eval_posid_by_rearange(box, rope_pos_id_list, gt_pos_id, num_tile, gt_pos_id.dtype)

            elif rope_pos_id_version == 'v3':
                # N 个子图共用跨度为 256 的 global id；一个缩略图正常分配 256 个 global id
                if num_sub_imgs > 0:
                    split_img_id_idxs = torch.linspace(last_record_pos_id, last_record_pos_id + num_image_token, num_sub_imgs * num_image_token + 1)[1:].to(dtype=dtype)  # 小数数值的 tensor 作为变换的数据取值
                    origin_split_img_id_idxs = split_img_id_idxs
                    ############################## 进行位置变换 ##############################
                    # 先计算第一个子图对应 index
                    rearange_idx_list = []
                    rearange_idx_list_list = []
                    base_index_list = []
                    # rearange_split_img_id_idxs_list = []
                    num_img_token_in_length = int(num_image_token ** 0.5)
                    num_patch_width = int(box[-1][2] // box[0][2])
                    num_patch_height = int(box[-1][3] // box[0][2])
                    assert num_patch_width * num_patch_height == len(box)

                    num_total_patch_width_token = num_patch_width * num_img_token_in_length
                    num_total_patch_height_token = num_patch_height * num_img_token_in_length
                    assert num_total_patch_width_token * num_total_patch_height_token == num_sub_imgs * num_image_token, (
                        num_total_patch_width_token * num_total_patch_height_token, num_sub_imgs * num_image_token)

                    for k in range(num_image_token):
                        map_idx = (k // num_img_token_in_length) * num_total_patch_width_token + (
                                k % num_img_token_in_length)
                        base_index_list.append(map_idx)

                    # 计算其他子图对应第一个子图的 offset

                    for k in range(num_sub_imgs):
                        patch_row = k // num_patch_width
                        patch_col = k % num_patch_width
                        offset = patch_row * (
                                num_image_token * num_patch_width) + patch_col * num_img_token_in_length
                        # print(f'{k=}, {offset=}')
                        dst_index_list = [base_index + offset for base_index in base_index_list]
                        rearange_idx_list.extend(dst_index_list)
                        rearange_idx_list_list.append(dst_index_list)
                        # rearange_split_img_id_idxs_list.append(split_img_id_idxs[dst_index_list])

                    ############################## plot 验证 ##############################

                    # img_boxes = [(deepcopy(img), cur_box, cur_posid) for img, cur_box, cur_posid in
                    #              zip(image[:-1], box, rearange_idx_list_list)]
                    # self.eval_posid_by_plot(img_boxes, rope_pos_id_version, None)
                    # img_boxes = [(deepcopy(img), cur_box, cur_posid) for img, cur_box, cur_posid in
                    #              zip(image[:-1], box, rearange_idx_list_list)]
                    # self.eval_posid_by_plot(img_boxes, rope_pos_id_version, split_img_id_idxs)

                    ############################## rearrange ##############################
                    split_img_id_idxs = split_img_id_idxs[rearange_idx_list]

                    rope_pos_id_list.append(split_img_id_idxs)
                    thumbnail_id_idxs = origin_split_img_id_idxs.reshape([num_image_token, -1]).to(dtype=dtype).mean(dim=1).view(-1)
                    rope_pos_id_list.append(thumbnail_id_idxs)
                    last_record_pos_id = origin_split_img_id_idxs[-1].long()
                else:
                    thumbnail_id_idxs = torch.linspace(last_record_pos_id, last_record_pos_id + num_image_token, num_image_token + 1)[1:].to(dtype=dtype)  # 缩略图，和 default 处理一致
                    rope_pos_id_list.append(thumbnail_id_idxs)
                    last_record_pos_id = thumbnail_id_idxs[-1].to(dtype=dtype)

                # 验证是否能够恢复为等差数列
                if num_tile > 1:
                    gt_pos_id = torch.linspace(last_record_pos_id - num_image_token - num_image_token,
                                               last_record_pos_id - num_image_token,
                                               num_sub_imgs * num_image_token + 1)[1:].to(dtype=dtype)
                    # self.eval_posid_by_rearange(box, rope_pos_id_list, gt_pos_id, num_tile, gt_pos_id.dtype)

            elif rope_pos_id_version == 'v4':
                # stride 是可变长的
                assert rope_pos_id_stride is not None, 'when rope_pos_id_version == v4, rope_pos_id_stride should not be None'
                if num_sub_imgs > 0:
                    num_sub_image_tokens = num_image_token * num_sub_imgs
                    split_img_id_idxs = torch.linspace(last_record_pos_id, last_record_pos_id + rope_pos_id_stride, num_sub_imgs * num_image_token + 1)[1:].to(dtype=dtype)  # 小数数值的 tensor 作为变换的数据取值
                    assert len(split_img_id_idxs) == num_sub_image_tokens
                    origin_split_img_id_idxs = split_img_id_idxs
                    ############################## 进行位置变换 ##############################
                    # 先计算第一个子图对应 index
                    rearange_idx_list = []
                    rearange_idx_list_list = []
                    base_index_list = []
                    # rearange_split_img_id_idxs_list = []
                    num_img_token_in_length = int(num_image_token ** 0.5)
                    num_patch_width = int(box[-1][2] // box[0][2])
                    num_patch_height = int(box[-1][3] // box[0][2])
                    assert num_patch_width * num_patch_height == len(box)

                    num_total_patch_width_token = num_patch_width * num_img_token_in_length
                    num_total_patch_height_token = num_patch_height * num_img_token_in_length
                    assert num_total_patch_width_token * num_total_patch_height_token == num_sub_imgs * num_image_token, (
                        num_total_patch_width_token * num_total_patch_height_token, num_sub_imgs * num_image_token)

                    for k in range(num_image_token):
                        map_idx = (k // num_img_token_in_length) * num_total_patch_width_token + (
                                k % num_img_token_in_length)
                        base_index_list.append(map_idx)

                    # 计算其他子图对应第一个子图的 offset

                    for k in range(num_sub_imgs):
                        patch_row = k // num_patch_width
                        patch_col = k % num_patch_width
                        offset = patch_row * (num_image_token * num_patch_width) + patch_col * num_img_token_in_length
                        # print(f'{k=}, {offset=}')
                        dst_index_list = [base_index + offset for base_index in base_index_list]
                        rearange_idx_list.extend(dst_index_list)
                        rearange_idx_list_list.append(dst_index_list)
                        # rearange_split_img_id_idxs_list.append(split_img_id_idxs[dst_index_list])

                    ############################## plot 验证 ##############################

                    # img_boxes = [(deepcopy(img), cur_box, cur_posid) for img, cur_box, cur_posid in
                    #              zip(image[:-1], box, rearange_idx_list_list)]
                    # self.eval_posid_by_plot(img_boxes, rope_pos_id_version, None)
                    # img_boxes = [(deepcopy(img), cur_box, cur_posid) for img, cur_box, cur_posid in
                    #              zip(image[:-1], box, rearange_idx_list_list)]
                    # self.eval_posid_by_plot(img_boxes, rope_pos_id_version, split_img_id_idxs)

                    ############################## rearrange ##############################
                    split_img_id_idxs = split_img_id_idxs[rearange_idx_list]

                    rope_pos_id_list.append(split_img_id_idxs)
                    thumbnail_id_idxs = origin_split_img_id_idxs.reshape([num_image_token, -1]).to(dtype=dtype).mean(dim=1).view(-1)
                    rope_pos_id_list.append(thumbnail_id_idxs)
                    last_record_pos_id = origin_split_img_id_idxs[-1].long()
                else:
                    thumbnail_id_idxs = torch.linspace(last_record_pos_id, last_record_pos_id + num_image_token, num_image_token + 1)[1:].to(dtype=dtype)  # 缩略图，和 default 处理一致
                    rope_pos_id_list.append(thumbnail_id_idxs)
                    last_record_pos_id = thumbnail_id_idxs[-1].to(dtype=dtype)

            elif rope_pos_id_version == 'v5':
                assert rope_pos_id_stride is not None, 'when rope_pos_id_version == v5, self.rope_pos_id_stride should not be None'
                small_stride = rope_pos_id_stride / num_image_token
                # split_img_id_idxs = torch.arange(last_record_pos_id, last_record_pos_id + small_stride * (num_image_token * num_tile + 1), small_stride)[1:].to(dtype=dtype)
                split_img_id_idxs = torch.linspace(last_record_pos_id,last_record_pos_id+small_stride*(num_image_token * num_tile ),(num_image_token * num_tile + 1))[1:].to(dtype=dtype)
                rope_pos_id_list.append(split_img_id_idxs)
                last_record_pos_id = torch.ceil(split_img_id_idxs[-1]).long()
            elif rope_pos_id_version == 'v6':
                random_from=[1,2,4,8,16,32,64,128,256]
                rope_pos_id_stride=random.choice(random_from)
                small_stride = rope_pos_id_stride / num_image_token
                # split_img_id_idxs = torch.arange(last_record_pos_id, last_record_pos_id + small_stride * (num_image_token * num_tile + 1), small_stride)[1:].to(dtype=dtype)
                split_img_id_idxs = torch.linspace(last_record_pos_id,last_record_pos_id+small_stride*(num_image_token * num_tile ),(num_image_token * num_tile + 1))[1:].to(dtype=dtype)
                rope_pos_id_list.append(split_img_id_idxs)
                last_record_pos_id = torch.ceil(split_img_id_idxs[-1]).long()
            elif rope_pos_id_version == 'default':
                # baseline
                # 无特殊处理的做法
                split_img_id_idxs = torch.linspace(last_record_pos_id,
                                                   last_record_pos_id + (num_tile - 1) * num_image_token,
                                                   (num_tile - 1) * num_image_token + 1)[1:].to(dtype=dtype)  # 子图
                rope_pos_id_list.append(split_img_id_idxs)
                thumbnail_id_idxs = torch.linspace(last_record_pos_id + (num_tile - 1) * num_image_token,
                                                   last_record_pos_id + num_tile * num_image_token,
                                                   num_image_token + 1)[1:].to(dtype=dtype)  # 缩略图
                rope_pos_id_list.append(thumbnail_id_idxs)
                last_record_pos_id = (last_record_pos_id + num_tile * num_image_token).long()
            else:
                raise NotImplementedError(f'not implement for {rope_pos_id_version}')
            try:
                start_index = image_start_token_id_idxs[i] + num_tile * num_image_token + 1
                assert input_ids_0[start_index] == image_end_token_id # 下一次迭代的开头应该是 IMG_END_TOKEN
                assert start_index == image_end_token_id_idxs[i] # 下一次迭代的开头应该是 IMG_END_TOKEN
            except:
                import ipdb
                ipdb.set_trace()

        if image_end_token_id_idxs[-1] != input_ids_0.shape[0] - 1:
            # 末尾还有待处理的非图像 id 的情况
            assert image_end_token_id_idxs[-1] == start_index # 应当从最后一个 IMG_END_TOKEN 开始
            rope_pos_id_pre = attention_mask_0[start_index:].long().cumsum(-1) - 1 + (last_record_pos_id + 1)
            rope_pos_id_pre.masked_fill_(attention_mask_0[start_index:] == 0, 1)
            rope_pos_id_list.append(rope_pos_id_pre)

        rope_pos_id_list=[_.to('cpu') for _ in rope_pos_id_list]
        rope_pos_id = torch.cat(rope_pos_id_list).to(dtype=dtype)
        if rope_pos_id_version == 'default':
            rope_pos_id = rope_pos_id.long() # 不做特殊处理的 rope_pos_id 应当等于 position_ids
            assert torch.equal(rope_pos_id, position_id.to(rope_pos_id.device)), (rope_pos_id, position_id.to(rope_pos_id.device))
            assert torch.allclose(rope_pos_id, position_id.to(rope_pos_id.device), atol=1e-32)
        
        assert rope_pos_id.shape == input_ids_0.shape

        return list(rope_pos_id.numpy())