Create processing_colgranitevision.py

processing_colgranitevision.py

@@ -0,0 +1,396 @@
+import math
+from typing import ClassVar, List, Optional, Tuple, Union
+
+import torch
+from PIL import Image, ImageOps
+from transformers import BatchFeature, LlavaNextProcessor
+
+
+def round_by_factor(number: float, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+
+def ceil_by_factor(number: float, factor: int) -> int:
+    """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+    return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: float, factor: int) -> int:
+    """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+    return math.floor(number / factor) * factor
+
+
+class ColGraniteVisionProcessor(LlavaNextProcessor):
+    """
+    Processor for ColPali.
+    """
+
+    visual_prompt_prefix: ClassVar[str] = "<|user|>\n<image>\nDescribe the image.\n"
+    system_message: ClassVar[
+        str] = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions."
+    query_prefix: ClassVar[str] = "Query: "
+    query_start: ClassVar[str] = "<|user|>\n"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.factor = 14
+        self.min_size = 384
+        self.max_size = 384 * 2
+        self.suffix_len = 10
+        self.patch_size = 14
+
+    @property
+    def query_augmentation_token(self) -> str:
+        """
+        Return the query augmentation token.
+        Query augmentation buffers are used as reasoning buffers during inference.
+        """
+        return self.tokenizer.pad_token
+
+    @staticmethod
+    def smart_resize_helper(
+            width: int,
+            height: int,
+            factor: int,
+            min_size: int,
+            max_size: int
+    ) -> Tuple[int, int]:
+        """
+        Returns the resized image dimensions such that:
+        1. The smaller dimension is set to 'min_size'.
+        2. The larger dimension is scaled proportionally to maintain aspect ratio.
+        3. If the larger dimension exceeds 'max_size', it is clipped to 'max_size',
+        and the smaller dimension is adjusted accordingly to maintain aspect ratio.
+        4. Both dimensions are divisible by 'factor'.
+        """
+
+        # Determine scale factor based on min_size
+        if height < width:
+            scale_factor = min_size / height
+        else:
+            scale_factor = min_size / width
+
+        new_width = round(width * scale_factor)
+        new_height = round(height * scale_factor)
+
+        # If the longer dimension exceeds max_size, adjust accordingly
+        if max(new_width, new_height) > max_size:
+            clip_factor = max_size / max(new_width, new_height)
+            new_width = round(new_width * clip_factor)
+            new_height = round(new_height * clip_factor)
+
+        # Ensure dimensions are divisible by factor
+        # new_width = round_by_factor(new_width, factor)
+        # new_height = round_by_factor(new_height, factor)
+
+        return new_width, new_height
+
+    @staticmethod
+    def pad_image_center(image: Image.Image,
+                         target_width: int,
+                         target_height: int,
+                         fill_color=(0, 0, 0)) -> Image.Image:
+        """
+        Pads the given image to be centered within the target dimensions.
+
+        :param image: PIL Image to be padded.
+        :param target_width: The desired width after padding.
+        :param target_height: The desired height after padding.
+        :param fill_color: Background color (default is black).
+        :return: Padded image with centered content.
+        """
+
+        # Get original image size
+        img_width, img_height = image.size
+
+        # Compute padding values
+        pad_left = (target_width - img_width) // 2
+        pad_top = (target_height - img_height) // 2
+        pad_right = target_width - img_width - pad_left
+        pad_bottom = target_height - img_height - pad_top
+
+        # Apply padding
+        padded_image = ImageOps.expand(image, (pad_left, pad_top, pad_right, pad_bottom), fill_color).convert("RGB")
+
+        return padded_image
+
+    def smart_resize(self, image: Image.Image) -> Image.Image:
+        """
+        Resize and convert the image to the required format.
+        """
+        image_size = image.size
+        resized_height, resized_width = self.smart_resize_helper(
+            width=image_size[0],
+            height=image_size[1],
+            factor=self.factor,
+            min_size=self.min_size,
+            max_size=self.max_size
+        )
+        return image.convert("RGB").resize((resized_width, resized_height))
+
+    def smart_resize_and_pad(self, image: Image.Image) -> Image.Image:
+        """
+        Resize and pad the image to the required format.
+        """
+        return self.resize_and_pad_centered(
+            image=image,
+            factor=self.factor,
+            min_size=self.min_size,
+            max_size=self.max_size,
+            fill_color=0
+        )
+
+    def resize_and_pad_centered(self,
+                                image: Image.Image,
+                                factor: int,
+                                min_size: int,
+                                max_size: int,
+                                fill_color=0
+                                ) -> Image.Image:
+        """
+        Resizes and pads an image such that:
+        - The short side is set to `min_size`.
+        - The long side is scaled proportionally but clipped to `max_size`.
+        - The image is centered within the final padded area.
+
+        :param image: PIL Image
+        :param factor: Factor to make dimensions divisible by
+        :param min_size: Minimum size for the short side
+        :param max_size: Maximum allowed size for the long side
+        :param fill_color: Background padding color (default black)
+        :return: Resized and padded image
+        """
+
+        # Get original size
+        width, height = image.size
+
+        if min_size == -1 or max_size == -1:
+            return image.convert("RGB")
+
+        # Determine scale factor based on the short side (min_size)
+        if width < height:
+            scale_factor = min_size / width
+            target_width = min_size
+            max_scale_factor = min(max_size / height, scale_factor)
+            target_height = round(height * max_scale_factor)
+        else:
+            scale_factor = min_size / height
+            target_height = min_size
+            max_scale_factor = min(max_size / width, scale_factor)
+            target_width = round(width * max_scale_factor)
+
+        # Ensure the longer side does not exceed max_size
+        # if max(target_width, target_height) > max_size:
+        #     clip_factor = max_size / max(target_width, target_height)
+        #     target_width = round(target_width * clip_factor)
+        #     target_height = round(target_height * clip_factor)
+
+        # Ensure dimensions are divisible by factor
+        # target_width = round_by_factor(target_width, factor)
+        # target_height = round_by_factor(target_height, factor)
+
+        # Resize the image
+        resized_image = image.resize((target_width, target_height), Image.LANCZOS)
+
+        # Determine final padded dimensions (aligned to short side)
+        if width < height:
+            final_width, final_height = min_size, max_size
+        else:
+            final_width, final_height = max_size, min_size
+
+        # Compute padding to center the image
+        pad_left = (final_width - target_width) // 2
+        pad_top = (final_height - target_height) // 2
+        pad_right = final_width - target_width - pad_left
+        pad_bottom = final_height - target_height - pad_top
+
+        # Apply centered padding
+        # final_image = ImageOps.expand(resized_image, (pad_left, pad_top, pad_right, pad_bottom), fill_color).convert("RGB")
+        final_image = resized_image.convert("RGB")
+
+        return final_image
+
+    def format_data(self, question, image):
+        return [
+            {
+                "role": "system",
+                "content": [{"type": "text", "text": self.system_message}],
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image",
+                        "image": image,
+                    },
+                    {
+                        "type": "text",
+                        "text": question,
+                    },
+                ],
+            }
+        ]
+
+    def format_data_wo_role(self, question, image=None):
+        return [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image",
+                        "image": image,
+                    },
+                    {
+                        "type": "text",
+                        "text": question,
+                    },
+                ],
+            }
+        ]
+
+    def process_images(
+            self,
+            images: List[Image.Image],
+    ) -> BatchFeature:
+        """
+        Process images for ColPali.
+        """
+        # texts_doc = [self.apply_chat_template(self.format_data_wo_role(self.visual_prompt_prefix, img),tokenize=False ) for img in images]
+        texts_doc = [self.visual_prompt_prefix for _ in images]
+        images = [self.smart_resize_and_pad(image) for image in images]
+
+        batch_doc = self(
+            text=texts_doc,
+            images=images,
+            return_tensors="pt",
+            padding="longest",
+        )
+        return batch_doc
+
+    def process_queries(self, queries, max_length=2048, suffix=None):
+        if suffix is None:
+            suffix = self.query_augmentation_token * self.suffix_len
+
+        processed = []
+        for q in queries:
+            q = self.query_start + self.query_prefix + q
+            # truncate before it eats actual query content
+            if len(q) + len(suffix) > max_length:
+                q = q[: max_length - len(suffix) - 1]
+            q += suffix + "\n"
+            processed.append(q)
+
+        return self(
+            text=processed,
+            images=None,
+            return_tensors="pt",
+            padding="longest",
+            truncation=True,
+            max_length=max_length,
+        )
+
+    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]:
+        n_patches_x = self.image_processor.size["width"] // patch_size
+        n_patches_y = self.image_processor.size["height"] // patch_size
+
+        return n_patches_x, n_patches_y
+
+    def get_image_mask(self, batch_images: BatchFeature) -> torch.Tensor:
+        return batch_images.input_ids == self.image_token_id
+
+    @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.
+        """
+
+        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: Union[torch.Tensor, List[torch.Tensor]],
+            ps: Union[torch.Tensor, List[torch.Tensor]],
+            batch_size: int = 128,
+            device: Optional[Union[str, torch.device]] = None,
+    ) -> torch.Tensor:
+        """
+        Compute the late-interaction/MaxSim score (ColBERT-like) for the given multi-vector
+        query embeddings (`qs`) and passage embeddings (`ps`). For ColPali, a passage is the
+        image of a document page.
+
+        Because the embedding tensors are multi-vector and can thus have different shapes, they
+        should be fed as:
+        (1) a list of tensors, where the i-th tensor is of shape (sequence_length_i, embedding_dim)
+        (2) a single tensor of shape (n_passages, max_sequence_length, embedding_dim) -> usually
+            obtained by padding the list of tensors.
+
+        Args:
+            qs (`Union[torch.Tensor, List[torch.Tensor]`): Query embeddings.
+            ps (`Union[torch.Tensor, List[torch.Tensor]`): Passage embeddings.
+            batch_size (`int`, *optional*, defaults to 128): Batch size for computing scores.
+            device (`Union[str, torch.device]`, *optional*): Device to use for computation. If not
+                provided, uses `get_torch_device("auto")`.
+
+        Returns:
+            `torch.Tensor`: A tensor of shape `(n_queries, n_passages)` containing the scores. The score
+            tensor is saved on the "cpu" device.
+        """
+
+        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