Buckets:

EmergentTwilight
/

MMCoT

Files
xet
EmergentTwilight/MMCoT / OURS /src /sft /data.py
EmergentTwilight
20 days ago
download
raw
30.3 kB
"""Canonical cached interleaved trajectories and transition metadata."""

from __future__ import annotations

import hashlib
import itertools
import json
import random
from collections import OrderedDict
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Literal

import torch
from shared.labels import canonical_autonomous_text_labels
from shared.supervised import unpack_visual_codes
from uniar import SPECIAL_TOKEN_IDS

from .config import InformationGainConfig

CACHE_SCHEMAS = {"zebra_uniar_cache_v1"}
Direction = Literal["text_to_visual", "visual_to_text"]


@dataclass(frozen=True, order=True)
class Span:
    """Half-open positions in either a serialized sequence or flattened codes."""

    start: int
    end: int

    def __post_init__(self) -> None:
        if self.start < 0 or self.end <= self.start:
            raise ValueError(f"invalid half-open span [{self.start}, {self.end})")

    @property
    def length(self) -> int:
        return self.end - self.start


@dataclass(frozen=True)
class _Block:
    modality: Literal["text", "visual"]
    source_spans: tuple[Span, ...]
    target_sequence_spans: tuple[Span, ...]
    target_visual_spans: tuple[Span, ...] = ()


@dataclass(frozen=True)
class TransitionEvent:
    """One independently blinded adjacent block transition."""

    sample_id: str
    sample_index: int
    index: int
    direction: Direction
    source_sequence_spans: tuple[Span, ...]
    target_sequence_spans: tuple[Span, ...]
    target_visual_spans: tuple[Span, ...] = ()

    @property
    def target_steps(self) -> int:
        spans = (
            self.target_visual_spans
            if self.direction == "text_to_visual"
            else self.target_sequence_spans
        )
        return sum(span.length for span in spans)


@dataclass
class CachedSample:
    sample_id: str
    config: str
    source: dict
    input_ids: torch.Tensor
    labels: torch.Tensor
    image_grid_thw: torch.Tensor
    generated_sequence_mask: torch.Tensor
    generated_flat_mask: torch.Tensor
    visual_codes: torch.Tensor
    transitions: tuple[TransitionEvent, ...]


@dataclass
class InterleavedBatch:
    sample_ids: list[str]
    input_ids: torch.Tensor
    attention_mask: torch.Tensor
    labels: torch.Tensor
    image_grid_thw: torch.Tensor
    generated_sequence_mask: torch.Tensor
    generated_flat_mask: torch.Tensor
    visual_codes: torch.Tensor
    visual_code_spans: list[Span]
    transitions: list[list[TransitionEvent]]

    def selected_events(
        self,
        config: InformationGainConfig,
        *,
        seed: int,
        update: int,
    ) -> list[list[TransitionEvent]]:
        """Select transitions deterministically, independently for each direction."""

        selected: list[list[TransitionEvent]] = []
        for sample_id, events in zip(self.sample_ids, self.transitions):
            sample_events: list[TransitionEvent] = []
            for direction, direction_config in (
                ("text_to_visual", config.text_to_visual),
                ("visual_to_text", config.visual_to_text),
            ):
                candidates = [event for event in events if event.direction == direction]
                limit = direction_config.max_events_per_sample
                if (
                    config.event_selection == "all"
                    or limit is None
                    or limit >= len(candidates)
                ):
                    sample_events.extend(candidates)
                    continue
                digest = hashlib.sha256(
                    f"{seed}\0{sample_id}\0{direction}".encode()
                ).digest()
                base = int.from_bytes(digest[:8], "big") % len(candidates)
                start = (base + update * limit) % len(candidates)
                sample_events.extend(
                    candidates[(start + offset) % len(candidates)]
                    for offset in range(limit)
                )
            selected.append(sorted(sample_events, key=lambda event: event.index))
        return selected


def _true_spans(mask: torch.Tensor) -> list[Span]:
    values = mask.to(torch.bool).flatten().tolist()
    spans: list[Span] = []
    start = None
    for index, value in enumerate([*values, False]):
        if value and start is None:
            start = index
        elif not value and start is not None:
            spans.append(Span(start, index))
            start = None
    return spans


