OpenGVLab/InternVL2_5-26B
331.5k
1# --------------------------------------------------------2# InternVL3# Copyright (c) 2024 OpenGVLab4# Licensed under The MIT License [see LICENSE for details]5# --------------------------------------------------------6 7import warnings8from typing import List, Optional, Tuple, Union9 10import torch.utils.checkpoint11import transformers12from torch import nn13from torch.nn import CrossEntropyLoss14from transformers import (AutoModel, GenerationConfig, LlamaForCausalLM,15 LlamaTokenizer)16from transformers.modeling_outputs import CausalLMOutputWithPast17from transformers.modeling_utils import PreTrainedModel18from transformers.utils import ModelOutput, logging19 20from .configuration_internvl_chat import InternVLChatConfig21from .conversation import get_conv_template22from .modeling_intern_vit import InternVisionModel, has_flash_attn23from .modeling_internlm2 import InternLM2ForCausalLM24 25logger = logging.get_logger(__name__)26 27 28def version_cmp(v1, v2, op='eq'):29 import operator30 31 from packaging import version32 op_func = getattr(operator, op)33 return op_func(version.parse(v1), version.parse(v2))34 35 36class InternVLChatModel(PreTrainedModel):37 config_class = InternVLChatConfig38 main_input_name = 'pixel_values'39 base_model_prefix = 'language_model'40 _supports_flash_attn_2 = True41 supports_gradient_checkpointing = True42 _no_split_modules = ['InternVisionModel', 'LlamaDecoderLayer', 'InternLM2DecoderLayer']43 44 def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None, use_flash_attn=True):45 super().__init__(config)46 47 assert version_cmp(transformers.__version__, '4.37.0', 'ge')48 image_size = config.force_image_size or config.vision_config.image_size49 patch_size = config.vision_config.patch_size50 self.patch_size = patch_size51 self.select_layer = config.select_layer52 self.template = config.template53 self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2))54 self.downsample_ratio = config.downsample_ratio55 self.ps_version = config.ps_version56 use_flash_attn = use_flash_attn if has_flash_attn else False57 config.vision_config.use_flash_attn = True if use_flash_attn else False58 config.llm_config.attn_implementation = 'flash_attention_2' if use_flash_attn else 'eager'59 60 logger.info(f'num_image_token: {self.num_image_token}')61 logger.info(f'ps_version: {self.ps_version}')62 if vision_model is not None:63 self.vision_model = vision_model64 else:65 self.vision_model = InternVisionModel(config.vision_config)66 if language_model is not None:67 self.language_model = language_model68 else:69 if config.llm_config.architectures[0] == 'LlamaForCausalLM':70 self.language_model = LlamaForCausalLM(config.llm_config)71 elif config.llm_config.architectures[0] == 'InternLM2ForCausalLM':72 self.language_model = InternLM2ForCausalLM(config.llm_config)73 else:74 raise NotImplementedError(f'{config.llm_config.architectures[0]} is not implemented.')75 76 vit_hidden_size = config.vision_config.hidden_size77 llm_hidden_size = config.llm_config.hidden_size78 79 self.mlp1 = nn.Sequential(80 nn.LayerNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2),81 nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size),82 nn.GELU(),83 nn.Linear(llm_hidden_size, llm_hidden_size)84 )85 86 self.img_context_token_id = None87 self.conv_template = get_conv_template(self.template)88 self.system_message = self.conv_template.system_message89 90 def forward(91 self,92 pixel_values: torch.FloatTensor,93 input_ids: torch.LongTensor = None,94 attention_mask: Optional[torch.Tensor] = None,95 position_ids: Optional[torch.LongTensor] = None,96 image_flags: Optional[torch.LongTensor] = None,97 past_key_values: Optional[List[torch.FloatTensor]] = None,98 labels: Optional[torch.LongTensor] = None,99 use_cache: Optional[bool] = None,100 output_attentions: Optional[bool] = None,101 output_hidden_states: Optional[bool] = None,102 return_dict: Optional[bool] = None,103 ) -> Union[Tuple, CausalLMOutputWithPast]:104 return_dict = return_dict if return_dict is not None else self.config.use_return_dict105 106 image_flags = image_flags.squeeze(-1)107 input_embeds = self.language_model.get_input_embeddings()(input_ids).clone()108 109 vit_embeds = self.extract_feature(pixel_values)110 vit_embeds = vit_embeds[image_flags == 1]111 vit_batch_size = pixel_values.shape[0]112 113 B, N, C = input_embeds.shape114 input_embeds = input_embeds.reshape(B * N, C)115 116 if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:117 print(f'dynamic ViT batch size: {vit_batch_size}, images per sample: {vit_batch_size / B}, dynamic token length: {N}')118 119 input_ids = input_ids.reshape(B * N)120 selected = (input_ids == self.img_context_token_id)121 try:122 input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)123 except Exception as e:124 vit_embeds = vit_embeds.reshape(-1, C)125 print(f'warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, '126 f'vit_embeds.shape={vit_embeds.shape}')127 n_token = selected.sum()128 input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds[:n_token]129 130 input_embeds = input_embeds.reshape(B, N, C)131 132 outputs = self.language_model(133 inputs_embeds=input_embeds,134 attention_mask=attention_mask,135 position_ids=position_ids,136 past_key_values=past_key_values,137 use_cache=use_cache,138 output_attentions=output_attentions,139 output_hidden_states=output_hidden_states,140 return_dict=return_dict,141 )142 logits = outputs.logits143 144 loss = None145 if labels is not None:146 # Shift so that tokens < n predict n147 shift_logits = logits[..., :-1, :].contiguous()148 shift_labels = labels[..., 1:].contiguous()149 # Flatten the tokens150 loss_fct = CrossEntropyLoss()151 shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)152 shift_labels = shift_labels.view(-1)153 # Enable model parallelism154 shift_labels = shift_labels.to(shift_logits.device)155 loss = loss_fct(shift_logits, shift_labels)156 157 if not return_dict:158 output = (logits,) + outputs[1:]159 return (loss,) + output if loss is not None else output160 161 return CausalLMOutputWithPast(162 loss=loss,163 logits=logits,164 past_key_values=outputs.past_key_values,165 hidden_states=outputs.hidden_states,166 attentions=outputs.attentions,167 )168 169 def pixel_shuffle(self, x, scale_factor=0.5):170 n, w, h, c = x.size()171 # N, W, H, C --> N, W, H * scale, C // scale172 x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))173 # N, W, H * scale, C // scale --> N, H * scale, W, C // scale174 x = x.permute(0, 2, 1, 3).contiguous()175 # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)176 x = x.view(n, int(h * scale_factor), int(w * scale_factor),177 int(c / (scale_factor * scale_factor)))178 if self.ps_version == 'v1':179 warnings.warn("In ps_version 'v1', the height and width have not been swapped back, "180 'which results in a transposed image.')181 else:182 x = x.permute(0, 2, 1, 3).contiguous()183 return x184 185 def extract_feature(self, pixel_values):186 if self.select_layer == -1:187 vit_embeds = self.vision_model(188 pixel_values=pixel_values,189 output_hidden_states=False,190 return_dict=True).last_hidden_state191 else:192 vit_embeds = self.vision_model(193 pixel_values=pixel_values,194 output_hidden_states=True,195 return_dict=True).hidden_states[self.select_layer]196 vit_embeds = vit_embeds[:, 1:, :]197 198 h = w = int(vit_embeds.shape[1] ** 0.5)199 vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)200 vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)201 vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])202 vit_embeds = self.mlp1(vit_embeds)203 return vit_embeds204 205 def batch_chat(self, tokenizer, pixel_values, questions, generation_config, num_patches_list=None,206 history=None, return_history=False, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>',207 IMG_CONTEXT_TOKEN='<IMG_CONTEXT>', verbose=False, image_counts=None):208 if history is not None or return_history:209 print('Now multi-turn chat is not supported in batch_chat.')210 raise NotImplementedError211 212 if image_counts is not None:213 num_patches_list = image_counts214 print('Warning: `image_counts` is deprecated. Please use `num_patches_list` instead.')215 216 img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)217 self.img_context_token_id = img_context_token_id218 219 if verbose and pixel_values is not None:220 image_bs = pixel_values.shape[0]221 print(f'dynamic ViT batch size: {image_bs}')222 223 queries = []224 for idx, num_patches in enumerate(num_patches_list):225 question = questions[idx]226 if pixel_values is not None and '<image>' not in question:227 question = '<image>\n' + question228 template = get_conv_template(self.template)229 template.system_message = self.system_message230 template.append_message(template.roles[0], question)231 template.append_message(template.roles[1], None)232 query = template.get_prompt()233 234 image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN235 query = query.replace('<image>', image_tokens, 1)236 queries.append(query)237 238 tokenizer.padding_side = 'left'239 model_inputs = tokenizer(queries, return_tensors='pt', padding=True)240 input_ids = model_inputs['input_ids'].to(self.device)241 attention_mask = model_inputs['attention_mask'].to(self.device)242 eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())243 generation_config['eos_token_id'] = eos_token_id244 generation_output = self.generate(245 pixel_values=pixel_values,246 input_ids=input_ids,247 attention_mask=attention_mask,248 **generation_config249 )250 responses = tokenizer.batch_decode(generation_output, skip_special_tokens=True)251 responses = [response.split(template.sep.strip())[0].strip() for response in responses]252 return responses253 254 def chat(self, tokenizer, pixel_values, question, generation_config, history=None, return_history=False,255 num_patches_list=None, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>',256 verbose=False):257 258 if history is None and pixel_values is not None and '<image>' not in question:259 question = '<image>\n' + question260 261 if num_patches_list is None:262 num_patches_list = [pixel_values.shape[0]] if pixel_values is not None else []263 assert pixel_values is None or len(pixel_values) == sum(num_patches_list)264 265 img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)266 self.img_context_token_id = img_context_token_id267 268 template = get_conv_template(self.template)269 template.system_message = self.system_message270 eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())271 272 history = [] if history is None else history273 for (old_question, old_answer) in history:274 template.append_message(template.roles[0], old_question)275 template.append_message(template.roles[1], old_answer)276 template.append_message(template.roles[0], question)277 template.append_message(template.roles[1], None)278 query = template.get_prompt()279 280 if verbose and pixel_values is not None:281 image_bs = pixel_values.shape[0]282 print(f'dynamic ViT batch size: {image_bs}')283 284 for num_patches in num_patches_list:285 image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN286 query = query.replace('<image>', image_tokens, 1)287 288 model_inputs = tokenizer(query, return_tensors='pt')289 input_ids = model_inputs['input_ids'].to(self.device)290 attention_mask = model_inputs['attention_mask'].to(self.device)291 generation_config['eos_token_id'] = eos_token_id292 generation_output = self.generate(293 pixel_values=pixel_values,294 input_ids=input_ids,295 attention_mask=attention_mask,296 **generation_config297 )298 response = tokenizer.batch_decode(generation_output, skip_special_tokens=True)[0]299 response = response.split(template.sep.strip())[0].strip()300 history.append((question, response))301 if return_history:302 return response, history303 else:304 query_to_print = query.replace(IMG_CONTEXT_TOKEN, '')305 query_to_print = query_to_print.replace(f'{IMG_START_TOKEN}{IMG_END_TOKEN}', '<image>')306 if verbose:307 print(query_to_print, response)308 return response309 310 @torch.no_grad()311 def generate(312 self,313 pixel_values: Optional[torch.FloatTensor] = None,314 input_ids: Optional[torch.FloatTensor] = None,315 attention_mask: Optional[torch.LongTensor] = None,316 visual_features: Optional[torch.FloatTensor] = None,317 generation_config: Optional[GenerationConfig] = None,318 output_hidden_states: Optional[bool] = None,319 **generate_kwargs,320 ) -> torch.LongTensor:321 322 assert self.img_context_token_id is not None323 if pixel_values is not None:324 if visual_features is not None:325 vit_embeds = visual_features326 else:327 vit_embeds = self.extract_feature(pixel_values)328 input_embeds = self.language_model.get_input_embeddings()(input_ids)329 B, N, C = input_embeds.shape330 input_embeds = input_embeds.reshape(B * N, C)331 332 input_ids = input_ids.reshape(B * N)333 selected = (input_ids == self.img_context_token_id)334 assert selected.sum() != 0335 input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)336 337 input_embeds = input_embeds.reshape(B, N, C)338 else:339 input_embeds = self.language_model.get_input_embeddings()(input_ids)340 341 outputs = self.language_model.generate(342 inputs_embeds=input_embeds,343 attention_mask=attention_mask,344 generation_config=generation_config,345 output_hidden_states=output_hidden_states,346 use_cache=True,347 **generate_kwargs,348 )349 350 return outputs351 352 @property353 def lm_head(self):354 return self.language_model.get_output_embeddings()355 356 def get_input_embeddings(self):357 return self.language_model.get_input_embeddings()358 359 def get_output_embeddings(self):360 return self.language_model.get_output_embeddings()361