""" 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 @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"]