processing_colqwen35_bidirection.py

"""
Processing utilities for ColQwen35Bidirection retrieval.

Wraps the Qwen 3.5 VL processor components (image_processor, tokenizer,
video_processor) with retrieval-specific helpers for prompt construction,
MaxSim scoring, and batch handling.

processor kwargs:
    doc_prompt: Document prompt text appended after the image token.
    max_num_visual_tokens: Cap on visual tokens per image, controls
        resolution via max_pixels = max_num_visual_tokens × tile².
"""

from __future__ import annotations

import os
from typing import Any, List, Optional, Union

import numpy as np
from PIL import Image

from transformers import BatchFeature
from transformers.processing_utils import ProcessorMixin
from transformers.tokenization_utils_base import TextInput
from transformers.utils import logging

try:
    import torch
except ImportError:
    torch = None

logger = logging.get_logger(__name__)


def _size_value(size: Any, key: str) -> Any:
    """Read size entries from dict-like or object-like containers."""
    if size is None:
        return None
    if isinstance(size, dict):
        return size.get(key)
    return getattr(size, key, None)


class ColQwen35BidirectionProcessor(ProcessorMixin):
    """
    Processor for ColQwen35Bidirection retrieval model.

    Wraps Qwen 3.5's image processor, tokenizer, and video processor
    with retrieval-specific prompt construction, ``mm_token_type_ids``
    generation (required by Qwen 3.5 for 3-D position computation),
    and MaxSim scoring utilities.

    Visual token budget (``max_num_visual_tokens``):
    Qwen 3.5 determines visual token count from ``max_pixels`` on the
    image processor.  This class converts ``max_num_visual_tokens`` into
    the equivalent ``max_pixels`` value using:

        tile = patch_size × merge_size          # 16 × 2 = 32
        max_pixels = max_num_visual_tokens × tile²  # e.g. 512 × 1024 = 524,288

    Lower token budgets (e.g. 512) give memory-efficient training;
    higher budgets (e.g. 2048) give finer visual granularity at inference.
    """

    attributes = ["image_processor", "tokenizer", "video_processor"]
    image_processor_class = "AutoImageProcessor"
    video_processor_class = "AutoVideoProcessor"
    tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")

    def __init__(
        self,
        image_processor=None,
        tokenizer=None,
        video_processor=None,
        chat_template=None,
        doc_prompt: str = "Describe the image.",
        max_num_visual_tokens: Optional[int] = None,
        query_augmentation_tokens: int = 10,
        **kwargs,
    ):
        super().__init__(
            image_processor, tokenizer, video_processor,
            chat_template=chat_template, **kwargs,
        )

        self.doc_prompt = doc_prompt
        self.max_num_visual_tokens = max_num_visual_tokens
        self.query_augmentation_tokens = query_augmentation_tokens

        if max_num_visual_tokens is not None:
            self._apply_max_pixels()

        self.image_token = (
            tokenizer.image_token
            if getattr(tokenizer, "image_token", None)
            else "<|image_pad|>"
        )
        self.image_token_id = (
            tokenizer.image_token_id
            if getattr(tokenizer, "image_token_id", None)
            else tokenizer.convert_tokens_to_ids(self.image_token)
        )
        self.vision_start_token = (
            tokenizer.vision_start_token
            if getattr(tokenizer, "vision_start_token", None)
            else "<|vision_start|>"
        )
        self.vision_end_token = (
            tokenizer.vision_end_token
            if getattr(tokenizer, "vision_end_token", None)
            else "<|vision_end|>"
        )

        self.tokenizer.padding_side = "left"

        self._doc_prompt_template = (
            "<|im_start|>user\n"
            f"{self.vision_start_token}{self.image_token}{self.vision_end_token}"
            f"{self.doc_prompt}"
            "<|im_end|><|endoftext|>"
        )

    # ------------------------------------------------------------------
    # max_pixels / visual token budget
    # ------------------------------------------------------------------

    def _apply_max_pixels(self) -> None:
        """Sync image_processor.max_pixels with max_num_visual_tokens.

        Sets ``max_pixels`` on the image processor attribute AND in the
        ``size`` dict (``longest_edge``), because different versions of
        the Qwen2VLImageProcessor read from different locations.  Also
        ensures ``min_pixels`` (``shortest_edge``) does not exceed
        ``max_pixels``, which would cause the resize to ignore the cap.
        """
        patch_size = getattr(self.image_processor, "patch_size", None)
        merge_size = (
            getattr(self.image_processor, "merge_size", None)
            or getattr(self.image_processor, "spatial_merge_size", None)
        )
        if patch_size is None or merge_size is None:
            logger.warning(
                "Cannot derive max_pixels: image_processor missing "
                "patch_size or merge_size/spatial_merge_size."
            )
            return
        tile = patch_size * merge_size
        max_pixels = self.max_num_visual_tokens * tile * tile

        self.image_processor.max_pixels = max_pixels
        size_obj = getattr(self.image_processor, "size", None)
        if size_obj is not None:
            if isinstance(size_obj, dict):
                size_obj["longest_edge"] = max_pixels
                cur_min = size_obj.get("shortest_edge")
                if cur_min is not None and cur_min > max_pixels:
                    size_obj["shortest_edge"] = max_pixels
            else:
                if hasattr(size_obj, "longest_edge"):
                    size_obj.longest_edge = max_pixels
                cur_min = getattr(size_obj, "shortest_edge", None)
                if cur_min is not None and cur_min > max_pixels and hasattr(size_obj, "shortest_edge"):
                    size_obj.shortest_edge = max_pixels

        cur_min_pixels = getattr(self.image_processor, "min_pixels", 0)
        if cur_min_pixels > max_pixels:
            self.image_processor.min_pixels = max_pixels

    def replace_image_token(self, image_inputs: dict, image_idx: int, **kwargs) -> str:
        """Expand one ``<|image_pad|>`` placeholder into its per-image token run.

        ``ProcessorMixin.__call__`` delegates placeholder expansion here, and
        ``apply_chat_template`` calls ``__call__``, so without this both raise
        ``NotImplementedError`` for image inputs.
        """
        merge_length = self.image_processor.merge_size ** 2
        num_image_tokens = image_inputs["image_grid_thw"][image_idx].prod() // merge_length
        return self.image_token * num_image_tokens

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: Union[str, os.PathLike],
        *,
        max_num_visual_tokens: Optional[int] = None,
        doc_prompt: Optional[str] = None,
        query_augmentation_tokens: Optional[int] = None,
        **kwargs,
    ) -> "ColQwen35BidirectionProcessor":
        extra_kwargs: dict[str, Any] = {}
        if doc_prompt is not None:
            extra_kwargs["doc_prompt"] = doc_prompt
        if max_num_visual_tokens is not None:
            extra_kwargs["max_num_visual_tokens"] = max_num_visual_tokens
        if query_augmentation_tokens is not None:
            extra_kwargs["query_augmentation_tokens"] = query_augmentation_tokens

        instance = super().from_pretrained(
            pretrained_model_name_or_path, **extra_kwargs, **kwargs,
        )

        if max_num_visual_tokens is not None:
            instance.max_num_visual_tokens = max_num_visual_tokens
            instance._apply_max_pixels()

        if query_augmentation_tokens is not None:
            instance.query_augmentation_tokens = query_augmentation_tokens

        return instance

    # ------------------------------------------------------------------
    # Retrieval protocol: process_images
    # ------------------------------------------------------------------

    @property
    def query_augmentation_token(self) -> str:
        return self.tokenizer.pad_token

    def process_images(
        self,
        images: Union[Image.Image, List[Image.Image]],
    ) -> BatchFeature:
        """
        Tokenize and encode document images for retrieval.

        Each image is independently processed with the doc_prompt,
        and ``mm_token_type_ids`` is computed for Qwen 3.5's 3-D
        positional encoding. Multiple images are left-padded and
        concatenated into a single batch.
        """
        if not isinstance(images, list):
            images = [images]
        if len(images) == 0:
            raise ValueError("No images provided")

        images = [img.convert("RGB") for img in images]

        per_image_features: list[BatchFeature] = []
        for image in images:
            features = self._process_single_image(image)
            per_image_features.append(features)

        if len(per_image_features) == 1:
            return per_image_features[0]

        return self._left_pad_and_concat(per_image_features)

    def _process_single_image(self, image: Image.Image) -> BatchFeature:
        """Process one image through the full pipeline with mm_token_type_ids."""
        size_obj = getattr(self.image_processor, "size", None)
        min_pixels = _size_value(size_obj, "shortest_edge")
        if min_pixels is None:
            min_pixels = getattr(self.image_processor, "min_pixels", None)

        max_pixels = _size_value(size_obj, "longest_edge")
        if max_pixels is None:
            max_pixels = getattr(self.image_processor, "max_pixels", None)
        ip_kwargs: dict[str, Any] = {
            "images": [[image]],
        }
        if min_pixels is not None:
            ip_kwargs["min_pixels"] = int(min_pixels)
        if max_pixels is not None:
            ip_kwargs["max_pixels"] = int(max_pixels)
        image_inputs = self.image_processor(**ip_kwargs)
        image_grid_thw = image_inputs["image_grid_thw"]

        merge_size = (
            getattr(self.image_processor, "merge_size", None)
            or getattr(self.image_processor, "spatial_merge_size", None)
        )
        if merge_size is None:
            raise ValueError(
                "Image processor missing merge_size/spatial_merge_size."
            )
        merge_length = merge_size ** 2

        prompt = self._doc_prompt_template
        for grid in image_grid_thw:
            num_image_tokens = int(grid.prod() // merge_length) if hasattr(grid, 'prod') else int(np.prod(grid) // merge_length)
            prompt = prompt.replace(
                self.image_token,
                "<|placeholder|>" * num_image_tokens,
                1,
            )
        prompt = prompt.replace("<|placeholder|>", self.image_token)

        text_inputs = self.tokenizer(
            [prompt], padding="longest", return_tensors="pt",
        )

        input_ids = text_inputs["input_ids"]
        mm_token_type_ids = (input_ids == self.image_token_id).to(torch.int32)

        data = {**text_inputs, **image_inputs}
        data["mm_token_type_ids"] = mm_token_type_ids

        for key in ("input_ids", "attention_mask"):
            if key in data and not isinstance(data[key], torch.Tensor):
                data[key] = torch.tensor(data[key])

        return BatchFeature(data=data, tensor_type="pt")

    # ------------------------------------------------------------------
    # Retrieval protocol: process_queries
    # ------------------------------------------------------------------

    def process_queries(
        self,
        texts: Union[TextInput, List[TextInput]],
    ) -> BatchFeature:
        """
        Process text queries for retrieval.

        Each query is wrapped in a simple chat template and tokenized.
        """
        if not isinstance(texts, list):
            texts = [texts]
        if len(texts) == 0:
            raise ValueError("No texts provided")

        suffix = self.query_augmentation_token * self.query_augmentation_tokens
        formatted: list[str] = []
        for text in texts:
            prompt = f"<|im_start|>user\nQuery: {text}{suffix}<|im_end|><|endoftext|>"
            formatted.append(prompt)

        return self.tokenizer(
            formatted,
            return_tensors="pt",
            padding="longest",
        )

    # ------------------------------------------------------------------
    # Scoring utilities
    # ------------------------------------------------------------------

    def score_retrieval(
        self,
        query_embeddings: Union[torch.Tensor, List[torch.Tensor]],
        passage_embeddings: Union[torch.Tensor, List[torch.Tensor]],
        batch_size: int = 128,
        output_dtype: Optional[torch.dtype] = None,
        output_device: Union[torch.device, str] = "cpu",
    ) -> torch.Tensor:
        """
        Compute late-interaction / MaxSim retrieval scores (ColBERT-like).

        Args:
            query_embeddings: Per-query multi-vector embeddings.
            passage_embeddings: Per-passage multi-vector embeddings.
            batch_size: Scoring batch size.
            output_dtype: Output tensor dtype.
            output_device: Output device.

        Returns:
            Tensor of shape ``(n_queries, n_passages)`` with scores.
        """
        if len(query_embeddings) == 0:
            raise ValueError("No queries provided")
        if len(passage_embeddings) == 0:
            raise ValueError("No passages provided")

        if output_dtype is None:
            output_dtype = query_embeddings[0].dtype

        scores: list[torch.Tensor] = []
        for i in range(0, len(query_embeddings), batch_size):
            batch_queries = torch.nn.utils.rnn.pad_sequence(
                query_embeddings[i : i + batch_size],
                batch_first=True, padding_value=0,
            )
            batch_scores: list[torch.Tensor] = []
            for j in range(0, len(passage_embeddings), batch_size):
                batch_passages = torch.nn.utils.rnn.pad_sequence(
                    passage_embeddings[j : j + batch_size],
                    batch_first=True, padding_value=0,
                )
                batch_scores.append(
                    torch.einsum("bnd,csd->bcns", batch_queries, batch_passages)
                    .max(dim=3)[0]
                    .sum(dim=2)
                )
            scores.append(
                torch.cat(batch_scores, dim=1)
                .to(output_dtype)
                .to(output_device)
            )
        return torch.cat(scores, dim=0)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _left_pad_and_concat(
        batch_features: list[BatchFeature],
    ) -> BatchFeature:
        """
        Left-pad variable-length BatchFeature dicts and stack them.

        Qwen 3.5 yields a variable number of visual tokens per image
        (resolution dependent), so we align them before concatenation.
        Padding is on the left (decoder convention).
        """
        all_keys = batch_features[0].keys()
        concatenated: dict[str, Any] = {}

        for key in all_keys:
            tensors = [bf[key] for bf in batch_features]

            if not isinstance(tensors[0], torch.Tensor):
                concatenated[key] = tensors[0]
                continue

            if tensors[0].ndim < 2:
                concatenated[key] = torch.cat(tensors, dim=0)
                continue

            max_seq_len = max(t.shape[1] for t in tensors)
            padded: list[torch.Tensor] = []
            for t in tensors:
                pad_len = max_seq_len - t.shape[1]
                if pad_len > 0:
                    zeros = torch.zeros(
                        *t.shape[:1], pad_len, *t.shape[2:],
                        dtype=t.dtype, device=t.device,
                    )
                    t = torch.cat([zeros, t], dim=1)
                padded.append(t)

            concatenated[key] = torch.cat(padded, dim=0)

        return BatchFeature(concatenated)


__all__ = ["ColQwen35BidirectionProcessor"]