OpenGVLab/InternVL2-40B

1# --------------------------------------------------------2# InternVL3# Copyright (c) 2024 OpenGVLab4# Licensed under The MIT License [see LICENSE for details]5# --------------------------------------------------------6 7import warnings8from typing import Any, List, Optional, Tuple, Union9 10import torch.utils.checkpoint11import transformers12from torch import nn13from torch.nn import CrossEntropyLoss14from transformers import AutoModel, GenerationConfig, LlamaForCausalLM15from transformers.modeling_outputs import CausalLMOutputWithPast16from transformers.modeling_utils import PreTrainedModel17from transformers.utils import ModelOutput, logging18 19from .configuration_internvl_chat import InternVLChatConfig20from .conversation import get_conv_template21from .modeling_intern_vit import InternVisionModel, has_flash_attn22 23logger = logging.get_logger(__name__)24 25 26def version_cmp(v1, v2, op='eq'):27    import operator28 29    from packaging import version30    op_func = getattr(operator, op)31    return op_func(version.parse(v1), version.parse(v2))32 33 34class InternVLChatModel(PreTrainedModel):35    config_class = InternVLChatConfig36    main_input_name = 'pixel_values'37    base_model_prefix = 'language_model'38    _supports_flash_attn_2 = True39    supports_gradient_checkpointing = True40    _no_split_modules = ['InternVisionModel', 'LlamaDecoderLayer']41 42    def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None, use_flash_attn=True):43        super().__init__(config)44 45        assert version_cmp(transformers.__version__, '4.37.0', 'ge')46        image_size = config.force_image_size or config.vision_config.image_size47        patch_size = config.vision_config.patch_size48        self.patch_size = patch_size49        self.select_layer = config.select_layer50        self.template = config.template51        self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2))52        self.downsample_ratio = config.downsample_ratio53        self.ps_version = config.ps_version54        use_flash_attn = use_flash_attn if has_flash_attn else False55        config.vision_config.use_flash_attn = True if use_flash_attn else False56        config.llm_config._attn_implementation = 'flash_attention_2' if use_flash_attn else 'eager'57 58        logger.info(f'num_image_token: {self.num_image_token}')59        logger.info(f'ps_version: {self.ps_version}')60        if vision_model is not None:61            self.vision_model = vision_model62        else:63            self.vision_model = InternVisionModel(config.vision_config)64        if language_model is not None:65            self.language_model = language_model66        else:67            if config.llm_config.architectures[0] == 'LlamaForCausalLM':68                self.language_model = LlamaForCausalLM(config.llm_config)69            else:70                raise NotImplementedError(f'{config.llm_config.architectures[0]} is not implemented.')71 72        vit_hidden_size = config.vision_config.hidden_size73        llm_hidden_size = config.llm_config.hidden_size74 75        self.mlp1 = nn.Sequential(76            nn.LayerNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2),77            nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size),78            nn.GELU(),79            nn.Linear(llm_hidden_size, llm_hidden_size)80        )81 82        self.img_context_token_id = None83        self.conv_template = get_conv_template(self.template)84        self.system_message = self.conv_template.system_message85 86    def forward(87            self,88            pixel_values: torch.FloatTensor,89            input_ids: torch.LongTensor = None,90            attention_mask: Optional[torch.Tensor] = None,91            position_ids: Optional[torch.LongTensor] = None,92            image_flags: Optional[torch.LongTensor] = None,93            past_key_values: Optional[List[torch.FloatTensor]] = None,94            labels: Optional[torch.LongTensor] = None,95            use_cache: Optional[bool] = None,96            output_attentions: Optional[bool] = None,97            output_hidden_states: Optional[bool] = None,98            return_dict: Optional[bool] = None,99    ) -> Union[Tuple, CausalLMOutputWithPast]:100        return_dict = return_dict if return_dict is not None else self.config.use_return_dict101 102        image_flags = image_flags.squeeze(-1)103        input_embeds = self.language_model.get_input_embeddings()(input_ids).clone()104 105        vit_embeds = self.extract_feature(pixel_values)106        vit_embeds = vit_embeds[image_flags == 1]107        vit_batch_size = pixel_values.shape[0]108 109        B, N, C = input_embeds.shape110        input_embeds = input_embeds.reshape(B * N, C)111 112        if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:113            print(f'dynamic ViT batch size: {vit_batch_size}, images per sample: {vit_batch_size / B}, dynamic token length: {N}')114 115        input_ids = input_ids.reshape(B * N)116        selected = (input_ids == self.img_context_token_id)117        try:118            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)119        except Exception as e:120            vit_embeds = vit_embeds.reshape(-1, C)121            print(f'warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, '122                  f'vit_embeds.shape={vit_embeds.shape}')123            n_token = selected.sum()124            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds[:n_token]125 126        input_embeds = input_embeds.reshape(B, N, C)127 128        outputs = self.language_model(129            inputs_embeds=input_embeds,130            attention_mask=attention_mask,131            position_ids=position_ids,132            past_key_values=past_key_values,133            use_cache=use_cache,134            output_attentions=output_attentions,135            output_hidden_states=output_hidden_states,136            return_dict=return_dict,137        )138        logits = outputs.logits139 140        loss = None141        if labels is not None:142            # Shift so that tokens < n predict n143            shift_logits = logits[..., :-1, :].contiguous()144            shift_labels = labels[..., 1:].contiguous()145            # Flatten the tokens146            loss_fct = CrossEntropyLoss()147            shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)148            shift_labels = shift_labels.view(-1)149            # Enable model parallelism150            shift_labels = shift_labels.to(shift_logits.device)151            loss = loss_fct(shift_logits, shift_labels)152 153        if not return_dict:154            output = (logits,) + outputs[1:]155            return (loss,) + output if loss is not None else output156 157        return CausalLMOutputWithPast(158            loss=loss,159            logits=logits,160            past_key_values=outputs.past_key_values,161            hidden_states=outputs.hidden_states,162            attentions=outputs.attentions,163        )164 165    def pixel_shuffle(self, x, scale_factor=0.5):166        n, w, h, c = x.size()167        # N, W, H, C --> N, W, H * scale, C // scale168        x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))169        # N, W, H * scale, C // scale --> N, H * scale, W, C // scale170        x = x.permute(0, 2, 1, 3).contiguous()171        # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)172        x = x.view(n, int(h * scale_factor), int(w * scale_factor),173                   int(c / (scale_factor * scale_factor)))174        if self.ps_version == 'v1':175            warnings.warn("In ps_version 'v1', the height and width have not been swapped back, "176                          'which results in a transposed image.')177        else:178            x = x.permute(0, 2, 1, 3).contiguous()179        return x180 181    def extract_feature(self, pixel_values):182        if self.select_layer == -1:183            vit_embeds = self.vision_model(184                pixel_values=pixel_values,185                output_hidden_states=False,186                return_dict=True).last_hidden_state187        else:188            vit_embeds = self.vision_model(189                pixel_values=pixel_values,190                output_hidden_states=True,191                return_dict=True).hidden_states[self.select_layer]192        vit_embeds = vit_embeds[:, 1:, :]193 194        h = w = int(vit_embeds.shape[1] ** 0.5)195        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)196        vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)197        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])198        vit_embeds = self.mlp1(vit_embeds)199        return vit_embeds200 201    def batch_chat(self, tokenizer, pixel_values, questions, generation_config, num_patches_list=None,202                   history=None, return_history=False, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>',203                   IMG_CONTEXT_TOKEN='<IMG_CONTEXT>', verbose=False, image_counts=None):204        if history is not None or return_history:205            print('Now multi-turn chat is not supported in batch_chat.')206            raise NotImplementedError207 208        if image_counts is not None:209            num_patches_list = image_counts210            print('Warning: `image_counts` is deprecated. Please use `num_patches_list` instead.')211 212        img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)213        self.img_context_token_id = img_context_token_id214 215        if verbose and pixel_values is not None:216            image_bs = pixel_values.shape[0]217            print(f'dynamic ViT batch size: {image_bs}')218 219        queries = []220        for idx, num_patches in enumerate(num_patches_list):221            question = questions[idx]222            if pixel_values is not None and '<image>' not in question:223                question = '<image>\n' + question224            template = get_conv_template(self.template)225            template.system_message = self.system_message226            template.append_message(template.roles[0], question)227            template.append_message(template.roles[1], None)228            query = template.get_prompt()229 230            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN231            query = query.replace('<image>', image_tokens, 1)232            queries.append(query)233 234        tokenizer.padding_side = 'left'235        model_inputs = tokenizer(queries, return_tensors='pt', padding=True)236        input_ids = model_inputs['input_ids'].to(self.device)237        attention_mask = model_inputs['attention_mask'].to(self.device)238        eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())239        generation_config['eos_token_id'] = eos_token_id240        generation_output = self.generate(241            pixel_values=pixel_values,242            input_ids=input_ids,243            attention_mask=attention_mask,244            **generation_config245        )246        responses = tokenizer.batch_decode(generation_output, skip_special_tokens=True)247        responses = [response.split(template.sep.strip())[0].strip() for response in responses]248        return responses249 250    def chat(self, tokenizer, pixel_values, question, generation_config, history=None, return_history=False,251             num_patches_list=None, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>',252             verbose=False):253 254        if history is None and pixel_values is not None and '<image>' not in question:255            question = '<image>\n' + question256 257        if num_patches_list is None:258            num_patches_list = [pixel_values.shape[0]] if pixel_values is not None else []259        assert pixel_values is None or len(pixel_values) == sum(num_patches_list)260 261        img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)262        self.img_context_token_id = img_context_token_id263 264        template = get_conv_template(self.template)265        template.system_message = self.system_message266        eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())267 268        history = [] if history is None else history269        for (old_question, old_answer) in history:270            template.append_message(template.roles[0], old_question)271            template.append_message(template.roles[1], old_answer)272        template.append_message(template.roles[0], question)273        template.append_message(template.roles[1], None)274        query = template.get_prompt()275 276        if verbose and pixel_values is not None:277            image_bs = pixel_values.shape[0]278            print(f'dynamic ViT batch size: {image_bs}')279 280        for num_patches in num_patches_list:281            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN282            query = query.replace('<image>', image_tokens, 1)283 284        model_inputs = tokenizer(query, return_tensors='pt')285        input_ids = model_inputs['input_ids'].to(self.device)286        attention_mask = model_inputs['attention_mask'].to(self.device)287        generation_config['eos_token_id'] = eos_token_id288        generation_output = self.generate(289            pixel_values=pixel_values,290            input_ids=input_ids,291            attention_mask=attention_mask,292            **generation_config293        )294        response = tokenizer.batch_decode(generation_output, skip_special_tokens=True)[0]295        response = response.split(template.sep.strip())[0].strip()296        history.append((question, response))297        if return_history:298            return response, history299        else:300            query_to_print = query.replace(IMG_CONTEXT_TOKEN, '')301            query_to_print = query_to_print.replace(f'{IMG_START_TOKEN}{IMG_END_TOKEN}', '<image>')302            if verbose:303                print(query_to_print, response)304            return response305 306    @torch.no_grad()307    def generate(308            self,309            pixel_values: Optional[torch.FloatTensor] = None,310            input_ids: Optional[torch.FloatTensor] = None,311            attention_mask: Optional[torch.LongTensor] = None,312            visual_features: Optional[torch.FloatTensor] = None,313            generation_config: Optional[GenerationConfig] = None,314            output_hidden_states: Optional[bool] = None,315            **generate_kwargs,316    ) -> torch.LongTensor:317 318        assert self.img_context_token_id is not None319        if pixel_values is not None:320            if visual_features is not None:321                vit_embeds = visual_features322            else:323                vit_embeds = self.extract_feature(pixel_values)324            input_embeds = self.language_model.get_input_embeddings()(input_ids)325            B, N, C = input_embeds.shape326            input_embeds = input_embeds.reshape(B * N, C)327 328            input_ids = input_ids.reshape(B * N)329            selected = (input_ids == self.img_context_token_id)330            assert selected.sum() != 0331            input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)332 333            input_embeds = input_embeds.reshape(B, N, C)334        else:335            input_embeds = self.language_model.get_input_embeddings()(input_ids)336 337        outputs = self.language_model.generate(338            inputs_embeds=input_embeds,339            attention_mask=attention_mask,340            generation_config=generation_config,341            output_hidden_states=output_hidden_states,342            use_cache=True,343            **generate_kwargs,344        )345 346        return outputs347 348    @property349    def lm_head(self):350        return self.language_model.get_output_embeddings()351 352    def get_input_embeddings(self):353        return self.language_model.get_input_embeddings()354 355    def get_output_embeddings(self):356        return self.language_model.get_output_embeddings()357