Create modeling_minicpmv.py

modeling_minicpmv.py

@@ -1,22 +1,21 @@
 import math
 from typing import List, Optional
 import json
+import timm
 import torch
 import torchvision
-from threading import Thread
-from copy import deepcopy
+import deepspeed
 from PIL import Image
+from timm.data import IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD
 from torchvision import transforms
-from transformers import LlamaTokenizer, LlamaPreTrainedModel, LlamaForCausalLM, AutoModel, PreTrainedTokenizerFast, TextIteratorStreamer
-from transformers.models.idefics2.modeling_idefics2 import Idefics2VisionTransformer
-
+from transformers import LlamaTokenizer
+from transformers.integrations import is_deepspeed_zero3_enabled
 from .configuration_minicpm import MiniCPMVConfig
+from .modeling_minicpm import MiniCPMForCausalLM, MiniCPMPreTrainedModel
 from .resampler import Resampler
 
-IMAGENET_INCEPTION_MEAN = (0.5, 0.5, 0.5) # timm.data.IMAGENET_INCEPTION_MEAN
-IMAGENET_INCEPTION_STD = (0.5, 0.5, 0.5)  # timm.data.IMAGENET_INCEPTION_STD
 
-class MiniCPMVPreTrainedModel(LlamaPreTrainedModel):
+class MiniCPMVPreTrainedModel(MiniCPMPreTrainedModel):
     config_class = MiniCPMVConfig
 
 
@@ -24,7 +23,7 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
 
-        self.llm = LlamaForCausalLM(config)
+        self.llm = MiniCPMForCausalLM(config)
         self.vpm = self.init_vision_module()
         self.vision_dim = self.vpm.embed_dim
         self.embed_dim = self.llm.config.hidden_size
@@ -32,19 +31,26 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
         self.transform = self.init_transform()
 
     def init_vision_module(self):
-        # same as HuggingFaceM4/siglip-so400m-14-980-flash-attn2-navit
-        model = Idefics2VisionTransformer(self.config.vision_config)
-        if self.config.drop_vision_last_layer:
-            model.encoder.layers = model.encoder.layers[:-1]
+        model = timm.create_model(
+            self.config.vision_encoder,
+            pretrained=False,
+            num_classes=0,
+            dynamic_img_size=True,
+            dynamic_img_pad=True
+        )
 
-        setattr(model, 'embed_dim', model.embeddings.embed_dim)
-        setattr(model, 'patch_size', model.embeddings.patch_size)
+        if isinstance(model, timm.models.VisionTransformer):
+            if model.attn_pool is not None:
+                model.attn_pool = torch.nn.Identity()
+
+        if self.config.drop_vision_last_layer:
+            model.blocks = model.blocks[:-1]
 
         return model
 
     def init_resampler(self, embed_dim, vision_dim):
         return Resampler(
-            num_queries=self.config.query_num,
+            grid_size=int(math.sqrt(self.config.query_num)),
             embed_dim=embed_dim,
             num_heads=embed_dim // 128,
             kv_dim=vision_dim,
@@ -67,94 +73,75 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
     def set_input_embeddings(self, value):
         self.llm.embed_tokens = value
 
+    def get_vision_embedding(self, pixel_values):
+        res = []
+        dtype = self.llm.lm_head.weight.dtype
+        def process_each_pixel(pixel_value, dtype, config, vpm, resampler):
+            H, W = pixel_value.shape[-2:]
+            target_size = (math.ceil(H / config.patch_size), math.ceil(W / config.patch_size))
+            vision_embedding = self.vpm.forward_features(pixel_value.unsqueeze(0).type(dtype))
+            if hasattr(vpm, 'num_prefix_tokens') and vpm.num_prefix_tokens > 0:
+                vision_embedding = vision_embedding[:, vpm.num_prefix_tokens:]
+            return resampler(vision_embedding, target_size)
+
+        if is_deepspeed_zero3_enabled():
+            with deepspeed.zero.GatheredParameters(self.vpm.pos_embed, modifier_rank=0):
+                for pixel_value in pixel_values:
+                    result = process_each_pixel(pixel_value, dtype, self.config, self.vpm, self.resampler)
+                    res.append(result)
+        else:
+            for pixel_value in pixel_values:
+                result = process_each_pixel(pixel_value, dtype, self.config, self.vpm, self.resampler)
+                res.append(result)
+        return torch.vstack(res)
+
     def get_vllm_embedding(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
-            tgt_sizes = data['tgt_sizes']
-            pixel_values_list = data['pixel_values']
+        if "vision_hidden_states" not in data:
+            pixel_values_list = data["pixel_values"]
             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])