def _generated_visual_spans(item: dict) -> list[Span]:
    """Split flattened generated codes using per-image counts, not contiguity."""

    positions = torch.nonzero(
        item["generated_flat_mask"].to(torch.bool), as_tuple=False
    ).flatten()
    counts = [int(value) for value in item["generated_image_token_counts"].tolist()]
    if sum(counts) != positions.numel():
        raise ValueError("generated image counts do not cover generated_flat_mask")
    spans: list[Span] = []
    cursor = 0
    for count in counts:
        if count <= 0:
            raise ValueError(
                "generated images must contain at least one visual position"
            )
        selected = positions[cursor : cursor + count]
        start = int(selected[0])
        end = int(selected[-1]) + 1
        if end - start != count:
            raise ValueError("one generated image is not contiguous in flattened codes")
        spans.append(Span(start, end))
        cursor += count
    return spans


def _paired_think_spans(ids: list[int], prefix_length: int) -> list[tuple[Span, Span]]:
    start_id = SPECIAL_TOKEN_IDS["<think>"]
    end_id = SPECIAL_TOKEN_IDS["</think>"]
    stack: list[int] = []
    pairs: list[tuple[Span, Span]] = []
    for index in range(prefix_length, len(ids)):
        if ids[index] == start_id:
            if stack:
                raise ValueError("nested <think> tags are not supported")
            stack.append(index)
        elif ids[index] == end_id:
            if not stack:
                raise ValueError("unmatched </think> token")
            start = stack.pop()
            if index == start + 1:
                raise ValueError("empty thought block")
            pairs.append((Span(start, index + 1), Span(start + 1, index)))
    if stack:
        raise ValueError("unclosed <think> token")
    return pairs


def _assistant_end(ids: list[int], prefix_length: int) -> int:
    im_end = SPECIAL_TOKEN_IDS["<|im_end|>"]
    for index in range(prefix_length, len(ids)):
        if ids[index] == im_end:
            return index
    raise ValueError("assistant turn has no <|im_end|>")


