colqwen2base-v0.1-hf / processing_colqwen2.py
manu's picture
Update processing_colqwen2.py
31bc9b2 verified
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)
# print(f"Resizing image from {image_size} to {(resized_height, resized_width)}")
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"
)
# The following code is a hack to make sure the scatter in DDP is done correctly when training on multiple GPUs
offsets = batch_doc["image_grid_thw"][:, 1] * batch_doc["image_grid_thw"][:, 2]
# separate pixel_values for each image
pixel_values = torch.split(batch_doc["pixel_values"], offsets.tolist())
# pad pixel_values to the same length to be able to make it into a tensor
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 # add suffix (pad tokens)
texts_query.append(query)
batch_query = self(
text=texts_query,
return_tensors="pt",
padding="longest",
# max_length=max_length + self.image_seq_length,
)
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