-
-            # exist image
-            if all_pixel_values:
-                tgt_sizes = torch.vstack(tgt_sizes).type(torch.int32)
-
-                if self.config.batch_vision_input:
-                    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)
-                    B, L, _ = all_pixel_values.shape
-                    all_pixel_values = all_pixel_values.permute(0, 2, 1).reshape(B, 3, -1, 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)
-                else:
-                    # get vision_embedding foreach
-                    vision_embedding = []
-                    for single_tgt_size, single_pixel_values in zip(tgt_sizes, all_pixel_values):
-                        single_pixel_values = single_pixel_values.unsqueeze(0)
-                        B, L, _ = single_pixel_values.shape
-                        single_pixel_values = single_pixel_values.permute(0, 2, 1).reshape(B, 3, -1, L)
-                        single_vision_embedding = self.vpm(single_pixel_values.type(dtype)).last_hidden_state
-                        single_vision_embedding = self.resampler(single_vision_embedding, single_tgt_size.unsqueeze(0))
-                        vision_embedding.append(single_vision_embedding)
-                    vision_embedding = torch.vstack(vision_embedding)
-
-                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:
+                if len(pixel_values) > 0:
+                    vision_hidden_states.append(self.get_vision_embedding(pixel_values))
+                elif self.training:
+                    dtype = self.llm.lm_head.weight.dtype
+                    device = self.llm.lm_head.weight.device
                     dummy_image = torch.zeros(
-                        (1, 3, 224, 224),
-                        device=device, dtype=dtype
+                        (1, 3, 224, 224), device=device, dtype=dtype
                     )
-                    tgt_sizes = torch.Tensor([[(224 // self.config.patch_size), math.ceil(224 / self.config.patch_size)]]).type(torch.int32)
-                    dummy_feature = self.resampler(self.vpm(dummy_image).last_hidden_state, tgt_sizes)
+                    vision_hidden_states.append(self.get_vision_embedding(dummy_image))
                 else:
-                    dummy_feature = []
-                for _ in range(len(pixel_values_list)):
-                    vision_hidden_states.append(dummy_feature)
-
-        else:
-            vision_hidden_states = data['vision_hidden_states']
+                    vision_hidden_states.append([])
 
-        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 = data["vision_hidden_states"]
 
-        vision_hidden_states = [i.type(vllm_embedding.dtype) if isinstance(
-            i, torch.Tensor) else i for i in vision_hidden_states]
+        vllm_embedding = (
+            self.llm.model.embed_tokens(data["input_ids"]) * self.llm.config.scale_emb
+        )
+        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'])
+        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]
+                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]
+                        [
+                            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]))
+                    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]),
+                    )
                 elif self.training:
                     cur_vllm_emb += cur_vs_hs[0].mean() * 0
 
@@ -174,8 +161,12 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
         )
 
     def _convert_to_tensors(
-        self, tokenizer, input_ids, max_inp_length: Optional[int] = None
+        self, tokenizer, input_str, max_inp_length: Optional[int] = None
     ):