def build_transition_events(
    item: dict,
    ig_config: InformationGainConfig,
) -> tuple[TransitionEvent, ...]:
    """Map typed assistant events to maximal adjacent text/visual blocks.
    Source masks contain only semantic thought content or image placeholders.
    Text targets include autonomous wrappers, image scaffold through
    ``<|vision_start|>``, and ``<|im_end|>``.
    """

    ids = item["input_ids"].to(torch.long).tolist()
    prefix = int(item["prefix_length"])
    assistant_end = _assistant_end(ids, prefix)
    thought_pairs = _paired_think_spans(ids, prefix)
    generated_sequence_spans = _true_spans(item["generated_sequence_mask"])
    generated_visual_spans = _generated_visual_spans(item)
    if len(generated_sequence_spans) != len(generated_visual_spans):
        raise ValueError("generated sequence/code image counts differ")
    for sequence, visual in zip(generated_sequence_spans, generated_visual_spans):
        if sequence.length != visual.length:
            raise ValueError("generated sequence/code image lengths differ")
        if ids[sequence.start - 1] != SPECIAL_TOKEN_IDS["<|vision_start|>"]:
            raise ValueError("generated image is not preceded by vision_start")
        if ids[sequence.end] != SPECIAL_TOKEN_IDS["<|vision_end|>"]:
            raise ValueError("generated image is not followed by vision_end")

    assistants = item["trajectory"]["assistant"]
    expected_thoughts = sum(event["role"] == "thought" for event in assistants)
    expected_images = sum(event["role"] == "generated_state" for event in assistants)
    if expected_thoughts != len(thought_pairs) or expected_images != len(
        generated_sequence_spans
    ):
        raise ValueError(
            "typed/tokenized event count mismatch: "
            f"thought={expected_thoughts}/{len(thought_pairs)}, "
            f"visual={expected_images}/{len(generated_sequence_spans)}"
        )

    atomics: list[_Block] = []
    thought_index = 0
    image_index = 0
    previous_structural_end = prefix
    for typed_index, typed in enumerate(assistants):
        role = typed["role"]
        if role == "thought":
            wrapped, content = thought_pairs[thought_index]
            thought_index += 1
            if wrapped.start < previous_structural_end:
                raise ValueError("typed thought order disagrees with token sequence")
            target_end = wrapped.end
            if (
                typed_index + 1 < len(assistants)
                and assistants[typed_index + 1]["role"] == "generated_state"
            ):
                if image_index >= len(generated_sequence_spans):
                    raise ValueError("thought image scaffold has no generated image")
                # The model owns the complete image scaffold through the
                # immediately preceding vision_start.  The first visual token
                # starts the following visual target.
                target_end = generated_sequence_spans[image_index].start
                if target_end < wrapped.end:
                    raise ValueError("image scaffold overlaps the thought payload")
            target = Span(previous_structural_end, target_end)
            atomics.append(
                _Block(
                    modality="text",
                    source_spans=(content,),
                    target_sequence_spans=(target,),
                )
            )
            previous_structural_end = target.end
        elif role == "generated_state":
            sequence = generated_sequence_spans[image_index]
            visual = generated_visual_spans[image_index]
            image_index += 1
            if sequence.start < previous_structural_end:
                raise ValueError("typed image order disagrees with token sequence")
            if sequence.start > previous_structural_end:
                # Adjacent generated_state records still have an autonomous
                # text scaffold between vision_end and the next image payload.
                # It is a V->T target but has no semantic source to blind for
                # the following T->V transition.
                atomics.append(
                    _Block(
                        modality="text",
                        source_spans=(),
                        target_sequence_spans=(
                            Span(previous_structural_end, sequence.start),
                        ),
                    )
                )
            atomics.append(
                _Block(
                    modality="visual",
                    source_spans=(sequence,),
                    target_sequence_spans=(sequence,),
                    target_visual_spans=(visual,),
                )
            )
            previous_structural_end = sequence.end + 1
        elif role == "answer":
            target_end = assistant_end + 1
            if previous_structural_end >= assistant_end:
                raise ValueError("empty final-answer token block")
            answer_payload = Span(previous_structural_end, assistant_end)
            answer_target = Span(previous_structural_end, target_end)
            atomics.append(
                _Block(
                    modality="text",
                    source_spans=(answer_payload,),
                    target_sequence_spans=(answer_target,),
                )
            )
            previous_structural_end = target_end
        else:
            raise ValueError(f"unknown assistant event role: {role!r}")

    if thought_index != len(thought_pairs) or image_index != len(
        generated_sequence_spans
    ):
        raise AssertionError("not all tokenized assistant events were consumed")

    blocks: list[_Block] = []
    for block in atomics:
        if blocks and blocks[-1].modality == block.modality:
            previous = blocks.pop()
            blocks.append(
                _Block(
                    modality=block.modality,
                    source_spans=previous.source_spans + block.source_spans,
                    target_sequence_spans=(
                        previous.target_sequence_spans + block.target_sequence_spans
                    ),
                    target_visual_spans=(
                        previous.target_visual_spans + block.target_visual_spans
                    ),
                )
            )
        else:
            blocks.append(block)

    transitions: list[TransitionEvent] = []
    for source, target in itertools.pairwise(blocks):
        if not source.source_spans:
            continue
        if source.modality == "text" and target.modality == "visual":
            direction: Direction = "text_to_visual"
        elif source.modality == "visual" and target.modality == "text":
            direction = "visual_to_text"
        else:
            raise AssertionError("maximal blocks must alternate modalities")
        transitions.append(
            TransitionEvent(
                sample_id=item["sample_id"],
                sample_index=-1,
                index=len(transitions),
                direction=direction,
                source_sequence_spans=source.source_spans,
                target_sequence_spans=target.target_sequence_spans,
                target_visual_spans=target.target_visual_spans,
            )
        )
    return tuple(transitions)


def cached_sample(item: dict, ig_config: InformationGainConfig) -> CachedSample:
    if item.get("schema") not in CACHE_SCHEMAS:
        raise ValueError(f"unsupported cache item schema: {item.get('schema')!r}")
    input_ids = item["input_ids"].to(torch.long)
    labels = canonical_autonomous_text_labels(
        input_ids, item["labels"].to(torch.long)
    )
    return CachedSample(
        sample_id=item["sample_id"],
        config=item["config"],
        source=item["source"],
        input_ids=input_ids,
        labels=labels,
        image_grid_thw=item["image_grid_thw"].to(torch.long),
        generated_sequence_mask=item["generated_sequence_mask"].to(torch.bool),
        generated_flat_mask=item["generated_flat_mask"].to(torch.bool),
        visual_codes=unpack_visual_codes(
            item["visual_codes_packed"], item["visual_codes_shape"]
        ),
        transitions=build_transition_events(item, ig_config),
    )


