Buckets:
| """Streaming Zebra-CoT validation and UniAR-native 960px cache utilities.""" | |
| from __future__ import annotations | |
| import hashlib | |
| import io | |
| import json | |
| import re | |
| from collections.abc import Iterator | |
| from pathlib import Path | |
| import pyarrow.parquet as pq | |
| import torch | |
| from PIL import Image | |
| from shared.labels import canonical_autonomous_text_labels | |
| from .bsq_codes import pack_visual_codes, unpack_visual_codes | |
| from .dataset import image_gen_token_id | |
| from .objective import encode_visual_codes | |
| from .trajectory import ( | |
| build_trajectory, | |
| render_messages, | |
| split_image_references, | |
| symbolic_render, | |
| ) | |
| TEXT_COLUMNS = ("Question", "Text Reasoning Trace", "Final Answer") | |
| CAP_CANDIDATES = (960, 896, 832, 768, 704, 640, 576, 512, 448, 384, 320, 256, 224, 192, 160, 128, 96, 64) | |
| GRAPH_CONFIG = "Scientific Reasoning - Graph Algorithms" | |
| CACHE_SCHEMA = "zebra_uniar_cache_v1" | |
| LEGACY_IMAGE_PREPROCESSING = "legacy_cap_qwen" | |
| UNDERSTANDING_INPUT_GRID16_TARGET = "understanding_input_grid16_target" | |
| IMAGE_PREPROCESSING_PROFILES = ( | |
| LEGACY_IMAGE_PREPROCESSING, | |
| UNDERSTANDING_INPUT_GRID16_TARGET, | |
| ) | |
| class InvalidZebraRow(ValueError): | |
| """A source row cannot be represented by the canonical training contract.""" | |
| def __init__(self, reason: str, detail: str): | |
| super().__init__(detail) | |
| self.reason = reason | |
| def source_shards(source_root: Path) -> list[Path]: | |
| """Return the canonical 163 consolidated Parquet shards.""" | |
| return sorted( | |
| path | |
| for path in source_root.glob("*/*.parquet") | |
| if path.parent.name != ".git" | |
| ) | |
| def _normalize_question(text: str) -> str: | |
| return re.sub(r"\s+", " ", text).strip().casefold() | |
| def _image_bytes(cell: dict | None, key: str) -> bytes: | |
| if not cell or not cell.get("bytes"): | |
| raise InvalidZebraRow("missing_image", f"missing bytes for {key}") | |
| return cell["bytes"] | |
| def _reference_keys(text: str) -> list[str]: | |
| return [event["key"] for event in split_image_references(text) if event["type"] == "image"] | |
| def _group_and_split( | |
| config: str, | |
| question: str, | |
| problem_payloads: list[bytes], | |
| *, | |
| test_fraction: float, | |
| validation_fraction: float, | |
| ) -> tuple[str, str, float]: | |
| group = hashlib.sha256() | |
| group.update(config.encode()) | |
| group.update(b"\0") | |
| group.update(_normalize_question(question).encode()) | |
| for payload in problem_payloads: | |
| group.update(b"\0") | |
| group.update(hashlib.sha256(payload).digest()) | |
| group_id = group.hexdigest() | |
| score = int(group_id[:16], 16) / 2**64 | |
| if score < test_fraction: | |
| split = "test" | |
| elif score < test_fraction + validation_fraction: | |
| split = "validation" | |
| else: | |
| split = "train" | |
| return group_id, split, score | |
| def iter_rows( | |
| source_root: Path, | |
| graph_root: Path, | |
| *, | |
| rank: int = 0, | |
| world_size: int = 1, | |
| batch_size: int = 8, | |
| max_source_shards: int | None = None, | |
| owned_shard_indices: set[int] | None = None, | |
| ) -> Iterator[tuple[int, Path, int, dict]]: | |
| """Stream rank-owned source rows and apply the verified Graph overlay.""" | |
| shards = source_shards(source_root) | |
| if max_source_shards is not None: | |
| shards = shards[:max_source_shards] | |
| for shard_index, path in enumerate(shards): | |
| owned = ( | |
| shard_index in owned_shard_indices | |
| if owned_shard_indices is not None | |
| else shard_index % world_size == rank | |
| ) | |
| if not owned: | |
| continue | |
| parquet = pq.ParquetFile(path) | |
| is_graph = path.parent.name == GRAPH_CONFIG | |
| graph_complete = all( | |
| f"reasoning_image_{index}" in parquet.schema_arrow.names | |
| for index in range(10, 20) | |
| ) | |
| overlay = None | |
| if is_graph and not graph_complete: | |
| overlay_path = graph_root / "data" / path.name | |
| if not overlay_path.exists(): | |
| raise FileNotFoundError(f"Graph overlay missing: {overlay_path}") | |
| overlay = pq.ParquetFile(overlay_path) | |
| if overlay.metadata.num_rows != parquet.metadata.num_rows: | |
| raise ValueError(f"Graph overlay row mismatch: {path}") | |
| source_batches = parquet.iter_batches(batch_size=batch_size) | |
| overlay_batches = overlay.iter_batches(batch_size=batch_size) if overlay else None | |
| row_index = 0 | |
| while True: | |
| try: | |
| batch = next(source_batches) | |
| except StopIteration: | |
| break | |
| extra = next(overlay_batches) if overlay_batches is not None else None | |
| source_rows = batch.to_pylist() | |
| overlay_rows = extra.to_pylist() if extra is not None else [None] * len(source_rows) | |
| if len(source_rows) != len(overlay_rows): | |
| raise ValueError(f"Graph overlay batch mismatch: {path}") | |
| for row, graph_row in zip(source_rows, overlay_rows): | |
| if graph_row is not None: | |
| for column in TEXT_COLUMNS: | |
| if row[column] != graph_row[column]: | |
| raise ValueError( | |
| f"Graph overlay text mismatch: {path}:{row_index}:{column}" | |
| ) | |
| for key, value in graph_row.items(): | |
| if key.startswith("reasoning_image_") and key not in row: | |
| row[key] = value | |
| yield shard_index, path, row_index, row | |
| row_index += 1 | |
| if overlay_batches is not None: | |
| try: | |
| next(overlay_batches) | |
| except StopIteration: | |
| pass | |
| else: | |
| raise ValueError(f"Graph overlay has extra rows: {path}") | |
| def validate_row( | |
| *, | |
| config: str, | |
| row: dict, | |
| test_fraction: float, | |
| validation_fraction: float, | |
| ) -> dict: | |
| """Validate roles/references and return ordered raw image payloads.""" | |
| for column in TEXT_COLUMNS: | |
| if not isinstance(row.get(column), str) or not row[column].strip(): | |
| raise InvalidZebraRow("empty_text", f"empty {column}") | |
| problem_keys = _reference_keys(row["Question"]) | |
| reasoning_keys = _reference_keys(row["Text Reasoning Trace"]) | |
| # Zebra legitimately contains text-only questions whose assistant trajectory | |
| # introduces the first visual state. Only the reasoning-image target is | |
| # required; zero problem/observation images is a valid text-to-image prefix. | |
| if not reasoning_keys: | |
| raise InvalidZebraRow("missing_reasoning_reference", "trace lacks a reasoning image") | |
| if any(not key.startswith("problem_image_") for key in problem_keys): | |
| raise InvalidZebraRow("question_role_anomaly", f"question refs={problem_keys}") | |
| if any( | |
| not (key.startswith("reasoning_image_") or key.startswith("problem_image_")) | |
| for key in reasoning_keys | |
| ): | |
| raise InvalidZebraRow("unknown_trace_image_role", f"trace refs={reasoning_keys}") | |
| if len(set(problem_keys + reasoning_keys)) != len(problem_keys + reasoning_keys): | |
| raise InvalidZebraRow("duplicate_reference", "duplicate image reference") | |
| problem_payloads = [_image_bytes(row.get(key), key) for key in problem_keys] | |
| reasoning_payloads = [_image_bytes(row.get(key), key) for key in reasoning_keys] | |
| group_id, split, split_score = _group_and_split( | |
| config, | |
| row["Question"], | |
| problem_payloads, | |
| test_fraction=test_fraction, | |
| validation_fraction=validation_fraction, | |
| ) | |
| return { | |
| "problem_keys": problem_keys, | |
| "reasoning_keys": reasoning_keys, | |
| "keys": problem_keys + reasoning_keys, | |
| "payloads": problem_payloads + reasoning_payloads, | |
| "roles": ["observation"] * len(problem_keys) | |
| + ["generated_state"] * len(reasoning_keys), | |
| # Eleven released rows place an assistant-generated diagram in a | |
| # problem_image_* column. Turn position is authoritative for semantic | |
| # role; preserve the source mismatch for audit instead of dropping it. | |
| "trace_source_role_mismatches": [ | |
| key for key in reasoning_keys if key.startswith("problem_image_") | |
| ], | |
| "group_id": group_id, | |
| "split": split, | |
| "split_score": split_score, | |
| } | |
| def _decode_images(payloads: list[bytes]) -> list[Image.Image]: | |
| images = [] | |
| for index, payload in enumerate(payloads): | |
| try: | |
| with Image.open(io.BytesIO(payload)) as image: | |
| images.append(image.convert("RGB").copy()) | |
| except Exception as error: | |
| raise InvalidZebraRow("undecodable_image", f"image {index}: {error}") from error | |
| return images | |
| def _cap_image(image: Image.Image, max_long_edge: int) -> Image.Image: | |
| width, height = image.size | |
| scale = min(1.0, max_long_edge / max(width, height)) | |
| if scale == 1.0: | |
| return image | |
| output = (max(1, round(width * scale)), max(1, round(height * scale))) | |
| return image.resize(output, Image.Resampling.BICUBIC) | |
| def _letterbox_square( | |
| image: Image.Image, | |
| *, | |
| size: int = 512, | |
| background: int = 128, | |
| ) -> Image.Image: | |
| """Fit an image to a centered square canvas without changing its aspect ratio.""" | |
| if size <= 0 or not 0 <= background <= 255: | |
| raise ValueError("letterbox size/background are out of range") | |
| image = image.convert("RGB") | |
| width, height = image.size | |
| scale = size / max(width, height) | |
| resized = image.resize( | |
| (max(1, round(width * scale)), max(1, round(height * scale))), | |
| Image.Resampling.BICUBIC, | |
| ) | |
| canvas = Image.new("RGB", (size, size), (background,) * 3) | |
| canvas.paste( | |
| resized, | |
| ((size - resized.width) // 2, (size - resized.height) // 2), | |
| ) | |
| return canvas | |
| def _understanding_grid16_images( | |
| originals: list[Image.Image], | |
| keys: list[str], | |
| reasoning_keys: list[str], | |
| *, | |
| target_canvas_size: int, | |
| target_background: int, | |
| ) -> list[Image.Image]: | |
| """Keep user images on the checkpoint processor path and letterbox GT images.""" | |
| targets = set(reasoning_keys) | |
| return [ | |
| _letterbox_square( | |
| image, | |
| size=target_canvas_size, | |
| background=target_background, | |
| ) | |
| if key in targets | |
| else image.convert("RGB") | |
| for key, image in zip(keys, originals, strict=True) | |
| ] | |
| def _qwen_vision_images(images: list[Image.Image]) -> list[Image.Image]: | |
| """Apply the same qwen-vl-utils resize stage used by released UniAR paths.""" | |
| from qwen_vl_utils import process_vision_info | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [{"type": "image", "image": image} for image in images], | |
| } | |
| ] | |
| processed, videos = process_vision_info(messages) | |
| if videos: | |
| raise ValueError("unexpected video output while preparing Zebra images") | |
| return processed | |
| def _expand_image_tokens(ids: list[int], image_token_id: int, counts: list[int]) -> list[int]: | |
| output: list[int] = [] | |
| cursor = 0 | |
| for token_id in ids: | |
| if token_id == image_token_id: | |
| if cursor >= len(counts): | |
| raise ValueError("more image placeholders than image grids") | |
| output.extend([token_id] * counts[cursor]) | |
| cursor += 1 | |
| else: | |
| output.append(token_id) | |
| if cursor != len(counts): | |
| raise ValueError(f"unused image grids: used={cursor}, total={len(counts)}") | |
| return output | |
| def render_cached_inputs( | |
| processor, | |
| trajectory: dict, | |
| image_grid_thw: torch.Tensor, | |
| image_roles: list[str], | |
| image_token_id: int, | |
| ) -> dict: | |
| """Render a mixed sequence from grids without reprocessing RGB pixels.""" | |
| messages, rendered_roles = render_messages(trajectory, lambda path: path) | |
| if rendered_roles != image_roles: | |
| raise ValueError(f"image role mismatch: {rendered_roles} != {image_roles}") | |
| full_text = processor.apply_chat_template( | |
| messages, tokenize=False, add_generation_prompt=False | |
| ) | |
| prefix_text = processor.apply_chat_template( | |
| messages[:-1], tokenize=False, add_generation_prompt=True | |
| ) | |
| merge = processor.image_processor.merge_size | |
| counts = (image_grid_thw.prod(-1) // (merge * merge)).tolist() | |
| observation_count = image_roles.count("observation") | |
| tokenizer = processor.tokenizer | |
| full_ids = _expand_image_tokens( | |
| tokenizer.encode(full_text, add_special_tokens=False), image_token_id, counts | |
| ) | |
| prefix_ids = _expand_image_tokens( | |
| tokenizer.encode(prefix_text, add_special_tokens=False), | |
| image_token_id, | |
| counts[:observation_count], | |
| ) | |
| if full_ids[: len(prefix_ids)] != prefix_ids: | |
| raise ValueError("rendered sequence does not share its expected ChatML prefix") | |
| input_ids = torch.tensor(full_ids, dtype=torch.int32) | |
| labels = input_ids.clone() | |
| labels[: len(prefix_ids)] = -100 | |
| labels = canonical_autonomous_text_labels(input_ids, labels) | |
| image_positions = torch.nonzero(input_ids == image_token_id, as_tuple=False).flatten() | |
| if image_positions.numel() != sum(counts): | |
| raise ValueError("expanded image-token count mismatch") | |
| generated_sequence_mask = torch.zeros(input_ids.shape, dtype=torch.bool) | |
| generated_flat_mask = torch.zeros(sum(counts), dtype=torch.bool) | |
| generated_image_token_counts: list[int] = [] | |
| cursor = 0 | |
| for role, count in zip(image_roles, counts): | |
| end = cursor + count | |
| if role == "generated_state": | |
| generated_sequence_mask[image_positions[cursor:end]] = True | |
| generated_flat_mask[cursor:end] = True | |
| generated_image_token_counts.append(count) | |
| cursor = end | |
| return { | |
| "input_ids": input_ids, | |
| "labels": labels, | |
| "image_grid_thw": image_grid_thw.to(torch.int16), | |
| "generated_sequence_mask": generated_sequence_mask, | |
| "generated_flat_mask": generated_flat_mask, | |
| "generated_image_token_counts": torch.tensor( | |
| generated_image_token_counts, dtype=torch.int16 | |
| ), | |
| "image_roles": image_roles, | |
| "prefix_length": len(prefix_ids), | |
| "rendered_text": full_text, | |
| } | |
| def supervised_generation_marker_count(item: dict, marker_id: int) -> int: | |
| """Count generation markers expected to supervise generated images.""" | |
| return int((item["labels"] == marker_id).sum()) | |
| def cached_batch( | |
| item: dict, *, code_rows: torch.Tensor | None = None | |
| ) -> tuple[dict, torch.Tensor]: | |
| """Convert one serialized cache item to the joint-objective contract.""" | |
| input_ids_1d = item["input_ids"].to(torch.long) | |
| input_ids = input_ids_1d.unsqueeze(0) | |
| batch = { | |
| "sample_id": item["sample_id"], | |
| "input_ids": input_ids, | |
| "attention_mask": torch.ones_like(input_ids), | |
| "labels": canonical_autonomous_text_labels( | |
| input_ids_1d, item["labels"].to(torch.long) | |
| ).unsqueeze(0), | |
| "image_grid_thw": item["image_grid_thw"].to(torch.long), | |
| "generated_sequence_mask": item["generated_sequence_mask"].unsqueeze(0), | |
| "generated_flat_mask": item["generated_flat_mask"], | |
| "generated_image_token_counts": item["generated_image_token_counts"].tolist(), | |
| "image_roles": item["image_roles"], | |
| } | |
| if code_rows is None: | |
| codes = unpack_visual_codes( | |
| item["visual_codes_packed"], item["visual_codes_shape"] | |
| ) | |
| else: | |
| code_rows = code_rows.to(torch.bool) | |
| if code_rows.numel() != int(item["visual_codes_shape"][0]): | |
| raise ValueError("visual code row selection has the wrong length") | |
| selected_shape = [int(code_rows.sum()), *item["visual_codes_shape"][1:]] | |
| codes = unpack_visual_codes( | |
| item["visual_codes_packed"][code_rows], selected_shape | |
| ) | |
| return batch, codes | |
| def collate_cached_items( | |
| items: list[dict], pad_token_id: int | |
| ) -> tuple[dict, torch.Tensor]: | |
| """Right-pad cached trajectories while preserving flattened image order.""" | |
| if not items: | |
| raise ValueError("cannot collate an empty cache batch") | |
| prepared = [cached_batch(item) for item in items] | |
| max_length = max(batch["input_ids"].shape[1] for batch, _ in prepared) | |
| batch_size = len(items) | |
| input_ids = torch.full( | |
| (batch_size, max_length), pad_token_id, dtype=torch.long | |
| ) | |
| attention_mask = torch.zeros((batch_size, max_length), dtype=torch.long) | |
| labels = torch.full((batch_size, max_length), -100, dtype=torch.long) | |
| generated_sequence_mask = torch.zeros( | |
| (batch_size, max_length), dtype=torch.bool | |
| ) | |
| for index, (batch, _) in enumerate(prepared): | |
| length = batch["input_ids"].shape[1] | |
| input_ids[index, :length] = batch["input_ids"][0] | |
| attention_mask[index, :length] = 1 | |
| labels[index, :length] = batch["labels"][0] | |
| generated_sequence_mask[index, :length] = batch[ | |
| "generated_sequence_mask" | |
| ][0] | |
| batch = { | |
| "sample_id": [item["sample_id"] for item in items], | |
| "input_ids": input_ids, | |
| "attention_mask": attention_mask, | |
| "labels": labels, | |
| "image_grid_thw": torch.cat( | |
| [batch["image_grid_thw"] for batch, _ in prepared] | |
| ), | |
| "generated_sequence_mask": generated_sequence_mask, | |
| "generated_flat_mask": torch.cat( | |
| [batch["generated_flat_mask"] for batch, _ in prepared] | |
| ), | |
| } | |
| codes = torch.cat([codes for _, codes in prepared]) | |
| return batch, codes | |
| def prepare_row( | |
| *, | |
| ar_model, | |
| processor, | |
| config: str, | |
| source_shard: str, | |
| source_row: int, | |
| row: dict, | |
| max_long_edge: int = 960, | |
| max_sequence_length: int = 32768, | |
| test_fraction: float = 0.01, | |
| validation_fraction: float = 0.0025, | |
| max_vision_patch_tokens: int = 4096, | |
| image_token_id: int | None = None, | |
| image_preprocessing_profile: str = LEGACY_IMAGE_PREPROCESSING, | |
| target_canvas_size: int = 512, | |
| target_background: int = 128, | |
| ) -> dict: | |
| """Validate, resize, encode, and serialize one full-corpus trajectory.""" | |
| validated = validate_row( | |
| config=config, | |
| row=row, | |
| test_fraction=test_fraction, | |
| validation_fraction=validation_fraction, | |
| ) | |
| originals = _decode_images(validated["payloads"]) | |
| marker_id = image_gen_token_id(processor) | |
| image_token_id = ( | |
| ar_model.config.image_token_id if image_token_id is None else image_token_id | |
| ) | |
| selected = None | |
| selected_pixels = None | |
| selected_trajectory = None | |
| if image_preprocessing_profile not in IMAGE_PREPROCESSING_PROFILES: | |
| raise ValueError( | |
| f"unknown image preprocessing profile: {image_preprocessing_profile!r}" | |
| ) | |
| if image_preprocessing_profile == LEGACY_IMAGE_PREPROCESSING: | |
| candidates = ( | |
| ( | |
| cap, | |
| [_cap_image(image, cap) for image in originals], | |
| ) | |
| for cap in CAP_CANDIDATES | |
| if cap <= max_long_edge | |
| ) | |
| else: | |
| candidates = iter( | |
| [ | |
| ( | |
| None, | |
| _understanding_grid16_images( | |
| originals, | |
| validated["keys"], | |
| validated["reasoning_keys"], | |
| target_canvas_size=target_canvas_size, | |
| target_background=target_background, | |
| ), | |
| ) | |
| ] | |
| ) | |
| for cap, profile_images in candidates: | |
| model_inputs = ( | |
| _qwen_vision_images(profile_images) | |
| if image_preprocessing_profile == LEGACY_IMAGE_PREPROCESSING | |
| else profile_images | |
| ) | |
| processed = processor.image_processor(images=model_inputs, return_tensors="pt") | |
| grids = processed.image_grid_thw.to(torch.long) | |
| patch = processor.image_processor.patch_size | |
| image_sizes = { | |
| key: (int(grid[2]) * patch, int(grid[1]) * patch) | |
| for key, grid in zip(validated["keys"], grids) | |
| } | |
| sample_seed = hashlib.sha256( | |
| f"{config}\0{source_shard}\0{source_row}".encode() | |
| ).hexdigest() | |
| trajectory = build_trajectory( | |
| sample_id=sample_seed[:24], | |
| config=config, | |
| source_shard=source_shard, | |
| source_row=source_row, | |
| question=row["Question"], | |
| reasoning_trace=row["Text Reasoning Trace"], | |
| final_answer=row["Final Answer"], | |
| image_paths={key: key for key in validated["keys"]}, | |
| image_sizes=image_sizes, | |
| ) | |
| rendered = render_cached_inputs( | |
| processor, | |
| trajectory, | |
| grids, | |
| validated["roles"], | |
| image_token_id, | |
| ) | |
| if rendered["input_ids"].numel() <= max_sequence_length: | |
| selected = (cap, rendered, grids, image_sizes, profile_images, model_inputs) | |
| selected_pixels = processed.pixel_values | |
| selected_trajectory = trajectory | |
| break | |
| if selected is None or selected_pixels is None or selected_trajectory is None: | |
| raise InvalidZebraRow( | |
| "context_overflow", | |
| f"sequence exceeds {max_sequence_length} even at {CAP_CANDIDATES[-1]}px", | |
| ) | |
| cap, rendered, grids, image_sizes, profile_images, model_inputs = selected | |
| code_parts = [] | |
| grid_parts = [] | |
| pixel_parts = [] | |
| patch_cursor = 0 | |
| patch_budget = 0 | |
| def flush() -> None: | |
| nonlocal patch_budget | |
| if not grid_parts: | |
| return | |
| chunk_grids = torch.stack(grid_parts) | |
| chunk_pixels = torch.cat(pixel_parts) | |
| code_parts.append(encode_visual_codes(ar_model, chunk_pixels, chunk_grids).cpu()) | |
| grid_parts.clear() | |
| pixel_parts.clear() | |
| patch_budget = 0 | |
| for grid in grids: | |
| patches = int(grid.prod()) | |
| if grid_parts and patch_budget + patches > max_vision_patch_tokens: | |
| flush() | |
| grid_parts.append(grid) | |
| pixel_parts.append(selected_pixels[patch_cursor : patch_cursor + patches]) | |
| patch_cursor += patches | |
| patch_budget += patches | |
| flush() | |
| if patch_cursor != selected_pixels.shape[0]: | |
| raise ValueError("pixel/grid patch alignment mismatch") | |
| visual_codes = torch.cat(code_parts) | |
| expected_visual_positions = int(grids.prod(-1).sum() // 4) | |
| if visual_codes.shape[0] != expected_visual_positions: | |
| raise ValueError( | |
| f"BSQ/grid mismatch: codes={visual_codes.shape[0]}, expected={expected_visual_positions}" | |
| ) | |
| packed, code_shape = pack_visual_codes(visual_codes) | |
| item = { | |
| "schema": CACHE_SCHEMA, | |
| "sample_id": selected_trajectory["sample_id"], | |
| # This threshold split is only a rank-local hint. prepare_full assigns | |
| # exact bounded per-config group splits in the canonical manifest after | |
| # every rank has finished. | |
| "split_hint": validated["split"], | |
| "group_id": validated["group_id"], | |
| "split_score": validated["split_score"], | |
| "config": config, | |
| "source": { | |
| "shard": source_shard, | |
| "row": source_row, | |
| }, | |
| "source_anomalies": { | |
| "trace_problem_image_columns": validated[ | |
| "trace_source_role_mismatches" | |
| ] | |
| }, | |
| "trajectory": selected_trajectory, | |
| **{key: value for key, value in rendered.items() if key != "rendered_text"}, | |
| "visual_codes_packed": packed, | |
| "visual_codes_shape": code_shape, | |
| "preprocessing": { | |
| "profile": image_preprocessing_profile, | |
| "max_long_edge": max_long_edge, | |
| "selected_cap": cap, | |
| "understanding_processor": { | |
| "min_pixels": getattr(processor.image_processor, "min_pixels", None), | |
| "max_pixels": getattr(processor.image_processor, "max_pixels", None), | |
| }, | |
| "target_canvas": ( | |
| { | |
| "size": target_canvas_size, | |
| "background_rgb": [target_background] * 3, | |
| "fit": "longest_edge", | |
| "placement": "center", | |
| } | |
| if image_preprocessing_profile == UNDERSTANDING_INPUT_GRID16_TARGET | |
| else None | |
| ), | |
| "source_sizes": { | |
| key: image.size for key, image in zip(validated["keys"], originals) | |
| }, | |
| "capped_sizes": { | |
| key: image.size | |
| for key, image in zip(validated["keys"], profile_images) | |
| }, | |
| "qwen_resized_sizes": { | |
| key: image.size for key, image in zip(validated["keys"], model_inputs) | |
| }, | |
| "model_ready_sizes": image_sizes, | |
| "patch_size": processor.image_processor.patch_size, | |
| "merge_size": processor.image_processor.merge_size, | |
| }, | |
| "target_render": symbolic_render(selected_trajectory), | |
| } | |
| marker_count = supervised_generation_marker_count(item, marker_id) | |
| if marker_count != len(validated["reasoning_keys"]): | |
| raise ValueError( | |
| f"generation marker mismatch: markers={marker_count}, images={len(validated['reasoning_keys'])}" | |
| ) | |
| return item | |
| def manifest_row(item: dict, cache_shard: str, cache_index: int) -> dict: | |
| """Compact index entry used for split audits and held-out sampling.""" | |
| return { | |
| "sample_id": item["sample_id"], | |
| "split": "unassigned", | |
| "group_id": item["group_id"], | |
| "split_score": item["split_score"], | |
| "config": item["config"], | |
| "source": item["source"], | |
| "cache_shard": cache_shard, | |
| "cache_index": cache_index, | |
| "sequence_length": int(item["input_ids"].numel()), | |
| "images": len(item["image_roles"]), | |
| "generated_images": item["image_roles"].count("generated_state"), | |
| "visual_positions": int(item["image_grid_thw"].prod(-1).sum() // 4), | |
| "selected_cap": item["preprocessing"]["selected_cap"], | |
| "trace_problem_image_columns": item["source_anomalies"][ | |
| "trace_problem_image_columns" | |
| ], | |
| } | |
| def read_manifest(path: Path) -> list[dict]: | |
| with path.open(encoding="utf-8") as handle: | |
| return [json.loads(line) for line in handle if line.strip()] | |
| def _read_parquet_row(path: Path, row_index: int) -> dict: | |
| parquet = pq.ParquetFile(path) | |
| offset = 0 | |
| for group in range(parquet.num_row_groups): | |
| rows = parquet.metadata.row_group(group).num_rows | |
| if row_index < offset + rows: | |
| return parquet.read_row_group(group).slice(row_index - offset, 1).to_pylist()[0] | |
| offset += rows | |
| raise IndexError(f"row {row_index} outside {path} ({offset} rows)") | |
| def load_item_images( | |
| item: dict, source_root: Path, graph_root: Path | |
| ) -> dict[str, Image.Image]: | |
| """Recover the exact capped/qwen-resized RGB inputs for inspection.""" | |
| path = source_root / item["source"]["shard"] | |
| row = _read_parquet_row(path, item["source"]["row"]) | |
| if item["config"] == GRAPH_CONFIG: | |
| graph_row = _read_parquet_row(graph_root / "data" / path.name, item["source"]["row"]) | |
| for key, value in graph_row.items(): | |
| if key.startswith("reasoning_image_") and key not in row: | |
| row[key] = value | |
| keys = [ | |
| event["key"] | |
| for turn in (item["trajectory"]["user"], item["trajectory"]["assistant"]) | |
| for event in turn | |
| if event["type"] == "image" | |
| ] | |
| originals = _decode_images([_image_bytes(row.get(key), key) for key in keys]) | |
| preprocessing = item["preprocessing"] | |
| profile = preprocessing.get("profile", LEGACY_IMAGE_PREPROCESSING) | |
| if profile == LEGACY_IMAGE_PREPROCESSING: | |
| cap = preprocessing["selected_cap"] | |
| processed = _qwen_vision_images([_cap_image(image, cap) for image in originals]) | |
| elif profile == UNDERSTANDING_INPUT_GRID16_TARGET: | |
| canvas = preprocessing["target_canvas"] | |
| processed = _understanding_grid16_images( | |
| originals, | |
| keys, | |
| [ | |
| event["key"] | |
| for event in item["trajectory"]["assistant"] | |
| if event["type"] == "image" | |
| ], | |
| target_canvas_size=int(canvas["size"]), | |
| target_background=int(canvas["background_rgb"][0]), | |
| ) | |
| else: | |
| raise ValueError(f"unknown cached image preprocessing profile: {profile!r}") | |
| return dict(zip(keys, processed)) | |
Xet Storage Details
- Size:
- 28.7 kB
- Xet hash:
- 074fd9444c0824b2ee698c25221a0b931616327f9d0a37f656e5874047fce7d3
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.