+        if tokenizer.add_bos_token:
+            input_ids = tokenizer.encode(input_str)
+        else:
+            input_ids = [tokenizer.bos_id] + tokenizer.encode(input_str)
         if max_inp_length is not None:
             input_ids = input_ids[:max_inp_length]
         input_ids = torch.tensor(input_ids, dtype=torch.int32)
@@ -199,13 +190,13 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
         return model_input
 
     def _process_list(
-        self, tokenizer, input_id_list, max_inp_length: Optional[int] = None
+        self, tokenizer, data_list: List[str], max_inp_length: Optional[int] = None
     ):
         pad_keys = ["input_ids"]
         input_tensors = []
-        for input_ids in input_id_list:
+        for data in data_list:
             input_tensors.append(
-                self._convert_to_tensors(tokenizer, input_ids, max_inp_length)
+                self._convert_to_tensors(tokenizer, data, max_inp_length)
             )
         padded = {}
         for key in pad_keys:
@@ -214,36 +205,13 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
         return padded
 
     def _decode(self, inputs_embeds, tokenizer, **kwargs):
-        terminators = [
-            tokenizer.eos_token_id,
-            tokenizer.convert_tokens_to_ids("<|eot_id|>")
-        ]
         output = self.llm.generate(
             inputs_embeds=inputs_embeds,
             pad_token_id=0,
-            eos_token_id=terminators,
+            eos_token_id=tokenizer.eos_token_id,
             **kwargs
         )
         return self._decode_text(output, tokenizer)
-    
-    def _decode_stream(self, inputs_embeds, tokenizer, **kwargs):
-        terminators = [
-            tokenizer.eos_token_id,
-            tokenizer.convert_tokens_to_ids("<|eot_id|>")
-        ]
-        streamer = TextIteratorStreamer(tokenizer=tokenizer)
-        generation_kwargs = {
-            'inputs_embeds': inputs_embeds,
-            'pad_token_id': 0,
-            'eos_token_id': terminators,
-            'streamer': streamer
-        }
-        generation_kwargs.update(kwargs)
-
-        thread = Thread(target=self.llm.generate, kwargs=generation_kwargs)
-        thread.start()
-    
-        return streamer
 
     def _decode_text(self, result_ids, tokenizer):
         result_text = []
@@ -251,7 +219,7 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
             result = result[result != 0]
             if result[0] == tokenizer.bos_id:
                 result = result[1:]
-            if result[-1] == tokenizer.eos_id or result[-1] == tokenizer.eot_id:
+            if result[-1] == tokenizer.eos_id:
                 result = result[:-1]
             result_text.append(tokenizer.decode(result).strip())
         return result_text
@@ -259,9 +227,9 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
     def slice_image(self, image):
         return slice_image(
             image,
-            self.config.slice_config.max_slice_nums,
-            self.config.slice_config.scale_resolution,
-            self.config.slice_config.patch_size,
+            self.config.max_slice_nums,
+            self.config.scale_resolution,
+            self.config.patch_size,
         )
 
     def get_slice_image_placeholder(self, image, tokenizer):
@@ -275,9 +243,9 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
 
         source_image, patches, best_grid = slice_image(
             image,
-            self.config.slice_config.max_slice_nums,
-            self.config.slice_config.scale_resolution,
-            self.config.slice_config.patch_size,
+            self.config.max_slice_nums,
+            self.config.scale_resolution,
+            self.config.patch_size,
         )
 
         slice_images.append(source_image)
@@ -294,56 +262,36 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
 
         return slice_images, final_placeholder
 
