Upload 4 files

modeling_colinternvl2.py

@@ -0,0 +1,115 @@
+from typing import ClassVar, List, Optional
+from typing import Any, List, Optional, Tuple, Union
+import torch
+from torch import nn
+from .modeling_internvl_chat import InternVLChatModel, InternVLChatConfig
+import math
+
+class ColInternVL2(InternVLChatModel):
+    """
+    ColInternVL2 model implementation from the "ColPali: Efficient Document Retrieval with Vision Language Models" paper.
+    """
+
+    # main_input_name: ClassVar[str] = "doc_input_ids"  # transformers-related
+
+    def __init__(self, config: InternVLChatConfig):
+        super().__init__(config=config)
+        self.dim = 128
+        self.custom_text_proj = nn.Linear(self.language_model.model.config.hidden_size, self.dim ) #, bias=False)
+        self.padding_side = "left"
+        self.img_context_token_id = 151648
+        # self.post_init()
+        self.init_linear()
+    
+    def init_linear(self): 
+        print(self.language_model.model.embed_tokens.weight)
+        stdv = 1. / math.sqrt(self.custom_text_proj.weight.size(1))
+        self.custom_text_proj.weight.data = self.custom_text_proj.weight.data.uniform_(-stdv, stdv)
+        if self.custom_text_proj.bias is not None:
+            self.custom_text_proj.bias.data = self.custom_text_proj.bias.data.uniform_(-stdv, stdv)
+
+
+    def forward(
+            self,
+            pixel_values: torch.FloatTensor = None,
+            input_ids: torch.LongTensor = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            image_flags: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+            statistics: Optional[torch.LongTensor] = None,
+            loss_weight: Optional[List] = None,
+            loss_reduction_all_gather: Optional[bool] = False,
+            **kwargs
+    ) -> torch.Tensor:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+
+        input_embeds = self.language_model.get_input_embeddings()(input_ids).clone()
+        B, N, C = input_embeds.shape
+        
+        if pixel_values is not None:
+            
+            pixel_values = pixel_values.type(self.vision_model.embeddings.patch_embedding.weight.dtype)
+            vit_embeds = self.extract_feature(pixel_values)
+            # image_flags = image_flags.squeeze(-1)
+            # vit_embeds = vit_embeds[image_flags == 1]
+            vit_batch_size = pixel_values.shape[0]
+            
+            input_embeds = input_embeds.reshape(B * N, C)
+            
+            if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:
+                print(f'dynamic ViT batch size: {vit_batch_size}, images per sample: {vit_batch_size / B}, dynamic token length: {N}')
+                if statistics is not None:
+                    num_samples, num_padding_tokens, num_padding_images = statistics.tolist()
+                    self.num_samples += num_samples
+                    print(f'total_samples={self.num_samples}, {num_samples=}, {num_padding_tokens=}, {num_padding_images=}')
+
+            input_ids = input_ids.reshape(B * N)
+            selected = (input_ids == self.img_context_token_id)
+            try:
+                input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)
+                ignore_flag = False
+            except Exception as e:
+                
+                vit_embeds = vit_embeds.reshape(-1, C)
+                print(f'warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, '
+                      f'vit_embeds.shape={vit_embeds.shape}')
+                n_token = selected.sum()
+                input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds[:n_token]
+                ignore_flag = True
+    
+            input_embeds = input_embeds.reshape(B, N, C)
+        
+        outputs = self.language_model.model(
+            inputs_embeds=input_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=True,
+            return_dict=return_dict,
+        )
+        
+        last_hidden_states = outputs[0].type(self.custom_text_proj.weight.dtype)
+        proj = self.custom_text_proj(last_hidden_states)  # (batch_size, sequence_length, dim)
+
+        # L2 normalization
+        proj = proj / proj.norm(dim=-1, keepdim=True)  # (batch_size, sequence_length, dim)
+        proj = proj * attention_mask.unsqueeze(-1)  # (batch_size, sequence_length, dim)
+        return proj
+
+    
+    @property
+    def get_patch_size(self) -> int:
+        return self.visual.config.patch_size
+
+    @property
+    def spatial_merge_size(self) -> int:
+        return self.visual.config.spatial_merge_size

processing_colinternvl2.py

