OpenGVLab/InternVL3-2B
4624k
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 Qwen2ForCausalLM)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_attn23 24logger = logging.get_logger(__name__)25 26 27def version_cmp(v1, v2, op='eq'):28 import operator29 30 from packaging import version31 op_func = getattr(operator, op)32 return op_func(version.parse(v1), version.parse(v2))33 34 35class InternVLChatModel(PreTrainedModel):36 config_class = InternVLChatConfig37 main_input_name = 'pixel_values'38 base_model_prefix = 'language_model'39 _supports_flash_attn_2 = True40 supports_gradient_checkpointing = True41 _no_split_modules = ['InternVisionModel', 'LlamaDecoderLayer', 'Qwen2DecoderLayer']42 43 def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None, use_flash_attn=True):44 super().__init__(config)45 46 assert version_cmp(transformers.__version__, '4.37.0', 'ge')47 image_size = config.force_image_size or config.vision_config.image_size48 patch_size = config.vision_config.patch_size49 self.patch_size = patch_size50 self.select_layer = config.select_layer51 self.template = config.template52 self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2))53 self.downsample_ratio = config.downsample_ratio54 self.ps_version = config.ps_version55 use_flash_attn = use_flash_attn if has_flash_attn else False56 config.vision_config.use_flash_attn = True if use_flash_attn else False57 config.llm_config._attn_implementation = 'flash_attention_2' if use_flash_attn else 'eager'58 59 logger.info(f'num_image_token: {self.num_image_token}')60 logger.info(f'ps_version: {self.ps_version}')61 if vision_model is not None:62 self.vision_model = vision_model63 else:64 self.vision_model = InternVisionModel(config.vision_config)65 if language_model is not None:66 self.language_model = language_model67 else:68 if config.llm_config.architectures[0] == 'LlamaForCausalLM':69 self.language_model = LlamaForCausalLM(config.llm_config)70 elif config.llm_config.architectures[0] == 'Qwen2ForCausalLM':71 self.language_model = Qwen2ForCausalLM(config.llm_config)72 else:73 raise NotImplementedError(f'{config.llm_config.architectures[0]} is not implemented.')74 75 vit_hidden_size = config.vision_config.hidden_size76 llm_hidden_size = config.llm_config.hidden_size77 78 self.mlp1 = nn.Sequential(79 nn.LayerNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2),80 nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size),81 nn.GELU(),82 nn.Linear(llm_hidden_size, llm_hidden_size)83 )84 85 self.img_context_token_id = None86 self.conv_template = get_conv_template(self.template)87 self.system_message = self.conv_template.system_message88 89 def forward(90 self,91 pixel_values: torch.FloatTensor,92 input_ids: torch.LongTensor = None,93 attention_mask: Optional[torch.Tensor] = None,94 position_ids: Optional[torch.LongTensor] = None,95 image_flags: Optional[torch.LongTensor] = None,96 past_key_values: Optional[List[torch.FloatTensor]] = None,97 labels: Optional[torch.LongTensor] = None,98 use_cache: Optional[bool] = None,99 output_attentions: Optional[bool] = None,100 output_hidden_states: Optional[bool] = None,101 return_dict: Optional[bool] = None,102 ) -> Union[Tuple, CausalLMOutputWithPast]:103 return_dict = return_dict if return_dict is not None else self.config.use_return_dict104 105 image_flags = image_flags.squeeze(-1)106 input_embeds = self.language_model.get_input_embeddings()(input_ids).clone()107 108 vit_embeds = self.extract_feature(pixel_values)109 vit_embeds = vit_embeds[image_flags == 1]110 vit_batch_size = pixel_values.shape[0]111 112 B, N, C = input_embeds.shape113 input_embeds = input_embeds.reshape(B * N, C)114 115 if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:116 print(f'dynamic ViT batch size: {vit_batch_size}, images per sample: {vit_batch_size / B}, dynamic token length: {N}')117 118 input_ids = input_ids.reshape(B * N)119 selected = (input_ids == self.img_context_token_id)120 try:121 input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)122 except Exception as e:123 vit_embeds = vit_embeds.reshape(-1, C)124 print(f'warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, '125 f'vit_embeds.shape={vit_embeds.shape}')126 n_token = min(selected.sum(), vit_embeds.size(0))127 input_embeds[selected][:n_token] = input_embeds[selected][:n_token] * 0.0 + vit_embeds[:n_token]128 129 input_embeds = input_embeds.reshape(B, N, C)130 131 outputs = self.language_model(132 inputs_embeds=input_embeds,133 attention_mask=attention_mask,134 position_ids=position_ids,135 past_key_values=past_key_values,136 use_cache=use_cache,137 output_attentions=output_attentions,138 output_hidden_states=output_hidden_states,139 return_dict=return_dict,140 )141 logits = outputs.logits142 143 loss = None144 if labels is not None:145 # Shift so that tokens < n predict n146 shift_logits = logits[..., :-1, :].contiguous()147 shift_labels = labels[..., 1:].contiguous()148 # Flatten the tokens149 loss_fct = CrossEntropyLoss()150 shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)151 shift_labels = shift_labels.view(-1)152 # Enable model parallelism153 shift_labels = shift_labels.to(shift_logits.device)154 loss = loss_fct(shift_logits, shift_labels)155 156 if not return_dict:157 output = (logits,) + outputs[1:]158 return (loss,) + output if loss is not None else output159 160 return CausalLMOutputWithPast(161 loss=loss,162 logits=logits,163 past_key_values=outputs.past_key_values,164 hidden_states=outputs.hidden_states,165 attentions=outputs.attentions,166 )167 168 def pixel_shuffle(self, x, scale_factor=0.5):169 n, w, h, c = x.size()170 # N, W, H, C --> N, W, H * scale, C // scale171 x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))172 # N, W, H * scale, C // scale --> N, H * scale, W, C // scale173 x = x.permute(0, 2, 1, 3).contiguous()174 # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)175 x = x.view(n, int(h * scale_factor), int(w * scale_factor),176 int(c / (scale_factor * scale_factor)))177 if self.ps_version == 'v1':178 warnings.warn("In ps_version 'v1', the height and width have not been swapped back, "179 'which results in a transposed image.')180 else:181 x = x.permute(0, 2, 1, 3).contiguous()182 return x183 184 def extract_feature(self, pixel_values):185 if self.select_layer == -1:186 vit_embeds = self.vision_model(187 pixel_values=pixel_values,188 output_hidden_states=False,189 return_dict=True).last_hidden_state190 else:191 vit_embeds = self.vision_model(192 pixel_values=pixel_values,193 output_hidden_states=True,194 return_dict=True).hidden_states[self.select_layer]195 vit_embeds = vit_embeds[:, 1:, :]196 197 h = w = int(vit_embeds.shape[1] ** 0.5)198 vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)199 vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)200 vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])201 vit_embeds = self.mlp1(vit_embeds)202 return vit_embeds203 204 def batch_chat(self, tokenizer, pixel_values, questions, generation_config, num_patches_list=None,205 history=None, return_history=False, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>',206 IMG_CONTEXT_TOKEN='<IMG_CONTEXT>', verbose=False, image_counts=None):207 if history is not None or return_history:208 print('Now multi-turn chat is not supported in batch_chat.')209 raise NotImplementedError210 211 if image_counts is not None:212 num_patches_list = image_counts213 print('Warning: `image_counts` is deprecated. Please use `num_patches_list` instead.')214 215 img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)216 self.img_context_token_id = img_context_token_id217 218 if verbose and pixel_values is not None:219 image_bs = pixel_values.shape[0]220 print(f'dynamic ViT batch size: {image_bs}')221 222 queries = []223 for idx, num_patches in enumerate(num_patches_list):224 question = questions[idx]225 if pixel_values is not None and '<image>' not in question:226 question = '<image>\n' + question227 template = get_conv_template(self.template)228 template.system_message = self.system_message229 template.append_message(template.roles[0], question)230 template.append_message(template.roles[1], None)231 query = template.get_prompt()232 233 image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN234 query = query.replace('<image>', image_tokens, 1)235 queries.append(query)236 237 tokenizer.padding_side = 'left'238 model_inputs = tokenizer(queries, return_tensors='pt', padding=True)239 input_ids = model_inputs['input_ids'].to(self.device)240 attention_mask = model_inputs['attention_mask'].to(self.device)241 eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())242 generation_config['eos_token_id'] = eos_token_id243 generation_output = self.generate(244 pixel_values=pixel_values,245 input_ids=input_ids,246 attention_mask=attention_mask,247 **generation_config248 )249 responses = tokenizer.batch_decode(generation_output, skip_special_tokens=True)250 responses = [response.split(template.sep.strip())[0].strip() for response in responses]251 return responses252 253 def chat(self, tokenizer, pixel_values, question, generation_config, history=None, return_history=False,254 num_patches_list=None, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>',255 verbose=False):256 257 if history is None and pixel_values is not None and '<image>' not in question:258 question = '<image>\n' + question259 260 if num_patches_list is None:261 num_patches_list = [pixel_values.shape[0]] if pixel_values is not None else []262 assert pixel_values is None or len(pixel_values) == sum(num_patches_list)263 264 img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)265 self.img_context_token_id = img_context_token_id266 267 template = get_conv_template(self.template)268 template.system_message = self.system_message269 eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())270 271 history = [] if history is None else history272 for (old_question, old_answer) in history:273 template.append_message(template.roles[0], old_question)274 template.append_message(template.roles[1], old_answer)275 template.append_message(template.roles[0], question)276 template.append_message(template.roles[1], None)277 query = template.get_prompt()278 279 if verbose and pixel_values is not None:280 image_bs = pixel_values.shape[0]281 print(f'dynamic ViT batch size: {image_bs}')282 283 for num_patches in num_patches_list:284 image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN285 query = query.replace('<image>', image_tokens, 1)286 287 model_inputs = tokenizer(query, return_tensors='pt')288 input_ids = model_inputs['input_ids'].to(self.device)289 attention_mask = model_inputs['attention_mask'].to(self.device)290 generation_config['eos_token_id'] = eos_token_id291 generation_output = self.generate(292 pixel_values=pixel_values,293 input_ids=input_ids,294 attention_mask=attention_mask,295 **generation_config296 )297 response = tokenizer.batch_decode(generation_output, skip_special_tokens=True)[0]298 response = response.split(template.sep.strip())[0].strip()299 history.append((question, response))300 if return_history:301 return response, history302 else:303 query_to_print = query.replace(IMG_CONTEXT_TOKEN, '')304 query_to_print = query_to_print.replace(f'{IMG_START_TOKEN}{IMG_END_TOKEN}', '<image>')305 if verbose:306 print(query_to_print, response)307 return response308 309 @torch.no_grad()310 def generate(311 self,312 pixel_values: Optional[torch.FloatTensor] = None,313 input_ids: Optional[torch.FloatTensor] = None,314 attention_mask: Optional[torch.LongTensor] = None,315 visual_features: Optional[torch.FloatTensor] = None,316 generation_config: Optional[GenerationConfig] = None,317 output_hidden_states: Optional[bool] = None,318 **generate_kwargs,319 ) -> torch.LongTensor:320 321 assert self.img_context_token_id is not None322 if pixel_values is not None:323 if visual_features is not None:324 vit_embeds = visual_features325 else:326 vit_embeds = self.extract_feature(pixel_values)327 input_embeds = self.language_model.get_input_embeddings()(input_ids)328 B, N, C = input_embeds.shape329 input_embeds = input_embeds.reshape(B * N, C)330 331 input_ids = input_ids.reshape(B * N)332 selected = (input_ids == self.img_context_token_id)333 assert selected.sum() != 0334 input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)335 336 input_embeds = input_embeds.reshape(B, N, C)337 else:338 input_embeds = self.language_model.get_input_embeddings()(input_ids)339 340 outputs = self.language_model.generate(341 inputs_embeds=input_embeds,342 attention_mask=attention_mask,343 generation_config=generation_config,344 output_hidden_states=output_hidden_states,345 use_cache=True,346 **generate_kwargs,347 )348 349 return outputs350 351 @property352 def lm_head(self):353 return self.language_model.get_output_embeddings()354 355 def get_input_embeddings(self):356 return self.language_model.get_input_embeddings()357 358 def get_output_embeddings(self):359 return self.language_model.get_output_embeddings()360