"""Contextual retrieval chunks from a whole image's merged visual tokens.""" from dataclasses import dataclass import torch from PIL import Image from torch.nn.utils.rnn import pad_sequence from transformers import BatchEncoding @dataclass(frozen=True) class ImageInputMapper: _processor: object tokenizer: object document_prefix: str boundary_marker: str @dataclass(frozen=True) class ImageEncodingConfig: chunks: int = 8 min_pixels: int = 16384 max_pixels: int = 1800000 def __post_init__(self) -> None: assert self.chunks > 0 assert 0 < self.min_pixels <= self.max_pixels def chunk_indices(self, indices: torch.Tensor, grid: torch.Tensor, merge_size: int) -> list[torch.Tensor]: frames, height, width = map(int, grid) assert frames == 1 and height % merge_size == 0 and width % merge_size == 0 height, width = height // merge_size, width // merge_size assert len(indices) == height * width, "Visual token count does not match merged image grid" return [band.flatten() for band in indices.reshape(height, width).tensor_split(min(self.chunks, height), dim=0)] def prepare_native_batch(self, mapper: ImageInputMapper, images: list[Image.Image]) -> BatchEncoding: processor = mapper._processor placeholder = processor.vision_start_token + processor.image_token + processor.vision_end_token batch = BatchEncoding(processor( text=[placeholder] * len(images), images=[[image] for image in images], return_tensors="pt", padding=True, min_pixels=self.min_pixels, max_pixels=self.max_pixels, )) prefix_id = mapper.tokenizer.convert_tokens_to_ids(mapper.document_prefix) prefix = torch.full((len(images), 1), prefix_id, dtype=torch.long) batch["input_ids"] = torch.cat((prefix, batch["input_ids"]), dim=1) batch["attention_mask"] = torch.cat((torch.ones_like(prefix), batch["attention_mask"]), dim=1) for key in ("mm_token_type_ids", "token_type_ids"): if key in batch: batch[key] = torch.cat((torch.zeros_like(prefix), batch[key]), dim=1) assert "position_ids" not in batch image_id = mapper.tokenizer.convert_tokens_to_ids(mapper._processor.image_token) visual = batch["input_ids"] == image_id assert torch.equal(batch["mm_token_type_ids"] == 1, visual) spans = [] for row, grid in zip(visual, batch["image_grid_thw"], strict=True): indices = row.nonzero().flatten() assert torch.equal(indices, torch.arange(indices[0], indices[-1] + 1)) groups = self.chunk_indices(indices, grid, mapper._processor.image_processor.merge_size) spans.append([(int(part[0]), int(part[-1]) + 1) for part in groups]) batch["chunk_token_spans"] = pad_sequence( [torch.tensor(row, dtype=torch.long) for row in spans], batch_first=True ) batch["chunk_mask"] = pad_sequence([torch.ones(len(row), dtype=torch.bool) for row in spans], batch_first=True) return batch def prepare(self, mapper: ImageInputMapper, images: list[Image.Image]) -> BatchEncoding: batch = self.prepare_native_batch(mapper, images) positions = torch.zeros(3, *batch["input_ids"].shape, dtype=torch.float32) merge = mapper._processor.image_processor.merge_size for row, grid in enumerate(batch["image_grid_thw"]): active = batch["attention_mask"][row].bool() visual = batch["mm_token_type_ids"][row, active] == 1 indices = visual.nonzero().flatten() start, end = int(indices[0]), int(indices[-1]) + 1 height, width = int(grid[1]) // merge, int(grid[2]) // merge native = torch.arange(int(active.sum())).float().expand(3, -1).clone() native[:, start:end] = torch.stack( ( torch.full((height * width,), start), torch.arange(height).repeat_interleave(width) + start, torch.arange(width).repeat(height) + start, ) ) native[:, end:] = torch.arange(native.shape[1] - end) + start + max(height, width) positions[:, row, active] = native batch["position_ids"] = positions return self.insert_native_separators(batch, mapper) def boundary_positions(self, native: torch.Tensor, start: int, visual: torch.Tensor) -> torch.Tensor: left, right = native[:, start - 1], native[:, start] assert left[0] == right[0] and right[1] == left[1] + 1, "Separator must lie between adjacent image rows" widths = native[2, visual] return torch.stack((left[0], (left[1] + right[1]) / 2, (widths.min() + widths.max()) / 2))[:, None] def insert_native_separators(self, batch: BatchEncoding, mapper: ImageInputMapper) -> BatchEncoding: separator = mapper.tokenizer.encode(mapper.boundary_marker, add_special_tokens=False, split_special_tokens=True) assert len(separator) == 1, "Image chunk boundaries require a single separator token" assert batch["position_ids"].shape[0] == 3 token_keys = ["input_ids", "mm_token_type_ids"] if "token_type_ids" in batch: token_keys.append("token_type_ids") rows = {key: [] for key in token_keys} positions, spans = [], [] for row, mask in enumerate(batch["attention_mask"]): active = mask.nonzero().flatten() assert torch.equal(active, torch.arange(active[0], active[-1] + 1, device=active.device)) native = batch["position_ids"][:, row, active].float() original = {key: batch[key][row, active] for key in token_keys} parts = {key: [] for key in token_keys} position_parts, row_spans = [], [] cursor, added = 0, 0 valid_spans = batch["chunk_token_spans"][row, batch["chunk_mask"][row]] for index, (start, end) in enumerate(valid_spans.tolist()): start, end = start - int(active[0]), end - int(active[0]) if index: sep_positions = self.boundary_positions(native, start, original["mm_token_type_ids"] == 1) count = sep_positions.shape[1] for key in token_keys: parts[key].append(original[key][cursor:start]) values = separator * (count // len(separator)) if key == "input_ids" else [0] * count parts[key].append(original[key].new_tensor(values)) position_parts.extend((native[:, cursor:start], sep_positions)) added += count cursor = start row_spans.append((start + added, end + added)) for key in token_keys: parts[key].append(original[key][cursor:]) rows[key].append(torch.cat(parts[key])) position_parts.append(native[:, cursor:]) positions.append(torch.cat(position_parts, dim=1).T) spans.append(row_spans) for key in token_keys: batch[key] = pad_sequence( rows[key], batch_first=True, padding_value=mapper.tokenizer.pad_token_id if key == "input_ids" else 0, ) batch["attention_mask"] = pad_sequence([torch.ones_like(row) for row in rows["input_ids"]], batch_first=True) batch["position_ids"] = pad_sequence(positions, batch_first=True).permute(2, 0, 1) batch["chunk_token_spans"] = pad_sequence( [batch["chunk_token_spans"].new_tensor(row) for row in spans], batch_first=True ) return batch