-    def reshape_by_patch(self, image_tensor):
-        """
-        :param image_tensor: shape [3, H, W]
-        :param patch_size:
-        :return: [3, patch_size, HW/patch_size]
-        """
-        patch_size = self.config.patch_size
-        patches = torch.nn.functional.unfold(
-            image_tensor,
-            (patch_size, patch_size),
-            stride=(patch_size, patch_size)
-        )
-
-        patches = patches.reshape(image_tensor.size(0), patch_size, patch_size, -1)
-        patches = patches.permute(0, 1, 3, 2).reshape(image_tensor.size(0), patch_size, -1)
-        return patches
-
     def generate(
         self,
-        input_id_list=None,
+        data_list=None,
         img_list=None,
-        tgt_sizes=None,
         tokenizer=None,
         max_inp_length: Optional[int] = None,
         vision_hidden_states=None,
         return_vision_hidden_states=False,
-        stream=False,
         **kwargs
     ):
 
-        assert input_id_list is not None
-        bs = len(input_id_list)
+        assert data_list is not None
+        bs = len(data_list)
         if img_list == None:
             img_list = [[] for i in range(bs)]
         assert bs == len(img_list)
 
-        model_inputs = self._process_list(tokenizer, input_id_list, max_inp_length)
+        model_inputs = self._process_list(tokenizer, data_list, max_inp_length)
 
         if vision_hidden_states is None:
             pixel_values = []
             for i in range(bs):
                 img_inps = []
                 for img in img_list[i]:
-                    img_inps.append(img.to(self.device))
+                    img_inps.append(self.transform(img).to(self.device))
                 if img_inps:
                     pixel_values.append(img_inps)
                 else:
                     pixel_values.append([])
             model_inputs["pixel_values"] = pixel_values
-            model_inputs['tgt_sizes'] = tgt_sizes
         else:
             model_inputs["vision_hidden_states"] = vision_hidden_states
 
@@ -353,10 +301,7 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
                 vision_hidden_states,
             ) = self.get_vllm_embedding(model_inputs)
 
-            if stream:
-                result = self._decode_stream(model_inputs["inputs_embeds"], tokenizer, **kwargs)
-            else:
-                result = self._decode(model_inputs["inputs_embeds"], tokenizer, **kwargs)
+            result = self._decode(model_inputs["inputs_embeds"], tokenizer, **kwargs)
 
         if return_vision_hidden_states:
             return result, vision_hidden_states
@@ -367,70 +312,42 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
         self,
         image,
         msgs,
+        context,
         tokenizer,
         vision_hidden_states=None,
         max_new_tokens=1024,
         sampling=True,
         max_inp_length=2048,
-        system_prompt='',
-        stream=False,
         **kwargs
     ):
         if isinstance(msgs, str):
             msgs = json.loads(msgs)
-
-        copy_msgs = deepcopy(msgs)
-        assert len(copy_msgs) > 0, 'msgs is empty'
-        assert sampling or not stream, 'if use stream mode, make sure sampling=True'
-
-        if image is not None and isinstance(copy_msgs[0]['content'], str):
-            copy_msgs[0]['content'] = [image, copy_msgs[0]['content']]
-
-        images = []
-        tgt_sizes = []
-        for i, msg in enumerate(copy_msgs):
+        # msgs to prompt
+        prompt = ""
+        for i, msg in enumerate(msgs):
             role = msg["role"]
             content = msg["content"]
             assert role in ["user", "assistant"]
             if i == 0:
                 assert role == "user", "The role of first msg should be user"