@@ -0,0 +1,211 @@
+import math
+from typing import ClassVar, List, Optional, Tuple, Union
+
+import torch
+from PIL import Image
+from transformers import BatchFeature
+
+from .processing_utils import BaseVisualRetrieverProcessor
+import numpy as np
+import torch
+import torchvision.transforms as T
+from decord import VideoReader, cpu
+from PIL import Image
+from torchvision.transforms.functional import InterpolationMode
+from transformers import AutoModel, AutoTokenizer
+from .conversation import get_conv_template
+from transformers import BatchFeature, ProcessorMixin
+
+class ColInternVL2Processor(BaseVisualRetrieverProcessor, ProcessorMixin):
+    """
+    Processor for ColInternVL2.
+    """
+    attributes = [ "tokenizer"]
+    image_processor_class = "InternVL2ImageProcessor"
+    tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
+    
+    def __init__(self, tokenizer,  **kwargs):
+        self.template = "Hermes-2"
+        self.num_image_token = 256        
+        # self.max_num = 6
+        self.max_num = 4
+
+        if isinstance(tokenizer, str):
+            self.tokenizer = AutoTokenizer.from_pretrained(tokenizer, trust_remote_code=True, use_fast=False)
+        else:
+            self.tokenizer = tokenizer
+            
+        self.tokenizer.padding_side = 'left'
+        self.IMAGENET_MEAN = (0.485, 0.456, 0.406)
+        self.IMAGENET_STD = (0.229, 0.224, 0.225)
+        self.IMG_CONTEXT_TOKEN='<IMG_CONTEXT>'
+        self.IMG_START_TOKEN='<img>'
+        self.IMG_END_TOKEN='</img>'
+        self.img_context_token_id = self.tokenizer.convert_tokens_to_ids(self.IMG_CONTEXT_TOKEN)
+        # self.system_message = '你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。'
+        self.system_message = 'Bạn là một mô hình trí tuệ nhân tạo đa phương thức Tiếng Việt có tên gọi là Vintern, được phát triển bởi người Việt. Bạn là một trợ lý trí tuệ nhân tạo hữu ích và không gây hại.'
+        super().__init__(tokenizer)
+        
+    # def from_pretrained(pretrained_model_name_or_path, template="Hermes-2", **kwargs):
+    #     return ColInternVL2Processor(pretrained_model_name_or_path, template=template, **kwargs)
+        
+    def build_transform(self, input_size):
+        MEAN, STD = self.IMAGENET_MEAN, self.IMAGENET_STD
+        transform = T.Compose([
+            T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
+            T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
+            T.ToTensor(),
+            T.Normalize(mean=MEAN, std=STD)
+        ])
+        return transform
+    
+    def find_closest_aspect_ratio(self, aspect_ratio, target_ratios, width, height, image_size):
+        best_ratio_diff = float('inf')
+        best_ratio = (1, 1)
+        area = width * height
+        for ratio in target_ratios:
+            target_aspect_ratio = ratio[0] / ratio[1]
+            ratio_diff = abs(aspect_ratio - target_aspect_ratio)
+            if ratio_diff < best_ratio_diff:
+                best_ratio_diff = ratio_diff
+                best_ratio = ratio
+            elif ratio_diff == best_ratio_diff:
+                if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
+                    best_ratio = ratio
+        return best_ratio
+    
+    def dynamic_preprocess(self, image, min_num=1, max_num=12, image_size=448, use_thumbnail=False):
+        orig_width, orig_height = image.size
+        aspect_ratio = orig_width / orig_height
+    
+        # calculate the existing image aspect ratio
+        target_ratios = set(
+            (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
+            i * j <= max_num and i * j >= min_num)
+        target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+    
+        # find the closest aspect ratio to the target
+        target_aspect_ratio = self.find_closest_aspect_ratio(
+            aspect_ratio, target_ratios, orig_width, orig_height, image_size)
+    
+        # calculate the target width and height
+        target_width = image_size * target_aspect_ratio[0]
+        target_height = image_size * target_aspect_ratio[1]
+        blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
+    
+        # resize the image
+        resized_img = image.resize((target_width, target_height))
+        processed_images = []
+        for i in range(blocks):
+            box = (
+                (i % (target_width // image_size)) * image_size,
+                (i // (target_width // image_size)) * image_size,
+                ((i % (target_width // image_size)) + 1) * image_size,
+                ((i // (target_width // image_size)) + 1) * image_size
+            )
+            # split the image
+            split_img = resized_img.crop(box)
+            processed_images.append(split_img)
+        assert len(processed_images) == blocks
+        if use_thumbnail and len(processed_images) != 1:
+            thumbnail_img = image.resize((image_size, image_size))
+            processed_images.append(thumbnail_img)
+        return processed_images
+    
+    def load_image(self, image, input_size=448, max_num=12):
+        transform = self.build_transform(input_size=input_size)
+        images = self.dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
+        pixel_values = [transform(image) for image in images]
+        pixel_values = torch.stack(pixel_values)
+        return pixel_values
+
+    
+    def process_images(
+        self,
+        images: List[Image.Image],
+    ) -> BatchFeature:
+        """
+        Process images for InternVl2.
+        """
+
+        pixel_values = [ self.load_image(image, max_num=self.max_num) for image in images]
+
+        num_patches_list = [ pixel_.size(0) for pixel_ in pixel_values]
+        image_flags = [ torch.tensor([1] * pixel_.shape[0], dtype=torch.long) for pixel_ in pixel_values ]
+        
+        queries = []
+        for idx, num_patches in enumerate(num_patches_list):
+            question = "<image>\nDescribe the image."
+
+            template = get_conv_template(self.template)
+            template.system_message = self.system_message
+            template.append_message(template.roles[0], question)
+            template.append_message(template.roles[1], None)
+            query = template.get_prompt()
+            image_tokens = self.IMG_START_TOKEN + self.IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + self.IMG_END_TOKEN
+            query = query.replace('<image>', image_tokens, 1)
+            queries.append(query)
+
+        model_inputs = self.tokenizer(queries, return_tensors='pt', padding=True)
+        input_ids = model_inputs['input_ids'] #.to(self.device)
+        attention_mask = model_inputs['attention_mask'] #.to(self.device)
+        pixel_values = torch.cat(pixel_values)
+        
+        batch_doc = BatchFeature({
+            "pixel_values" : pixel_values,
+            "input_ids" : input_ids,
+            "attention_mask" : attention_mask,
+            # "image_flags" : image_flags
+        })
+        return batch_doc    
+    
+    def process_queries(
+        self,
+        queries: List[str],
+        max_length: int = 100,
+        suffix: Optional[str] = None,
+    ) -> BatchFeature:
+        """
+        Process queries for InternVl2.
+        """
+
+        texts_query: List[str] = []
+
+        for query in queries:
+            query = f"Query: {query}"
+            template = get_conv_template(self.template)
+            template.system_message = self.system_message
+            template.append_message(template.roles[0], query)
+            template.append_message(template.roles[1], None)
+            query = template.get_prompt()
+            texts_query.append(query)
+            
+        model_inputs = self.tokenizer(texts_query, return_tensors='pt', max_length=max_length, padding="longest")
+        input_ids = model_inputs['input_ids'] #.to(self.device)
+        attention_mask = model_inputs['attention_mask'] #.to(self.device)
+        
+        batch_query = BatchFeature({
+            "pixel_values" : None,
+            "input_ids" : input_ids,
+            "attention_mask" : attention_mask,
+        })
+        return batch_query
+
+    def score(
+        self,
+        qs: List[torch.Tensor],
+        ps: List[torch.Tensor],
+        device: Optional[Union[str, torch.device]] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Compute the MaxSim score (ColBERT-like) for the given multi-vector query and passage embeddings.
+        """
+        return self.score_multi_vector(qs, ps, device=device, **kwargs)
+
+    def get_n_patches(
+        self,
+        image_size: Tuple[int, int],
+        patch_size: int,
+    ) -> Tuple[int, int]:
+        raise NotImplementedError("This method is not implemented for ColInternVL2.")

processing_utils.py

@@ -0,0 +1,117 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional, Tuple, Union
+
+import torch
+from PIL import Image
+from transformers import BatchEncoding, BatchFeature
+
+from .torch_utils import get_torch_device
+
+
+class BaseVisualRetrieverProcessor(ABC):
+    """
+    Base class for visual retriever processors.
+    """
+
+    @abstractmethod
+    def process_images(
+        self,
+        images: List[Image.Image],
+    ) -> Union[BatchFeature, BatchEncoding]:
+        pass
+
+    @abstractmethod
+    def process_queries(
+        self,
+        queries: List[str],
+        max_length: int = 50,
+        suffix: Optional[str] = None,
+    ) -> Union[BatchFeature, BatchEncoding]:
+        pass
+
+    @abstractmethod
+    def score(
+        self,
+        qs: List[torch.Tensor],
+        ps: List[torch.Tensor],
+        device: Optional[Union[str, torch.device]] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        pass
+
+    @staticmethod
+    def score_single_vector(
+        qs: List[torch.Tensor],
+        ps: List[torch.Tensor],
+        device: Optional[Union[str, torch.device]] = None,
+    ) -> torch.Tensor:
+        """
+        Compute the dot product score for the given single-vector query and passage embeddings.
+        """
+        device = device or get_torch_device("auto")
+
+        if len(qs) == 0:
+            raise ValueError("No queries provided")
+        if len(ps) == 0:
+            raise ValueError("No passages provided")
+
+        qs_stacked = torch.stack(qs).to(device)
+        ps_stacked = torch.stack(ps).to(device)
+
+        scores = torch.einsum("bd,cd->bc", qs_stacked, ps_stacked)
+        assert scores.shape[0] == len(qs), f"Expected {len(qs)} scores, got {scores.shape[0]}"
+
+        scores = scores.to(torch.float32)
+        return scores
+
+    @staticmethod
+    def score_multi_vector(
+        qs: List[torch.Tensor],
+        ps: List[torch.Tensor],
+        batch_size: int = 128,
+        device: Optional[Union[str, torch.device]] = None,
+    ) -> torch.Tensor:
+        """
+        Compute the MaxSim score (ColBERT-like) for the given multi-vector query and passage embeddings.
+        """
+        device = device or get_torch_device("auto")
+
+        if len(qs) == 0:
+            raise ValueError("No queries provided")
+        if len(ps) == 0:
+            raise ValueError("No passages provided")
+
+        scores_list: List[torch.Tensor] = []
+
+        for i in range(0, len(qs), batch_size):
+            scores_batch = []
+            qs_batch = torch.nn.utils.rnn.pad_sequence(qs[i : i + batch_size], batch_first=True, padding_value=0).to(
+                device
+            )
+            for j in range(0, len(ps), batch_size):
+                ps_batch = torch.nn.utils.rnn.pad_sequence(
+                    ps[j : j + batch_size], batch_first=True, padding_value=0
+                ).to(device)
+                scores_batch.append(torch.einsum("bnd,csd->bcns", qs_batch, ps_batch).max(dim=3)[0].sum(dim=2))
+            scores_batch = torch.cat(scores_batch, dim=1).cpu()
+            scores_list.append(scores_batch)
+
+        scores = torch.cat(scores_list, dim=0)
+        assert scores.shape[0] == len(qs), f"Expected {len(qs)} scores, got {scores.shape[0]}"
+
+        scores = scores.to(torch.float32)
+        return scores
+
+    @abstractmethod
+    def get_n_patches(
+        self,
+        image_size: Tuple[int, int],
+        patch_size: int = 14,
+        *args,
+        **kwargs,
+    ) -> Tuple[int, int]:
+        """
+        Get the number of patches (n_patches_x, n_patches_y) that will be used to process an
+        image of size (height, width) with the given patch size.
+        """
+        pass

torch_utils.py

@@ -0,0 +1,52 @@
+import gc
+import logging
+from typing import List, TypeVar
+
+import torch
+from torch.utils.data import Dataset
+
+logger = logging.getLogger(__name__)
+T = TypeVar("T")
+
+
+def get_torch_device(device: str = "auto") -> str:
+    """
+    Returns the device (string) to be used by PyTorch.
+
+    `device` arg defaults to "auto" which will use:
+    - "cuda:0" if available
+    - else "mps" if available
+    - else "cpu".
+    """
+
+    if device == "auto":
+        if torch.cuda.is_available():
+            device = "cuda:0"
+        elif torch.backends.mps.is_available():  # for Apple Silicon
+            device = "mps"
+        else:
+            device = "cpu"
+        logger.info(f"Using device: {device}")
+
+    return device
+
+
+def tear_down_torch():
+    """
+    Teardown for PyTorch.
+    Clears GPU cache for both CUDA and MPS.
+    """
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.mps.empty_cache()
+
+
+class ListDataset(Dataset[T]):
+    def __init__(self, elements: List[T]):
+        self.elements = elements
+
+    def __len__(self) -> int:
+        return len(self.elements)
+
+    def __getitem__(self, idx: int) -> T:
+        return self.elements[idx]