Buckets:
| """Canonical Zebra-CoT trajectory parsing and UniAR chat rendering.""" | |
| from __future__ import annotations | |
| import re | |
| from pathlib import Path | |
| from typing import Callable | |
| from .protocol import SYSTEM_PROMPT | |
| IMAGE_REFERENCE = re.compile(r"<image_start>\[([^\]]+)\]<image_end>") | |
| THOUGHT_LABEL = re.compile(r"THOUGHT\s*(?:\d+|N)\s*:\s*", re.IGNORECASE) | |
| ANSWER_PREFIX = re.compile(r"^\s*FINAL\s+ANSWER\s*:\s*", re.IGNORECASE) | |
| def split_image_references(text: str) -> list[dict]: | |
| """Split Zebra placeholder text into ordered text/image events.""" | |
| events: list[dict] = [] | |
| cursor = 0 | |
| for match in IMAGE_REFERENCE.finditer(text): | |
| if match.start() > cursor: | |
| events.append({"type": "text", "text": text[cursor : match.start()]}) | |
| events.append({"type": "image", "key": match.group(1)}) | |
| cursor = match.end() | |
| if cursor < len(text): | |
| events.append({"type": "text", "text": text[cursor:]}) | |
| return events | |
| def clean_thought(text: str) -> str: | |
| """Remove Zebra step labels without changing reasoning content.""" | |
| return THOUGHT_LABEL.sub("", text).strip() | |
| def clean_answer(text: str) -> str: | |
| """Remove Zebra answer wrappers that would otherwise be duplicated.""" | |
| value = text.strip() | |
| wrapped = re.fullmatch(r"<answer>\s*(.*?)\s*</answer>", value, re.IGNORECASE | re.DOTALL) | |
| if wrapped: | |
| value = wrapped.group(1) | |
| return ANSWER_PREFIX.sub("", value).strip() | |
| def build_trajectory( | |
| *, | |
| sample_id: str, | |
| config: str, | |
| source_shard: str, | |
| source_row: int, | |
| question: str, | |
| reasoning_trace: str, | |
| final_answer: str, | |
| image_paths: dict[str, str], | |
| image_size: int | None = None, | |
| image_sizes: dict[str, tuple[int, int]] | None = None, | |
| downsample_factor: int = 32, | |
| ) -> dict: | |
| """Convert one validated Zebra row into the project trajectory contract.""" | |
| def dimensions(key: str) -> tuple[int, int]: | |
| if image_sizes is not None: | |
| width, height = image_sizes[key] | |
| elif image_size is not None: | |
| width = height = image_size | |
| else: | |
| raise ValueError("image_size or image_sizes is required") | |
| if width % downsample_factor or height % downsample_factor: | |
| raise ValueError( | |
| f"model-ready image dimensions must be divisible by {downsample_factor}: " | |
| f"{key}={width}x{height}" | |
| ) | |
| return width, height | |
| user: list[dict] = [] | |
| for event in split_image_references(question): | |
| if event["type"] == "text" and event["text"]: | |
| user.append({"type": "text", "text": event["text"]}) | |
| elif event["type"] == "image": | |
| key = event["key"] | |
| if not key.startswith("problem_image_") or key not in image_paths: | |
| raise ValueError(f"invalid problem-image reference: {key}") | |
| width, height = dimensions(key) | |
| user.append( | |
| { | |
| "type": "image", | |
| "role": "observation", | |
| "key": key, | |
| "path": image_paths[key], | |
| "height": height, | |
| "width": width, | |
| } | |
| ) | |
| assistant: list[dict] = [] | |
| for event in split_image_references(reasoning_trace): | |
| if event["type"] == "text": | |
| thought = clean_thought(event["text"]) | |
| if thought: | |
| assistant.append({"type": "text", "role": "thought", "text": thought}) | |
| else: | |
| key = event["key"] | |
| if key not in image_paths or not ( | |
| key.startswith("reasoning_image_") or key.startswith("problem_image_") | |
| ): | |
| raise ValueError(f"invalid reasoning-image reference: {key}") | |
| width, height = dimensions(key) | |
| assistant.append( | |
| { | |
| "type": "image", | |
| "role": "generated_state", | |
| "key": key, | |
| "path": image_paths[key], | |
| "height": height, | |
| "width": width, | |
| "grid_h": height // downsample_factor, | |
| "grid_w": width // downsample_factor, | |
| "source_column_role_mismatch": key.startswith("problem_image_"), | |
| } | |
| ) | |
| assistant.append( | |
| {"type": "text", "role": "answer", "text": clean_answer(final_answer)} | |
| ) | |
| return { | |
| "schema": "spatial_mmcot_trajectory_v1", | |
| "sample_id": sample_id, | |
| "source": { | |
| "dataset": "Zebra-CoT", | |
| "config": config, | |
| "shard": source_shard, | |
| "row": source_row, | |
| }, | |
| "system": SYSTEM_PROMPT, | |
| "user": user, | |
| "assistant": assistant, | |
| } | |
| def render_messages( | |
| trajectory: dict, | |
| load_image: Callable[[str], object], | |
| ) -> tuple[list[dict], list[str]]: | |
| """Render typed events as processor-compatible messages and image roles.""" | |
| messages: list[dict] = [ | |
| { | |
| "role": "system", | |
| "content": [{"type": "text", "text": trajectory["system"]}], | |
| } | |
| ] | |
| image_roles: list[str] = [] | |
| user_content: list[dict] = [] | |
| for event in trajectory["user"]: | |
| if event["type"] == "text": | |
| user_content.append({"type": "text", "text": event["text"]}) | |
| else: | |
| user_content.append({"type": "image", "image": load_image(event["path"])}) | |
| image_roles.append(event["role"]) | |
| messages.append({"role": "user", "content": user_content}) | |
| assistant_content: list[dict] = [] | |
| def append_text_block(text: str) -> None: | |
| assistant_content.append( | |
| {"type": "text", "text": ("\n" if assistant_content else "") + text} | |
| ) | |
| for event in trajectory["assistant"]: | |
| if event["role"] == "thought": | |
| append_text_block(f"<think>\n{event['text']}\n</think>") | |
| elif event["role"] == "generated_state": | |
| append_text_block( | |
| f"<image_gen> generate image {event['grid_h']} {event['grid_w']}" | |
| ) | |
| assistant_content.append({"type": "image", "image": load_image(event["path"])}) | |
| image_roles.append(event["role"]) | |
| elif event["role"] == "answer": | |
| append_text_block(f"Final answer: {event['text']}") | |
| else: | |
| raise ValueError(f"unknown assistant event role: {event['role']}") | |
| messages.append({"role": "assistant", "content": assistant_content}) | |
| return messages, image_roles | |
| def symbolic_render(trajectory: dict) -> str: | |
| """Human-readable trajectory rendering saved beside prepared samples.""" | |
| lines = ["<|im_start|>system", trajectory["system"], "<|im_end|>", "<|im_start|>user"] | |
| for event in trajectory["user"]: | |
| lines.append( | |
| event["text"] if event["type"] == "text" else f"<observation:{event['key']}>" | |
| ) | |
| lines.extend(["<|im_end|>", "<|im_start|>assistant"]) | |
| for event in trajectory["assistant"]: | |
| if event["role"] == "thought": | |
| lines.extend(["<think>", event["text"], "</think>"]) | |
| elif event["role"] == "generated_state": | |
| lines.append( | |
| f"<image_gen> generate image {event['grid_h']} {event['grid_w']}" | |
| f"<|vision_start|><generated_state:{event['key']}><|vision_end|>" | |
| ) | |
| else: | |
| lines.append(f"Final answer: {event['text']}") | |
| lines.append("<|im_end|>") | |
| return "\n".join(lines) + "\n" | |
| def resolve_image_path(samples_root: Path, relative_path: str) -> Path: | |
| """Resolve a prepared image while preventing paths outside the sample root.""" | |
| root = samples_root.resolve() | |
| path = (root / relative_path).resolve() | |
| if not path.is_relative_to(root): | |
| raise ValueError(f"image path escapes sample root: {relative_path}") | |
| return path | |
Xet Storage Details
- Size:
- 7.9 kB
- Xet hash:
- 9a9f6cec159be598e5f3489c837e6a7b7665b8244a20defe94dd7a8abba7b040
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.