-            if isinstance(content, str):
-                content = [content]
-
-            cur_msgs = []
-            for c in content:
-                if isinstance(c, Image.Image):
-                    image = c
-                    if self.config.slice_mode:
-                        slice_images, image_placeholder = self.get_slice_image_placeholder(
-                            image, tokenizer
-                        )
-                        cur_msgs.append(image_placeholder)
-                        for slice_image in slice_images:
-                            slice_image = self.transform(slice_image)
-                            H, W = slice_image.shape[1:]
-                            images.append(self.reshape_by_patch(slice_image))
-                            tgt_sizes.append(torch.Tensor([H // self.config.patch_size, W // self.config.patch_size]).type(torch.int32))
-                    else:
-                        images.append(self.transform(image))
-                        cur_msgs.append(
-                            tokenizer.im_start
-                            + tokenizer.unk_token * self.config.query_num
-                            + tokenizer.im_end
-                        )
-                elif isinstance(c, str):
-                    cur_msgs.append(c)
-
-
-            msg['content'] = '\n'.join(cur_msgs)
-        if tgt_sizes:
-            tgt_sizes = torch.vstack(tgt_sizes)
-
-        if system_prompt:
-            sys_msg = {'role': 'system', 'content': system_prompt}
-            copy_msgs = [sys_msg] + copy_msgs
-
-        input_ids = tokenizer.apply_chat_template(copy_msgs, tokenize=True, add_generation_prompt=False)
+                if self.config.slice_mode:
+                    images, final_placeholder = self.get_slice_image_placeholder(
+                        image, tokenizer
+                    )
+                    content = final_placeholder + "\n" + content
+                else:
+                    images = [image]
+                    content = (
+                        tokenizer.im_start
+                        + tokenizer.unk_token * self.config.query_num
+                        + tokenizer.im_end
+                        + "\n"
+                        + content
+                    )
+            prompt += "<用户>" if role == "user" else "<AI>"
+            prompt += content
+        prompt += "<AI>"
+        final_input = prompt
 
         if sampling:
             generation_config = {
@@ -452,34 +369,25 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
 
         with torch.inference_mode():
             res, vision_hidden_states = self.generate(
-                input_id_list=[input_ids],
+                data_list=[final_input],
                 max_inp_length=max_inp_length,
                 img_list=[images],
-                tgt_sizes=[tgt_sizes],
                 tokenizer=tokenizer,
                 max_new_tokens=max_new_tokens,
                 vision_hidden_states=vision_hidden_states,
                 return_vision_hidden_states=True,
-                stream=stream,
                 **generation_config
             )
+        answer = res[0]
+        context = msgs.copy()
+        context.append({"role": "assistant", "content": answer})
 
-        if stream:
-            def stream_gen():
-                for text in res:
-                    text = text.replace(tokenizer.eot_token, '').replace(tokenizer.eos_token, '')
-                    yield text
-            return stream_gen()
-
-        else:
-            answer = res[0]
-            return answer
+        return answer, context, generation_config
 
 
-class PreTrainedTokenizerFastWrapper(PreTrainedTokenizerFast):
+class LlamaTokenizerWrapper(LlamaTokenizer):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        self.eot_token = "<|eot_id|>"
         self.im_start = "<image>"
         self.im_end = "</image>"
         self.ref_start = "<ref>"
@@ -488,40 +396,30 @@ class PreTrainedTokenizerFastWrapper(PreTrainedTokenizerFast):
         self.box_end = "</box>"
         self.quad_start = "<quad>"
         self.quad_end = "</quad>"
+        self.point_start = "<point>"
+        self.point_end = "</point>"
         self.slice_start = "<slice>"
         self.slice_end = "</slice>"
 
     @property
     def eos_id(self):
-        return self.eos_token_id
+        return self.sp_model.eos_id()
 
     @property
     def bos_id(self):
-        return self.bos_token_id
+        return self.sp_model.bos_id()
 
     @property
     def unk_id(self):
-        return self.unk_token_id
-
-    @property
-    def eot_id(self):
-        return self.convert_tokens_to_ids(self.eot_token)
+        return self.sp_model.unk_id()
 
     @property
     def im_start_id(self):
-        return self.convert_tokens_to_ids(self.im_start)
+        return self._convert_token_to_id(self.im_start)
 
     @property
     def im_end_id(self):
-        return self.convert_tokens_to_ids(self.im_end)
-
-    @staticmethod
-    def escape(text: str) -> str:
-        return text
-
-    @staticmethod
-    def unescape(text: str) -> str:
-        return text
+        return self._convert_token_to_id(self.im_end)
 
 
 def pad(orig_items, key, max_length=None, padding_value=0, padding_side="left"):