Upload modeling_minicpmv.py with huggingface_hub

modeling_minicpmv.py

@@ -0,0 +1,221 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright @2023 AI, ZHIHU Inc. (zhihu.com)
+#
+# @author: wangchongyi <[email protected]>
+# @date: 2023/9/1
+#
+
+# coding=utf-8
+# Copyright 2024 RhapsodyAI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+from torch import nn
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+from transformers.utils import ModelOutput
+from transformers.modeling_utils import PreTrainedModel
+
+from .configuration_siglip import SiglipVisionConfig
+from .configuration_minicpm import MiniCPMConfig
+from .configuration_minicpmv import MiniCPMVConfig
+
+from .resampler import Resampler
+from .modeling_minicpm import MiniCPMForCausalLM
+from .modeling_siglip import SiglipVisionModel
+
+from transformers import LlamaTokenizer # for text processing
+
+
+@dataclass
+class CausalVLMOutput(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    vision_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+class MiniCPMVForCausalLM(PreTrainedModel):
+    model_type = "minicpm"
+    _supports_flash_attn_2 = True
+    
+    def __init__(self, config: MiniCPMVConfig, adaptive=False):
+        super().__init__(config)
+        
+        llm_config = config.llm_config
+        vpm_config = config.vpm_config
+        
+        self.query_num = config.query_num
+        self.patch_size = vpm_config.patch_size
+        self.adaptive = adaptive
+        self.slice_mode = config.slice_mode
+        self.max_slice_nums = config.max_slice_nums
+        self.mm_use_im_start_end = config.mm_use_im_start_end
+        
+        drop_vision_last_layer = config.drop_vision_last_layer
+        
+        # should assert vpm_config is SiglipVisionConfig
+        vpm = SiglipVisionModel(vpm_config).vision_model
+        
+        if drop_vision_last_layer: # drop last vision layer
+            vpm.encoder.layers = nn.ModuleList(vpm.encoder.layers[:-1])
+        
+        self.vpm = vpm
+        
+        # should assert llm_config is minicpmconfig
+        self.llm = MiniCPMForCausalLM(llm_config)
+        
+        embed_dim = llm_config.hidden_size
+        
+        self.resampler = Resampler(
+            num_queries=config.query_num,
+            embed_dim=embed_dim,
+            num_heads=embed_dim // 128,
+            kv_dim=vpm_config.hidden_size,
+            adaptive=adaptive
+        )
+        
+        return
+
+    def vpm_forward(self, data):
+        if 'vision_hidden_states' not in data:
+            dtype = self.vpm.embeddings.position_embedding.weight.dtype
+            device = self.vpm.embeddings.position_embedding.weight.device
+
+            pixel_values_list = data['pixel_values']
+            tgt_sizes = data['tgt_sizes']
+
+            vision_hidden_states = []
+
+            all_pixel_values = []
+            img_cnt = []
+
+            for pixel_values in pixel_values_list:
+                img_cnt.append(len(pixel_values))
+                all_pixel_values.extend([i.flatten(end_dim=1).permute(1, 0) for i in pixel_values])  # 42 * L
+
+            # exist image
+            if all_pixel_values:
+                tgt_sizes = torch.vstack(tgt_sizes).type(torch.int32)
+                max_patches = torch.max(tgt_sizes[:, 0] * tgt_sizes[:, 1])
+
+                all_pixel_values = torch.nn.utils.rnn.pad_sequence(all_pixel_values, batch_first=True, padding_value=0.0)
+                all_pixel_values = all_pixel_values.to(device) # here we finally could put `all_pixel_values` to device.
+                
+                B, L, _ = all_pixel_values.shape
+                all_pixel_values = all_pixel_values.permute(0, 2, 1).reshape(B, 3, -1, L)  # B, 3, 14, L
+
+                patch_attn_mask = torch.zeros((B, 1, max_patches), dtype=torch.bool, device=device)
+                for i in range(B):
+                    patch_attn_mask[i, :tgt_sizes[i][0] * tgt_sizes[i][1]] = True
+
+                vision_embedding = self.vpm(all_pixel_values.type(dtype), patch_attention_mask=patch_attn_mask).last_hidden_state
+                vision_embedding = self.resampler(vision_embedding, tgt_sizes)
+
+                start = 0
+                for pixel_values in pixel_values_list:
+                    img_cnt = len(pixel_values)
+                    if img_cnt > 0:
+                        vision_hidden_states.append(vision_embedding[start: start + img_cnt])
+                        start += img_cnt
+                    else:
+                        vision_hidden_states.append([])
+            else: # no image
+                if self.training:
+                    dummy_image = torch.zeros(
+                        (1, 3, 224, 224),
+                        device=device, dtype=dtype
+                    )
+                    # 这是一个 dummy feature
+                    tgt_sizes = torch.Tensor([[(224 // self.patch_size), math.ceil(224 / self.patch_size)]]).type(torch.int32)
+                    dummy_feature = self.resampler(self.vpm(dummy_image).last_hidden_state, tgt_sizes)
+                else:
+                    dummy_feature = []
+                for _ in range(len(pixel_values_list)):
+                    vision_hidden_states.append(dummy_feature)
+
+        else:
+            vision_hidden_states = data['vision_hidden_states']
+
+        if hasattr(self.llm.config, 'scale_emb'):
+            vllm_embedding = self.llm.model.embed_tokens(data['input_ids']) * self.llm.config.scale_emb
+        else:
+            vllm_embedding = self.llm.model.embed_tokens(data['input_ids'])
+
+        vision_hidden_states = [i.type(vllm_embedding.dtype) if isinstance(
+            i, torch.Tensor) else i for i in vision_hidden_states]
+
+        bs = len(data['input_ids'])
+        for i in range(bs):
+            cur_vs_hs = vision_hidden_states[i]
+            
+            if len(cur_vs_hs) > 0:
+                
+                cur_vllm_emb = vllm_embedding[i]
+                
+                cur_image_bound = data['image_bound'][i]
+                
+                if len(cur_image_bound) > 0:
+                    
+                    image_indices = torch.stack(
+                        [torch.arange(r[0], r[1], dtype=torch.long) for r in cur_image_bound]
+                    ).to(vllm_embedding.device)
+                    
+                    cur_vllm_emb.scatter_(
+                        0, 
+                        image_indices.view(-1, 1).repeat(1, cur_vllm_emb.shape[-1]),
+                        cur_vs_hs.view(-1, cur_vs_hs.shape[-1])
+                    )
+
+        return vllm_embedding, vision_hidden_states
+
+    def forward(self, data, **kwargs):
+        vllm_embedding, vision_hidden_states = self.vpm_forward(data)
+
+        output = self.llm(
+            inputs_embeds=vllm_embedding,
+            attention_mask=data["attention_mask"],
+            return_dict=True
+        )
+
+        return CausalVLMOutput(
+            logits=output.logits,
+            hidden_states=output.hidden_states,
+            vision_hidden_states=vision_hidden_states
+        )
+    
+    def generate(self, data, **kwargs):
+        vllm_embedding, vision_hidden_states = self.vpm_forward(data)
+
+        # position_ids = torch.arange(data["input_ids"].size(1), dtype=torch.long).to(data["input_ids"].device)
+        # position_ids = position_ids.unsqueeze(0).expand_as(data["input_ids"])
+
+        # 使用attention_mask将填充位置的position_ids设置为0
+        # position_ids = position_ids * data["attention_mask"]
+        output = self.llm.generate(
+            inputs_embeds=vllm_embedding,
+            # position_ids=position_ids,
+            attention_mask=data["attention_mask"],
+            **kwargs
+        )
+
+        return output
+
+