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