def collate_samples(samples: list[CachedSample], pad_token_id: int) -> InterleavedBatch:
    if not samples:
        raise ValueError("cannot collate an empty sample list")
    batch_size = len(samples)
    max_length = max(sample.input_ids.numel() for sample in samples)
    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)
    grids: list[torch.Tensor] = []
    codes: list[torch.Tensor] = []
    generated_flat: list[torch.Tensor] = []
    visual_code_spans: list[Span] = []
    transitions: list[list[TransitionEvent]] = []
    code_offset = 0
    for sample_index, sample in enumerate(samples):
        length = sample.input_ids.numel()
        input_ids[sample_index, :length] = sample.input_ids
        attention_mask[sample_index, :length] = 1
        labels[sample_index, :length] = sample.labels
        generated_sequence_mask[sample_index, :length] = sample.generated_sequence_mask
        grids.append(sample.image_grid_thw)
        codes.append(sample.visual_codes)
        generated_flat.append(sample.generated_flat_mask)
        visual_code_spans.append(
            Span(code_offset, code_offset + sample.visual_codes.shape[0])
        )
        adjusted = []
        for event in sample.transitions:
            adjusted.append(
                replace(
                    event,
                    sample_index=sample_index,
                    target_visual_spans=tuple(
                        Span(span.start + code_offset, span.end + code_offset)
                        for span in event.target_visual_spans
                    ),
                )
            )
        transitions.append(adjusted)
        code_offset += sample.visual_codes.shape[0]
    return InterleavedBatch(
        sample_ids=[sample.sample_id for sample in samples],
        input_ids=input_ids,
        attention_mask=attention_mask,
        labels=labels,
        image_grid_thw=torch.cat(grids),
        generated_sequence_mask=generated_sequence_mask,
        generated_flat_mask=torch.cat(generated_flat),
        visual_codes=torch.cat(codes),
        visual_code_spans=visual_code_spans,
        transitions=transitions,
    )


class InterleavedCache:
    """Manifest-indexed loader with a bounded or mmap-unbounded shard cache."""

    def __init__(
        self, root: str | Path, *, max_open_shards: int | None = 2
    ):
        self.root = Path(root)
        if max_open_shards is not None and max_open_shards <= 0:
            raise ValueError("max_open_shards must be positive or null")
        if not (self.root / "READY").exists():
            raise ValueError(f"cache is incomplete: {self.root}")
        with (self.root / "manifest.jsonl").open(encoding="utf-8") as handle:
            self.rows = [json.loads(line) for line in handle if line.strip()]
        self.by_id = {row["sample_id"]: row for row in self.rows}
        if len(self.by_id) != len(self.rows):
            raise ValueError("cache manifest contains duplicate sample IDs")
        self.max_open_shards = max_open_shards
        self._shards: OrderedDict[str, list[dict]] = OrderedDict()
        self._subsets: dict[str, list[dict]] = {}

    def _items(self, relative: str) -> list[dict]:
        if relative in self._shards:
            self._shards.move_to_end(relative)
            return self._shards[relative]
        payload = torch.load(
            self.root / relative,
            map_location="cpu",
            weights_only=False,
            mmap=True,
        )
        if payload.get("schema") not in CACHE_SCHEMAS:
            raise ValueError(f"unsupported shard schema in {relative}")
        self._shards[relative] = payload["items"]
        while (
            self.max_open_shards is not None
            and len(self._shards) > self.max_open_shards
        ):
            self._shards.popitem(last=False)
        return self._shards[relative]

    def load(self, row: dict, ig_config: InformationGainConfig) -> CachedSample:
        item = self._items(row["cache_shard"])[row["cache_index"]]
        if item["sample_id"] != row["sample_id"]:
            raise ValueError("manifest/cache sample mismatch")
        return cached_sample(item, ig_config)

    def load_ids(
        self, sample_ids: list[str], ig_config: InformationGainConfig
    ) -> list[CachedSample]:
        missing = sorted(set(sample_ids) - set(self.by_id))
        if missing:
            raise KeyError(f"sample IDs absent from cache: {missing}")
        return [self.load(self.by_id[sample_id], ig_config) for sample_id in sample_ids]

    def subset_rows(self, manifest_path: str | Path | None) -> list[dict]:
        if manifest_path is None:
            return self.rows
        key = str(Path(manifest_path).resolve())
        if key in self._subsets:
            return self._subsets[key]
        coordinates = small_subset_coordinates(manifest_path)
        selected = []
        matched = set()
        for row in self.rows:
            source = row["source"]
            coordinate = (
                row["config"],
                Path(source["shard"]).name,
                int(source["row"]),
            )
            if coordinate in coordinates:
                selected.append(row)
                matched.add(coordinate)
        missing = sorted(set(coordinates) - matched)
        if missing:
            raise ValueError(
                f"subset manifest has {len(missing)} rows absent from cache; "
                f"first missing coordinate: {missing[0]}"
            )
        if len(selected) != len(coordinates):
            raise ValueError("cache contains duplicate rows for subset coordinates")
        self._subsets[key] = selected
        return selected

    def split_rows(
        self, splits: str | list[str], subset_manifest: str | Path | None = None
    ) -> list[dict]:
        names = [splits] if isinstance(splits, str) else splits
        rows = [
            row
            for row in self.subset_rows(subset_manifest)
            if row["split"] in names
        ]
        if not rows:
            raise ValueError(f"cache/subset contains no rows for splits {names!r}")
        return rows

    def validation_rows(
        self,
        splits: list[str],
        count: int,
        subset_manifest: str | Path | None = None,
    ) -> list[dict]:
        rows = self.split_rows(splits, subset_manifest)
        selected: list[dict] = []
        seen: set[str] = set()
        by_config: dict[str, list[dict]] = {}
        for row in rows:
            by_config.setdefault(row["config"], []).append(row)
        for name in sorted(by_config):
            row = min(by_config[name], key=lambda value: value["split_score"])
            selected.append(row)
            seen.add(row["sample_id"])
        for row in sorted(rows, key=lambda value: value["split_score"]):
            if len(selected) >= count:
                break
            if row["sample_id"] not in seen:
                selected.append(row)
                seen.add(row["sample_id"])
        return selected[:count]


