|
import math |
|
from typing import List, Optional, Union |
|
|
|
import torch |
|
from PIL import Image |
|
from transformers import BatchFeature |
|
from transformers.models.qwen2_vl import Qwen2VLProcessor |
|
|
|
|
|
class ColQwen2Processor(Qwen2VLProcessor): |
|
""" |
|
Processor for ColQwen2. |
|
""" |
|
|
|
def __init__(self, *args, **kwargs): |
|
super().__init__(*args, **kwargs) |
|
self.tokenizer.padding_side = "left" |
|
self.min_pixels = 4 * 28 * 28 |
|
self.max_pixels = 768 * 28 * 28 |
|
self.factor = 28 |
|
self.max_ratio = 200 |
|
|
|
@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 smart_resize(self, height: int, width: int, factor: int, min_pixels: int, max_pixels: int) -> tuple[int, int]: |
|
""" |
|
Rescales the image so that the following conditions are met: |
|
|
|
1. Both dimensions (height and width) are divisible by 'factor'. |
|
|
|
2. The total number of pixels is within the range ['min_pixels', 'max_pixels']. |
|
|
|
3. The aspect ratio of the image is maintained as closely as possible. |
|
""" |
|
if max(height, width) / min(height, width) > self.max_ratio: |
|
raise ValueError( |
|
f"absolute aspect ratio must be smaller than {self.max_ratio}, " |
|
f"got {max(height, width) / min(height, width)}" |
|
) |
|
h_bar = max(factor, self.round_by_factor(height, factor)) |
|
w_bar = max(factor, self.round_by_factor(width, factor)) |
|
if h_bar * w_bar > max_pixels: |
|
beta = math.sqrt((height * width) / max_pixels) |
|
h_bar = self.floor_by_factor(height / beta, factor) |
|
w_bar = self.floor_by_factor(width / beta, factor) |
|
elif h_bar * w_bar < min_pixels: |
|
beta = math.sqrt(min_pixels / (height * width)) |
|
h_bar = self.ceil_by_factor(height * beta, factor) |
|
w_bar = self.ceil_by_factor(width * beta, factor) |
|
return h_bar, w_bar |
|
|
|
def process_images( |
|
self, |
|
images: List[Image.Image], |
|
) -> BatchFeature: |
|
""" |
|
Process images for ColPali. |
|
""" |
|
texts_doc = (["<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n"] |
|
* len(images)) |
|
|
|
def resize_and_convert(image: Image.Image) -> Image.Image: |
|
image_size = image.size |
|
resized_height, resized_width = self.smart_resize(image_size[1], |
|
image_size[0], |
|
factor=self.factor, |
|
min_pixels=self.min_pixels, |
|
max_pixels=self.max_pixels) |
|
|
|
return image.convert("RGB").resize((resized_width, resized_height)) |
|
|
|
images = [resize_and_convert(image) for image in images] |
|
|
|
|
|
batch_doc = self( |
|
text=texts_doc, |
|
images=images, |
|
padding="longest", |
|
return_tensors="pt" |
|
) |
|
|
|
|
|
|
|
offsets = batch_doc["image_grid_thw"][:, 1] * batch_doc["image_grid_thw"][:, 2] |
|
|
|
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_queries( |
|
self, |
|
queries: List[str], |
|
max_length: int = 50, |
|
suffix: Optional[str] = None, |
|
) -> BatchFeature: |
|
""" |
|
Process queries for ColPali. |
|
""" |
|
if suffix is None: |
|
suffix = "<pad>" * 10 |
|
texts_query: List[str] = [] |
|
|
|
for query in queries: |
|
query = f"Query: {query}" |
|
query += suffix |
|
texts_query.append(query) |
|
|
|
batch_query = self( |
|
text=texts_query, |
|
return_tensors="pt", |
|
padding="longest", |
|
|
|
) |
|
|
|
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) |
|
|
|
|
|
@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 |
|
|