modeling_minicpmv.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright @2023 AI, ZHIHU Inc. (zhihu.com)
#
# @author: wangchongyi <wangchongyi@zhihu.com>
# @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