def plan_epoch_batches(
    rows: list[dict],
    *,
    seed: int,
    epoch: int,
    micro_batch_size: int,
    batch_padded_token_budget: int,
    max_sequence_tokens: int,
    world_size: int,
    distributed_tail: Literal["drop", "repeat"],
    shuffle: bool,
) -> list[list[list[dict]]]:
    """Build deterministic shard-local batches grouped one per DDP rank."""

    by_shard: dict[str, list[dict]] = {}
    for row in rows:
        if row["sequence_length"] > max_sequence_tokens:
            raise ValueError(
                f"sample {row['sample_id']} has {row['sequence_length']} tokens, "
                f"above data.max_sequence_tokens={max_sequence_tokens}"
            )
        by_shard.setdefault(row["cache_shard"], []).append(row)
    batches: list[list[dict]] = []
    randomizer = random.Random(seed + epoch * 1_000_003)
    shard_order = sorted(by_shard)
    if shuffle:
        randomizer.shuffle(shard_order)
    for relative in shard_order:
        shard_rows = sorted(by_shard[relative], key=lambda row: row["sequence_length"])
        shard_batches: list[list[dict]] = []
        current: list[dict] = []
        current_max = 0
        for row in shard_rows:
            candidate_max = max(current_max, row["sequence_length"])
            if current and (
                len(current) >= micro_batch_size
                or candidate_max * (len(current) + 1) > batch_padded_token_budget
            ):
                shard_batches.append(current)
                current = []
                current_max = 0
            current.append(row)
            current_max = max(current_max, row["sequence_length"])
        if current:
            shard_batches.append(current)
        if shuffle:
            randomizer.shuffle(shard_batches)
        # Preserve shard locality for the two-shard LRU. Global batch shuffling
        # would reload/unpickle a ~14 MiB shard for nearly every sample.
        batches.extend(shard_batches)
    remainder = len(batches) % world_size
    if remainder:
        if distributed_tail == "drop":
            batches = batches[:-remainder]
        else:
            needed = world_size - remainder
            original = list(batches)
            batches.extend(original[index % len(original)] for index in range(needed))
    if not batches:
        raise ValueError("no complete distributed batch groups")
    return [
        batches[index : index + world_size]
        for index in range(0, len(batches), world_size)
    ]


