Buckets:
| """Prepared Zebra trajectory loading and UniAR processor input construction.""" | |
| from __future__ import annotations | |
| import json | |
| from pathlib import Path | |
| import torch | |
| from PIL import Image | |
| from shared.labels import canonical_autonomous_text_labels | |
| from .trajectory import render_messages, resolve_image_path | |
| class ZebraTrajectoryDataset: | |
| """Small JSONL-backed dataset with stable sample ordering.""" | |
| def __init__(self, samples_root: str | Path, manifest: str = "prepared/trajectories.jsonl"): | |
| self.root = Path(samples_root).resolve() | |
| path = self.root / manifest | |
| with path.open(encoding="utf-8") as handle: | |
| self.records = [json.loads(line) for line in handle if line.strip()] | |
| if not self.records: | |
| raise ValueError(f"no trajectories found in {path}") | |
| def __len__(self) -> int: | |
| return len(self.records) | |
| def __getitem__(self, index: int) -> dict: | |
| return self.records[index] | |
| def _load_rgb(samples_root: Path, relative_path: str): | |
| with Image.open(resolve_image_path(samples_root, relative_path)) as image: | |
| return image.convert("RGB").copy() | |
| def prepare_processor_inputs(processor, trajectory: dict, samples_root: Path, image_token_id: int) -> dict: | |
| """Build one teacher-forced mixed-modality sequence on CPU.""" | |
| from qwen_vl_utils import process_vision_info | |
| messages, image_roles = render_messages( | |
| trajectory, lambda path: _load_rgb(samples_root, path) | |
| ) | |
| full_text = processor.apply_chat_template( | |
| messages, tokenize=False, add_generation_prompt=False | |
| ) | |
| prefix_messages = messages[:-1] | |
| prefix_text = processor.apply_chat_template( | |
| prefix_messages, tokenize=False, add_generation_prompt=True | |
| ) | |
| full_images, full_videos = process_vision_info(messages) | |
| prefix_images, prefix_videos = process_vision_info(prefix_messages) | |
| if full_videos or prefix_videos: | |
| raise ValueError("Zebra SFT fixture does not support videos") | |
| full = processor( | |
| text=[full_text], images=full_images, padding=False, return_tensors="pt" | |
| ) | |
| prefix = processor( | |
| text=[prefix_text], images=prefix_images, padding=False, return_tensors="pt" | |
| ) | |
| full.pop("mm_token_type_ids", None) | |
| prefix.pop("mm_token_type_ids", None) | |
| prefix_length = prefix.input_ids.shape[1] | |
| if not torch.equal(full.input_ids[:, :prefix_length], prefix.input_ids): | |
| raise ValueError("full sequence does not share the expected ChatML prefix") | |
| input_ids = full.input_ids | |
| labels = input_ids.clone() | |
| labels[:, :prefix_length] = -100 | |
| labels = canonical_autonomous_text_labels(input_ids, labels) | |
| labels[full.attention_mask == 0] = -100 | |
| merge = 2 | |
| token_counts = (full.image_grid_thw.prod(-1) // (merge * merge)).tolist() | |
| image_positions = torch.nonzero(input_ids[0] == image_token_id, as_tuple=False).flatten() | |
| if sum(token_counts) != image_positions.numel() or len(token_counts) != len(image_roles): | |
| raise ValueError( | |
| "image placeholder expansion mismatch: " | |
| f"tokens={image_positions.numel()} grids={token_counts} roles={image_roles}" | |
| ) | |
| generated_sequence_mask = torch.zeros_like(input_ids, dtype=torch.bool) | |
| generated_flat_mask = torch.zeros(image_positions.numel(), dtype=torch.bool) | |
| generated_image_token_counts: list[int] = [] | |
| cursor = 0 | |
| for role, count in zip(image_roles, token_counts): | |
| next_cursor = cursor + count | |
| if role == "generated_state": | |
| generated_sequence_mask[0, image_positions[cursor:next_cursor]] = True | |
| generated_flat_mask[cursor:next_cursor] = True | |
| generated_image_token_counts.append(count) | |
| cursor = next_cursor | |
| return { | |
| "sample_id": trajectory["sample_id"], | |
| "input_ids": input_ids, | |
| "attention_mask": full.attention_mask, | |
| "labels": labels, | |
| "pixel_values": full.pixel_values, | |
| "image_grid_thw": full.image_grid_thw, | |
| "generated_sequence_mask": generated_sequence_mask, | |
| "generated_flat_mask": generated_flat_mask, | |
| "image_roles": image_roles, | |
| "generated_image_token_counts": generated_image_token_counts, | |
| "rendered_text": full_text, | |
| "prefix_length": prefix_length, | |
| } | |
| def image_gen_token_id(processor) -> int: | |
| """Return the runtime marker ID and reject tokenizer ambiguity.""" | |
| from uniar import SPECIAL_TOKEN_IDS | |
| ids = processor.tokenizer.encode("<image_gen>", add_special_tokens=False) | |
| expected = [SPECIAL_TOKEN_IDS["<image_gen>"]] | |
| if ids != expected: | |
| raise ValueError(f"<image_gen> must encode as {expected}, got {ids}") | |
| return ids[0] | |
Xet Storage Details
- Size:
- 4.72 kB
- Xet hash:
- 5a553775bc0a376acfc72bbe7ceb3c9bd554aae3cc5c56af2e5c094b28014bf4
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.