Buckets:
| """Render assistant token pieces directly above their one-to-one token IDs.""" | |
| from __future__ import annotations | |
| import argparse | |
| import difflib | |
| import json | |
| import re | |
| from pathlib import Path | |
| from transformers import AutoProcessor | |
| from uniar import configure_processor | |
| from .dataset import ZebraTrajectoryDataset, prepare_processor_inputs | |
| from .trajectory import IMAGE_REFERENCE | |
| def visible_piece(tokenizer, token_id: int) -> str: | |
| """Decode one ID and make whitespace visible without losing content.""" | |
| piece = tokenizer.decode( | |
| [token_id], skip_special_tokens=False, clean_up_tokenization_spaces=False | |
| ) | |
| if not piece: | |
| piece = tokenizer.convert_ids_to_tokens(token_id) | |
| return ( | |
| piece.replace("\\", "\\\\") | |
| .replace("\r", "↵") | |
| .replace("\n", "↵") | |
| .replace("\t", "⇥") | |
| .replace(" ", "·") | |
| ) | |
| def render_alignment( | |
| tokenizer, | |
| token_ids: list[int], | |
| *, | |
| columns: int = 6, | |
| raw_labels: list[str] | None = None, | |
| ) -> str: | |
| """Return fixed-width paired TOK/ID rows with exact column alignment.""" | |
| lines = [ | |
| "Legend: · = space, ↵ = newline, ⇥ = tab", | |
| "Blank RAW cell = inserted by preprocessing/controller, not present in Zebra-CoT", | |
| "IMG[n] = <image_start>[reasoning_image_n]<image_end>", | |
| f"Token count: {len(token_ids)}", | |
| "", | |
| ] | |
| for start in range(0, len(token_ids), columns): | |
| ids = token_ids[start : start + columns] | |
| pieces = [visible_piece(tokenizer, token_id) for token_id in ids] | |
| positions = [str(index) for index in range(start, start + len(ids))] | |
| raw = raw_labels[start : start + len(ids)] if raw_labels is not None else None | |
| widths = [] | |
| for index, (position, piece, token_id) in enumerate(zip(positions, pieces, ids)): | |
| values = [len(position), len(piece), len(str(token_id))] | |
| if raw is not None: | |
| values.append(len(raw[index])) | |
| widths.append(max(values)) | |
| def row(label: str, values: list[str]) -> str: | |
| cells = [value.ljust(width) for value, width in zip(values, widths)] | |
| return f"{label:<4}│ " + " │ ".join(cells) + " │" | |
| lines.append(row("POS", positions)) | |
| if raw is not None: | |
| lines.append(row("RAW", raw)) | |
| lines.append(row("TOK", pieces)) | |
| lines.append(row("ID", [str(token_id) for token_id in ids])) | |
| lines.append("────┼─" + "─┼─".join("─" * width for width in widths) + "─┼") | |
| return "\n".join(lines) + "\n" | |
| def _raw_record(samples_root: Path, sample_id: str) -> dict: | |
| with (samples_root / "raw" / "records.jsonl").open(encoding="utf-8") as handle: | |
| records = [json.loads(line) for line in handle] | |
| matches = [record for record in records if record["sample_id"] == sample_id] | |
| if len(matches) != 1: | |
| raise ValueError(f"expected one raw record for {sample_id}, found {len(matches)}") | |
| return matches[0] | |
| def training_raw_labels(tokenizer, token_ids: list[int], raw_record: dict) -> list[str]: | |
| """Map rendered training IDs back to Zebra trace fields and placeholders.""" | |
| labels = [""] * len(token_ids) | |
| text_chunks = re.split(IMAGE_REFERENCE, raw_record["Text Reasoning Trace"])[::2] | |
| thought_groups = [ | |
| ",".join(re.findall(r"THOUGHT\s*(\d+)\s*:", chunk, flags=re.IGNORECASE)) | |
| for chunk in text_chunks | |
| if chunk.strip() | |
| ] | |
| image_keys = IMAGE_REFERENCE.findall(raw_record["Text Reasoning Trace"]) | |
| thought_index = 0 | |
| image_index = 0 | |
| cursor = 0 | |
| last_think_end = -1 | |
| while cursor < len(token_ids): | |
| token_id = token_ids[cursor] | |
| if token_id == 151667: | |
| end = token_ids.index(151668, cursor + 1) | |
| group = thought_groups[thought_index] if thought_index < len(thought_groups) else "?" | |
| labels[cursor] = f"THOUGHT·{group}:" | |
| for index in range(cursor + 1, end): | |
| piece = visible_piece(tokenizer, token_ids[index]) | |
| labels[index] = "" if index == cursor + 1 and piece == "↵" else piece | |
| last_think_end = end | |
| thought_index += 1 | |
| cursor = end + 1 | |
| continue | |
| if token_id == 151669: | |
| vision_start = token_ids.index(151652, cursor + 1) | |
| vision_end = token_ids.index(151653, vision_start + 1) | |
| pad_positions = [ | |
| index | |
| for index in range(vision_start + 1, vision_end) | |
| if token_ids[index] == 151655 | |
| ] | |
| if pad_positions: | |
| labels[pad_positions[0]] = f"IMG[{image_index + 1}]" | |
| for index in pad_positions[1:]: | |
| labels[index] = "↳" | |
| image_index += 1 | |
| cursor = vision_end + 1 | |
| continue | |
| cursor += 1 | |
| # Final Answer is a separate Zebra field, not part of Text Reasoning Trace. | |
| in_answer = False | |
| for index in range(last_think_end + 1, len(token_ids)): | |
| piece = visible_piece(tokenizer, token_ids[index]) | |
| if piece == "Final": | |
| labels[index] = "Final·Answer(field)" | |
| elif piece == ":": | |
| in_answer = True | |
| elif in_answer and labels[index] == "" and piece not in {"·", "<|im_end|>", "↵"}: | |
| labels[index] = piece | |
| return labels | |
| def align_test_raw_labels( | |
| training_ids: list[int], training_labels: list[str], test_ids: list[int] | |
| ) -> list[str]: | |
| """Copy provenance through equal spans; novel generated test tokens stay blank.""" | |
| labels = [""] * len(test_ids) | |
| matcher = difflib.SequenceMatcher(a=training_ids, b=test_ids, autojunk=False) | |
| for tag, left_start, left_end, right_start, right_end in matcher.get_opcodes(): | |
| if tag == "equal": | |
| labels[right_start:right_end] = training_labels[left_start:left_end] | |
| return labels | |
| def load_ids_and_raw(args, processor) -> tuple[list[int], list[str]]: | |
| dataset = ZebraTrajectoryDataset(args.samples_root) | |
| matches = [record for record in dataset.records if record["sample_id"] == args.sample_id] | |
| if len(matches) != 1: | |
| raise ValueError(f"expected one record for {args.sample_id}, found {len(matches)}") | |
| if args.mode == "training": | |
| batch = prepare_processor_inputs( | |
| processor, | |
| matches[0], | |
| dataset.root, | |
| image_token_id=args.image_token_id, | |
| ) | |
| ids = batch["input_ids"][0, batch["prefix_length"] :].tolist() | |
| return ids, training_raw_labels( | |
| processor.tokenizer, ids, _raw_record(dataset.root, args.sample_id) | |
| ) | |
| result_path = args.evaluation_root / args.sample_id / "result.json" | |
| test_ids = json.loads(result_path.read_text(encoding="utf-8"))["sequence_ids"] | |
| batch = prepare_processor_inputs( | |
| processor, | |
| matches[0], | |
| dataset.root, | |
| image_token_id=args.image_token_id, | |
| ) | |
| training_ids = batch["input_ids"][0, batch["prefix_length"] :].tolist() | |
| training_labels = training_raw_labels( | |
| processor.tokenizer, | |
| training_ids, | |
| _raw_record(dataset.root, args.sample_id), | |
| ) | |
| return test_ids, align_test_raw_labels(training_ids, training_labels, test_ids) | |
| def build_parser() -> argparse.ArgumentParser: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--mode", choices=["training", "test"], required=True) | |
| parser.add_argument("--sample-id", default="zebra-3d-counting-000011") | |
| parser.add_argument("--columns", type=int, default=6) | |
| parser.add_argument("--model-path", type=Path, default=Path("models/UniAR-SFT")) | |
| parser.add_argument("--samples-root", type=Path, default=Path("data/Zebra-CoT-samples")) | |
| parser.add_argument( | |
| "--evaluation-root", | |
| type=Path, | |
| default=Path( | |
| "experiments/output/uniar-sft/zebra-cot-samples/unversioned/smoke/eval/20260818-231640_zebra_overfit16_full_vs_base/" | |
| "finetuned_interleaved" | |
| ), | |
| ) | |
| parser.add_argument("--image-token-id", type=int, default=151655) | |
| parser.add_argument("--output", type=Path) | |
| return parser | |
| def main() -> None: | |
| args = build_parser().parse_args() | |
| processor = configure_processor( | |
| AutoProcessor.from_pretrained(args.model_path, use_fast=True) | |
| ) | |
| token_ids, raw_labels = load_ids_and_raw(args, processor) | |
| rendered = render_alignment( | |
| processor.tokenizer, | |
| token_ids, | |
| columns=args.columns, | |
| raw_labels=raw_labels, | |
| ) | |
| if args.output: | |
| args.output.parent.mkdir(parents=True, exist_ok=True) | |
| args.output.write_text(rendered, encoding="utf-8") | |
| print(rendered, end="") | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 8.8 kB
- Xet hash:
- 1070deaf77f1c8e4f599148b778c3ff10cfe98335a1f6d0ed8f00f582028fa6f
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.