# -------------------------------------------------------- # 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_END_TOKEN='', IMG_CONTEXT_TOKEN='', 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 '' not in question: question = '\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_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_END_TOKEN='', IMG_CONTEXT_TOKEN='', verbose=False,**kwargs): if history is None and pixel_values is not None and '' not in question: question = '\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_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}', '') 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_END_TOKEN='',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())