processing_colqwen_duo.py

import torch
import math

from abc import ABC, abstractmethod
from PIL import Image

from typing import Dict, List, Optional, Union, cast
from typing_extensions import Unpack

from transformers import BatchEncoding, BatchFeature
from transformers.processing_utils import (
    AllKwargsForChatTemplate,
    ImageInput,
    PreTokenizedInput,
    TextInput,
    VideoInput,
)
from transformers.models.qwen2_vl import Qwen2VLProcessor


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"
    return 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_texts(
        self,
        texts: List[str],
        max_length: int = 50,
        suffix: Optional[str] = None,
        prefix: 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



class QwenVLProcessor(ABC):

    def __call__(
        self,
        images: Optional[ImageInput] = None,
        text: Optional[Union[TextInput, PreTokenizedInput, List[PreTokenizedInput]]] = None,
        videos: Optional[VideoInput] = None,
        **kwargs,
    ) -> BatchFeature:
        return super().__call__(images=images, text=text, videos=videos, **kwargs)  # type: ignore

    def apply_chat_template(
        self,
        conversation: Union[List[Dict[str, str]], List[List[Dict[str, str]]]],
        chat_template: Optional[str] = None,
        **kwargs: Unpack[AllKwargsForChatTemplate],
    ) -> str:
        return super().apply_chat_template(conversation=conversation, chat_template=chat_template, **kwargs)  # type: ignore


class QwenVLEmbeddingProcessorBase(BaseVisualRetrieverProcessor, QwenVLProcessor):

    assistant_prefix_len: int = 58  # length of prefix created by
    # super().apply_chat_template(conversation=conversation, chat_template=chat_template, **kwargs)

    @staticmethod
    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

    @staticmethod
    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

    @staticmethod
    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

    def process_images(
        self,
        images: Union[List[Image.Image], List[List[Image.Image]]],
    ) -> BatchFeature:

        if isinstance(images[0], list):
            images = cast(List[List[Image.Image]], images)
            text_doc = []
            for i in range(len(images)):
                conversation = [{"role": "user", "content": [{"type": "image"}] * len(images[i])}]
                template = self.apply_chat_template(conversation, add_generation_prompt=False)
                text_doc.append(template[self.assistant_prefix_len :])

        else:
            images = cast(List[Image.Image], images)
            text_doc = [
                "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n"
            ] * len(images)

        # The following code is a hack to make sure the scatter in DDP is done correctly when training on multiple GPUs
        batch_doc = self(text=text_doc, images=images, padding="longest", return_tensors="pt")  # type: ignore
        # Separate pixel_values for each image
        offsets = batch_doc["image_grid_thw"][:, 1] * batch_doc["image_grid_thw"][:, 2]
        # Pad pixel_values to the same length to be able to make it into a tensor
        pixel_values = torch.split(batch_doc["pixel_values"], offsets.tolist())

        max_length = max([len(pv) for pv in pixel_values])

        pixel_values = [
            torch.cat([pv, torch.zeros((max_length - len(pv), pv.shape[1]), dtype=pv.dtype, device=pv.device)])
            for pv in pixel_values
        ]

        batch_doc["pixel_values"] = torch.stack(pixel_values)
        return batch_doc

    def process_texts(
        self,
        texts: List[str],
        max_length: int,
        suffix: Optional[str] = None,
        prefix: Optional[str] = None,
        padding: Optional[str] = None,
    ) -> BatchFeature:

        if suffix is None:
            suffix = "<pad>" * 10

        padded_texts: List[str] = []

        for text in texts:
            if prefix:
                text = f"{prefix}: {text}"
            text += suffix
            padded_texts.append(text)

        text_batch = self(
            text=padded_texts,
            return_tensors="pt",
            padding=padding or "longest",
            max_length=max_length,
            truncation=True,
        )

        return text_batch




class ColQwenDuoProcessorBase(QwenVLEmbeddingProcessorBase):
    """
    Processor for ColQwenDuo. Mirrors the `ColQwen2Processor` class.
    """

    def score(
        self,
        qs: List[torch.Tensor],
        ps: List[torch.Tensor],
        vector_type: str,
        device: Optional[Union[str, torch.device]] = None,
        truncate: Optional[int] = None,
        **kwargs,
    ) -> torch.Tensor:
        """
        Compute the MaxSim score (ColBERT-like) for the given multi-vector query and passage embeddings.
        """
        if truncate:
            qs = [q[..., :truncate] for q in qs]
            ps = [p[..., :truncate] for p in ps]

        if vector_type == "single_vector":
            return self.score_single_vector(qs, ps, device=device)
        elif vector_type == "multi_vector":
            return self.score_multi_vector(qs, ps, device=device, **kwargs)
        else:
            raise ValueError('vector_type must be one of the following: [`single_vector`, `multi_vector`]')


class ColQwen2DuoProcessor(ColQwenDuoProcessorBase, Qwen2VLProcessor):
    def __init__(self, *args, **kwargs) -> None:
        Qwen2VLProcessor.__init__(self, *args, **kwargs)