def plan_evaluation_batches(
    rows: list[dict],
    *,
    micro_batch_size: int,
    batch_padded_token_budget: int,
    rank: int,
    world_size: int,
) -> list[list[dict]]:
    """Assign whole shards to ranks, then length-pack rows within each shard."""

    if micro_batch_size <= 0 or batch_padded_token_budget <= 0:
        raise ValueError("evaluation batch limits must be positive")
    if world_size <= 0 or rank < 0 or rank >= world_size:
        raise ValueError("invalid evaluation rank/world size")
    by_shard: dict[str, list[dict]] = {}
    for row in rows:
        by_shard.setdefault(row["cache_shard"], []).append(row)

    assignments: list[list[str]] = [[] for _ in range(world_size)]
    loads = [0] * world_size
    shard_costs = {
        relative: sum(int(row["sequence_length"]) for row in shard_rows)
        for relative, shard_rows in by_shard.items()
    }
    for relative in sorted(by_shard, key=lambda key: (-shard_costs[key], key)):
        target = min(range(world_size), key=lambda index: (loads[index], index))
        assignments[target].append(relative)
        loads[target] += shard_costs[relative]

    batches = []
    for relative in sorted(assignments[rank]):
        current = []
        current_max = 0
        for row in sorted(
            by_shard[relative],
            key=lambda value: (value["sequence_length"], value["sample_id"]),
        ):
            candidate_max = max(current_max, int(row["sequence_length"]))
            if current and (
                len(current) >= micro_batch_size
                or candidate_max * (len(current) + 1) > batch_padded_token_budget
            ):
                batches.append(current)
                current = []
                current_max = 0
            current.append(row)
            current_max = max(current_max, int(row["sequence_length"]))
        if current:
            batches.append(current)
    return batches


def small_subset_coordinates(
    path: str | Path,
) -> dict[tuple[str, str, int], tuple[int, str]]:
    """Map original coordinates to materialized local row and Parquet path."""

    payload = json.loads(Path(path).read_text(encoding="utf-8"))
    if payload.get("format") != "zebra_cot_raw_parquet_subset_v1":
        raise ValueError("unexpected Zebra-CoT-small manifest format")
    coordinates: dict[tuple[str, str, int], tuple[int, str]] = {}
    for config in payload["configs"]:
        for local_row, source in enumerate(config["selected_rows"]):
            key = (config["name"], source["file"], int(source["row"]))
            if key in coordinates:
                raise ValueError(f"duplicate small-subset coordinate: {key}")
            coordinates[key] = (local_row, config["output_file"])
    return coordinates


def verify_small_subset_rows(rows: list[dict], manifest_path: str | Path) -> list[dict]:
    """Prove cache provenance and materialized small/source row equality."""

    manifest_path = Path(manifest_path)
    payload = json.loads(manifest_path.read_text(encoding="utf-8"))
    source_root = Path(payload["source_root"])
    coordinates = small_subset_coordinates(manifest_path)

    def parquet_row(path: Path, row_index: int):
        import pyarrow.parquet as pq

        parquet = pq.ParquetFile(path)
        offset = 0
        for group in range(parquet.num_row_groups):
            count = parquet.metadata.row_group(group).num_rows
            if row_index < offset + count:
                return parquet.read_row_group(group).slice(row_index - offset, 1)
            offset += count
        raise IndexError(f"row {row_index} is outside {path} ({offset} rows)")

    verified = []
    for row in rows:
        source = row["source"]
        key = (row["config"], Path(source["shard"]).name, int(source["row"]))
        if key not in coordinates:
            raise ValueError(f"cache row is not in Zebra-CoT-small: {row['sample_id']}")
        expected_id = hashlib.sha256(
            f"{row['config']}\0{source['shard']}\0{source['row']}".encode()
        ).hexdigest()[:24]
        if expected_id != row["sample_id"]:
            raise ValueError(f"cache provenance ID mismatch: {row['sample_id']}")
        local_row, relative = coordinates[key]
        materialized = Path(manifest_path).parent / relative
        if not materialized.is_file() or materialized.is_symlink():
            raise ValueError(
                f"small-subset Parquet is missing or linked: {materialized}"
            )
        source_parquet = source_root / source["shard"]
        local_value = parquet_row(materialized, local_row)
        source_value = parquet_row(source_parquet, int(source["row"]))
        if not local_value.equals(source_value, check_metadata=True):
            raise ValueError(
                f"materialized small row differs from source: {row['sample_id']}"
            )
        verified.append(
            {
                "sample_id": row["sample_id"],
                "config": row["config"],
                "small_local_row": local_row,
                "small_parquet": str(materialized),
                "source_parquet": str(source_parquet),
                "materialized_row_exact": True,
                "source": source,
            }
        )
    return verified


__all__ = [
    "CachedSample",
    "InterleavedBatch",
    "InterleavedCache",
    "Span",
    "TransitionEvent",
    "build_transition_events",
    "cached_sample",
    "collate_samples",
    "plan_epoch_batches",
    "plan_evaluation_batches",
    "small_subset_coordinates",
    "unpack_visual_codes",
    "verify_small_subset_rows",
]
Xet Storage Details

Size:: 30.3 kB
Xet hash:: 8bce689429fe23402c96b0bd4782b83952fc22700eabd2eef29b7233933f0025
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.