File size: 7,796 Bytes
697df6e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | """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
|