| |
| """Multimodal tokenizer processor.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass, field |
| from typing import Any |
|
|
| from Tokenizer.unified.encoded import EncodedToken |
|
|
| from .image_placeholders import ( |
| image_patch_count, |
| image_placeholder_tokens, |
| ) |
| from .video_placeholders import ( |
| video_patch_count, |
| video_placeholder_tokens, |
| ) |
| from .tokens import IMAGE_PLACEHOLDER, VIDEO_PLACEHOLDER |
|
|
|
|
| @dataclass |
| class MultimodalEncoding: |
| input_ids: list[int] |
| attention_mask: list[int] |
| tokens: list[EncodedToken] |
| image_token_spans: list[tuple[int, int]] |
| video_token_spans: list[tuple[int, int]] = field(default_factory=list) |
| images: list[Any] = field(default_factory=list) |
| videos: list[Any] = field(default_factory=list) |
| pixel_values: Any = None |
|
|
|
|
| @dataclass(frozen=True) |
| class _PlaceholderRecord: |
| start: int |
| end: int |
| patch_count: int |
|
|
|
|
| class MultimodalProcessor: |
| def __init__( |
| self, |
| tokenizer, |
| image_processor=None, |
| patch_size: int = 14, |
| merge_size: int = 2, |
| temporal_patch_size: int = 2, |
| ): |
| self.tokenizer = tokenizer |
| self.image_processor = image_processor |
| self.patch_size = patch_size |
| self.merge_size = merge_size |
| self.temporal_patch_size = temporal_patch_size |
|
|
| def __call__( |
| self, |
| text: str, |
| images: list[Any] | None = None, |
| image_sizes: list[Any] | None = None, |
| videos: list[Any] | None = None, |
| video_sizes: list[Any] | None = None, |
| add_bos: bool = False, |
| add_eos: bool = False, |
| ) -> MultimodalEncoding: |
| image_list = list(images or []) |
| sizes = list(image_sizes or []) |
| video_list = list(videos or []) |
| vsizes = list(video_sizes or []) |
|
|
| n_img_holes = text.count(IMAGE_PLACEHOLDER) |
| n_vid_holes = text.count(VIDEO_PLACEHOLDER) |
|
|
| if len(image_list) != n_img_holes: |
| raise ValueError( |
| f"<image> placeholder count ({n_img_holes}) does not match " |
| f"images length ({len(image_list)})" |
| ) |
| if len(video_list) != n_vid_holes: |
| raise ValueError( |
| f"<video> placeholder count ({n_vid_holes}) does not match " |
| f"videos length ({len(video_list)})" |
| ) |
| if sizes and len(sizes) != n_img_holes: |
| raise ValueError( |
| f"<image> placeholder count ({n_img_holes}) does not match " |
| f"image_sizes length ({len(sizes)})" |
| ) |
| if vsizes and len(vsizes) != n_vid_holes: |
| raise ValueError( |
| f"<video> placeholder count ({n_vid_holes}) does not match " |
| f"video_sizes length ({len(vsizes)})" |
| ) |
|
|
| |
| if not sizes and n_img_holes: |
| sizes = [None] * n_img_holes |
| if not vsizes and n_vid_holes: |
| vsizes = [None] * n_vid_holes |
|
|
| expanded, image_records, video_records, boundary_map = self._expand_placeholders( |
| text, sizes, vsizes |
| ) |
|
|
| result = self.tokenizer.encode_with_spans( |
| expanded, add_bos=add_bos, add_eos=add_eos |
| ) |
| tokens = self._annotate_patches( |
| result.tokens, image_records, video_records, boundary_map |
| ) |
| img_spans = self._collect_spans( |
| tokens, "<image_start>", "<image_patch>", "<image_end>", image_records |
| ) |
| vid_spans = self._collect_spans( |
| tokens, "<video_start>", "<video_patch>", "<video_end>", video_records |
| ) |
|
|
| |
| assert len(result.attention_mask) == len(result.input_ids) |
|
|
| processed_images = image_list |
| if self.image_processor is not None and image_list: |
| processed_images = self.image_processor(image_list) |
|
|
| return MultimodalEncoding( |
| input_ids=result.input_ids, |
| attention_mask=result.attention_mask, |
| tokens=tokens, |
| image_token_spans=img_spans, |
| video_token_spans=vid_spans, |
| images=processed_images, |
| videos=video_list, |
| ) |
|
|
| |
|
|
| def _expand_placeholders( |
| self, text: str, image_sizes: list[Any], video_sizes: list[Any] |
| ) -> tuple[str, list[_PlaceholderRecord], list[_PlaceholderRecord], list[int]]: |
| out: list[str] = [] |
| boundary_map = [0] |
| image_records: list[_PlaceholderRecord] = [] |
| video_records: list[_PlaceholderRecord] = [] |
| img_idx = 0 |
| vid_idx = 0 |
| i = 0 |
| while i < len(text): |
| if text.startswith(IMAGE_PLACEHOLDER, i): |
| n = self._patches_image(image_sizes[img_idx]) |
| end = i + len(IMAGE_PLACEHOLDER) |
| image_records.append(_PlaceholderRecord(i, end, n)) |
| expanded = "".join(image_placeholder_tokens(n)) |
| out.append(expanded) |
| boundary_map.extend([end] * len(expanded)) |
| img_idx += 1 |
| i = end |
| continue |
| if text.startswith(VIDEO_PLACEHOLDER, i): |
| n = self._patches_video(video_sizes[vid_idx]) |
| end = i + len(VIDEO_PLACEHOLDER) |
| video_records.append(_PlaceholderRecord(i, end, n)) |
| expanded = "".join(video_placeholder_tokens(n)) |
| out.append(expanded) |
| boundary_map.extend([end] * len(expanded)) |
| vid_idx += 1 |
| i = end |
| continue |
| out.append(text[i]) |
| boundary_map.append(i + 1) |
| i += 1 |
| return "".join(out), image_records, video_records, boundary_map |
|
|
| def _expand_images(self, text: str, sizes: list[Any]) -> str: |
| parts = text.split(IMAGE_PLACEHOLDER) |
| if len(parts) == 1: |
| return text |
| out = [parts[0]] |
| for i, size in enumerate(sizes): |
| n = self._patches_image(size) |
| out.append("".join(image_placeholder_tokens(n))) |
| out.append(parts[i + 1]) |
| return "".join(out) |
|
|
| def _expand_videos(self, text: str, sizes: list[Any]) -> str: |
| parts = text.split(VIDEO_PLACEHOLDER) |
| if len(parts) == 1: |
| return text |
| out = [parts[0]] |
| for i, size in enumerate(sizes): |
| n = self._patches_video(size) |
| out.append("".join(video_placeholder_tokens(n))) |
| out.append(parts[i + 1]) |
| return "".join(out) |
|
|
| def _patches_image(self, size: Any) -> int: |
| if size is None: |
| return 1 |
| if isinstance(size, dict): |
| w = int(size.get("width") or size.get("w")) |
| h = int(size.get("height") or size.get("h")) |
| else: |
| w, h = int(size[0]), int(size[1]) |
| return image_patch_count(w, h, self.patch_size, self.merge_size) |
|
|
| def _patches_video(self, size: Any) -> int: |
| if size is None: |
| return 1 |
| if isinstance(size, dict): |
| f = int(size.get("frames") or size.get("num_frames")) |
| w = int(size.get("width") or size.get("w")) |
| h = int(size.get("height") or size.get("h")) |
| else: |
| f, w, h = int(size[0]), int(size[1]), int(size[2]) |
| return video_patch_count( |
| f, w, h, self.patch_size, self.temporal_patch_size, self.merge_size |
| ) |
|
|
| |
|
|
| def _annotate_patches( |
| self, |
| tokens: list[EncodedToken], |
| image_records: list[_PlaceholderRecord], |
| video_records: list[_PlaceholderRecord], |
| boundary_map: list[int], |
| ) -> list[EncodedToken]: |
| """Attach image_index/video_index metadata to patch tokens.""" |
| out: list[EncodedToken] = [] |
| img_idx = -1 |
| vid_idx = -1 |
| in_img = False |
| in_vid = False |
| for t in tokens: |
| meta = t.metadata |
| start, end = self._remap_offsets(t.start, t.end, boundary_map) |
| if t.token == "<image_start>": |
| img_idx += 1 |
| in_img = True |
| meta = self._placeholder_metadata(meta, image_records, img_idx, "image_index") |
| elif t.token == "<image_end>": |
| meta = self._placeholder_metadata(meta, image_records, img_idx, "image_index") |
| in_img = False |
| elif t.token == "<image_patch>" and in_img: |
| meta = self._placeholder_metadata(meta, image_records, img_idx, "image_index") |
| elif t.token == "<video_start>": |
| vid_idx += 1 |
| in_vid = True |
| meta = self._placeholder_metadata(meta, video_records, vid_idx, "video_index") |
| elif t.token == "<video_end>": |
| meta = self._placeholder_metadata(meta, video_records, vid_idx, "video_index") |
| in_vid = False |
| elif t.token == "<video_patch>" and in_vid: |
| meta = self._placeholder_metadata(meta, video_records, vid_idx, "video_index") |
| if meta is not t.metadata: |
| source_span = meta.get("source_span") if meta else None |
| if source_span is not None: |
| start, end = int(source_span[0]), int(source_span[1]) |
| out.append( |
| EncodedToken( |
| t.id, t.token, t.track, start, end, t.surface, meta |
| ) |
| ) |
| elif start != t.start or end != t.end: |
| out.append( |
| EncodedToken( |
| t.id, t.token, t.track, start, end, t.surface, t.metadata |
| ) |
| ) |
| else: |
| out.append(t) |
| return out |
|
|
| def _remap_offsets( |
| self, start: int, end: int, boundary_map: list[int] |
| ) -> tuple[int, int]: |
| if start == -1 and end == -1: |
| return start, end |
| if not boundary_map: |
| return start, end |
| start_idx = min(max(start, 0), len(boundary_map) - 1) |
| end_idx = min(max(end, 0), len(boundary_map) - 1) |
| return boundary_map[start_idx], boundary_map[end_idx] |
|
|
| def _placeholder_metadata( |
| self, |
| meta: dict | None, |
| records: list[_PlaceholderRecord], |
| index: int, |
| index_key: str, |
| ) -> dict: |
| merged = {**(meta or {}), index_key: index} |
| if 0 <= index < len(records): |
| record = records[index] |
| merged["source_span"] = [record.start, record.end] |
| return merged |
|
|
| def _collect_spans( |
| self, |
| tokens: list[EncodedToken], |
| start_tok: str, |
| patch_tok: str, |
| end_tok: str, |
| records: list[_PlaceholderRecord], |
| ) -> list[tuple[int, int]]: |
| spans: list[tuple[int, int]] = [] |
| start: int | None = None |
| for idx, t in enumerate(tokens): |
| if t.token == start_tok: |
| start = idx |
| elif t.token == end_tok and start is not None: |
| spans.append((start, idx + 1)) |
| start = None |
| if len(spans) != len(records): |
| raise ValueError( |
| f"{start_tok}/{end_tok} span count ({len(spans)}) does not match " |
| f"placeholder count ({len(records)})" |
| ) |
| for i, (span_start, span_end) in enumerate(spans): |
| patch_count = sum(1 for token in tokens[span_start:span_end] if token.token == patch_tok) |
| if patch_count != records[i].patch_count: |
| raise ValueError( |
| f"{patch_tok} count in span {i} is {patch_count}, " |
| f"expected {records[i].patch_count}" |
| ) |
| if span_end - span_start != records[i].patch_count + 2: |
| raise ValueError( |
| f"{start_tok}/{end_tok} span length for placeholder {i} is " |
| f"{span_end - span_start}, expected {records[i].patch_count + 2}" |
| ) |
| return spans |
|
|
