| """Shared mmap-backed packed training row boundary. |
| |
| This module is import-safe for both the learn loop and the corpus-training |
| readers: it owns no state and imports nothing from either, so neither side can |
| create an import cycle by returning these rows across the I/O boundary. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import struct |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any, cast |
|
|
| import torch |
|
|
|
|
| _TARGET_FREE_CALIBRATION_INPUT_IDS_DOMAIN = ( |
| b"nnf.resynthesis.target_free_calibration_input_ids.v1\x00" |
| ) |
| _TARGET_FREE_CALIBRATION_INPUT_MASK_DOMAIN = ( |
| b"nnf.resynthesis.target_free_calibration_input_mask.v1\x00" |
| ) |
| _TARGET_FREE_CALIBRATION_PACKET_DOMAIN = ( |
| b"nnf.resynthesis.target_free_calibration_prompt_packet.v1\x00" |
| ) |
|
|
|
|
| def _domain_separated_tensor_sha256_t_boundary( |
| value_t: torch.Tensor, |
| *, |
| domain: bytes, |
| ) -> torch.Tensor: |
| """Hash one exact tensor at an explicit authority boundary. |
| |
| Calibration, cold proof, and acceptance must compare one identity domain. |
| Including dtype and geometry prevents equal raw bytes under a different |
| tensor interpretation from acquiring the same authority. |
| """ |
|
|
| if ( |
| not isinstance(value_t, torch.Tensor) |
| or value_t.device.type != "cpu" |
| or not value_t.is_contiguous() |
| or not isinstance(domain, bytes) |
| or not domain.endswith(b"\x00") |
| ): |
| raise ValueError("target-free calibration tensor authority differs") |
| digest = hashlib.sha256() |
| digest.update(domain) |
| dtype = str(value_t.dtype).encode("ascii") |
| digest.update(struct.pack("<Q", len(dtype))) |
| digest.update(dtype) |
| digest.update(struct.pack("<Q", value_t.ndim)) |
| for extent in value_t.shape: |
| digest.update(struct.pack("<Q", extent)) |
| digest.update(value_t.reshape(-1).view(torch.uint8).numpy().tobytes()) |
| return torch.frombuffer( |
| bytearray(digest.digest()), |
| dtype=torch.uint8, |
| ).clone() |
|
|
|
|
| def target_free_calibration_input_ids_sha256_t_boundary( |
| input_ids_t: torch.Tensor, |
| ) -> torch.Tensor: |
| """Return the shared exact identity for calibration prompt token IDs.""" |
|
|
| return _domain_separated_tensor_sha256_t_boundary( |
| input_ids_t, |
| domain=_TARGET_FREE_CALIBRATION_INPUT_IDS_DOMAIN, |
| ) |
|
|
|
|
| def target_free_calibration_input_mask_sha256_t_boundary( |
| input_mask_t: torch.Tensor, |
| ) -> torch.Tensor: |
| """Return the shared exact identity for a calibration prompt mask.""" |
|
|
| return _domain_separated_tensor_sha256_t_boundary( |
| input_mask_t, |
| domain=_TARGET_FREE_CALIBRATION_INPUT_MASK_DOMAIN, |
| ) |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class AllKnowledgeTargetFreeCalibrationPromptPacket: |
| """Tensor-only authority for one immutable, independently selected prompt. |
| |
| The prompt comes from a packed remaining-corpus snapshot, never historical |
| branch receipts, stores, scopes, prompts, or tensors. Filesystem paths and |
| JSON records remain outside this packet at explicit external-I/O |
| boundaries; every content identity needed by the model path is a tensor. |
| """ |
|
|
| input_ids_t: torch.Tensor |
| input_mask_t: torch.Tensor |
| global_cursor_t: torch.Tensor |
| prompt_sha256_t: torch.Tensor |
| window_sha256_t: torch.Tensor |
| row_authority_sha256_t: torch.Tensor |
| snapshot_file_sha256_t: torch.Tensor |
| collection_authority_sha256_t: torch.Tensor |
| federation_authority_sha256_t: torch.Tensor |
| exclusion_collection_file_sha256_t: torch.Tensor |
| excluded_prior_work_ids_sha256_t: torch.Tensor |
| input_ids_sha256_t: torch.Tensor |
| input_mask_sha256_t: torch.Tensor |
| physical_source_ranges_sha256_t: torch.Tensor |
| authority_sha256_t: torch.Tensor |
| target_entered_forward_t: torch.Tensor |
|
|
| def __post_init__(self) -> None: |
| tensors = ( |
| self.input_ids_t, |
| self.input_mask_t, |
| self.global_cursor_t, |
| self.prompt_sha256_t, |
| self.window_sha256_t, |
| self.row_authority_sha256_t, |
| self.snapshot_file_sha256_t, |
| self.collection_authority_sha256_t, |
| self.federation_authority_sha256_t, |
| self.exclusion_collection_file_sha256_t, |
| self.excluded_prior_work_ids_sha256_t, |
| self.input_ids_sha256_t, |
| self.input_mask_sha256_t, |
| self.physical_source_ranges_sha256_t, |
| self.authority_sha256_t, |
| self.target_entered_forward_t, |
| ) |
| digest_tensors = ( |
| self.prompt_sha256_t, |
| self.window_sha256_t, |
| self.snapshot_file_sha256_t, |
| self.collection_authority_sha256_t, |
| self.federation_authority_sha256_t, |
| self.exclusion_collection_file_sha256_t, |
| self.excluded_prior_work_ids_sha256_t, |
| self.input_ids_sha256_t, |
| self.input_mask_sha256_t, |
| self.physical_source_ranges_sha256_t, |
| self.authority_sha256_t, |
| ) |
| if ( |
| any( |
| not isinstance(value_t, torch.Tensor) |
| or value_t.device.type != "cpu" |
| or not value_t.is_contiguous() |
| for value_t in tensors |
| ) |
| or self.input_ids_t.dtype != torch.long |
| or self.input_ids_t.ndim != 2 |
| or self.input_ids_t.shape[0] != 1 |
| or self.input_ids_t.numel() < 1 |
| or bool(self.input_ids_t.lt(0).any()) |
| or self.input_mask_t.dtype != torch.bool |
| or self.input_mask_t.shape != self.input_ids_t.shape |
| or not bool(self.input_mask_t.all()) |
| or self.global_cursor_t.dtype != torch.int64 |
| or self.global_cursor_t.shape != (1,) |
| or bool(self.global_cursor_t.lt(0).any()) |
| or any( |
| value_t.dtype != torch.uint8 |
| or value_t.shape != (32,) |
| for value_t in digest_tensors |
| ) |
| or self.row_authority_sha256_t.dtype != torch.uint8 |
| or self.row_authority_sha256_t.shape != (2, 32) |
| or self.target_entered_forward_t.dtype != torch.bool |
| or self.target_entered_forward_t.shape != () |
| or bool(self.target_entered_forward_t) |
| or not torch.equal( |
| self.input_ids_sha256_t, |
| target_free_calibration_input_ids_sha256_t_boundary( |
| self.input_ids_t |
| ), |
| ) |
| or not torch.equal( |
| self.input_mask_sha256_t, |
| target_free_calibration_input_mask_sha256_t_boundary( |
| self.input_mask_t |
| ), |
| ) |
| or not torch.equal( |
| self.authority_sha256_t, |
| target_free_calibration_prompt_authority_sha256_t_boundary( |
| self |
| ), |
| ) |
| ): |
| raise ValueError( |
| "all-knowledge target-free calibration prompt differs" |
| ) |
|
|
|
|
| def target_free_calibration_prompt_authority_sha256_t_boundary( |
| packet: AllKnowledgeTargetFreeCalibrationPromptPacket, |
| ) -> torch.Tensor: |
| """Recompute one prompt packet's content authority without host metadata.""" |
|
|
| if not isinstance( |
| packet, |
| AllKnowledgeTargetFreeCalibrationPromptPacket, |
| ): |
| raise TypeError("target-free calibration prompt packet is malformed") |
| return _target_free_calibration_prompt_fields_authority_sha256_t_boundary( |
| ( |
| packet.input_ids_t, |
| packet.input_mask_t, |
| packet.global_cursor_t, |
| packet.prompt_sha256_t, |
| packet.window_sha256_t, |
| packet.row_authority_sha256_t, |
| packet.snapshot_file_sha256_t, |
| packet.collection_authority_sha256_t, |
| packet.federation_authority_sha256_t, |
| packet.exclusion_collection_file_sha256_t, |
| packet.excluded_prior_work_ids_sha256_t, |
| packet.input_ids_sha256_t, |
| packet.input_mask_sha256_t, |
| packet.physical_source_ranges_sha256_t, |
| packet.target_entered_forward_t, |
| ) |
| ) |
|
|
|
|
| def validated_target_free_calibration_prompt_packet_boundary( |
| packet: AllKnowledgeTargetFreeCalibrationPromptPacket, |
| ) -> AllKnowledgeTargetFreeCalibrationPromptPacket: |
| """Return an isolated, fully revalidated copy of one prompt packet.""" |
|
|
| if not isinstance( |
| packet, |
| AllKnowledgeTargetFreeCalibrationPromptPacket, |
| ): |
| raise TypeError("target-free calibration prompt packet is malformed") |
| return AllKnowledgeTargetFreeCalibrationPromptPacket( |
| input_ids_t=packet.input_ids_t.detach().cpu().clone().contiguous(), |
| input_mask_t=packet.input_mask_t.detach().cpu().clone().contiguous(), |
| global_cursor_t=( |
| packet.global_cursor_t.detach().cpu().clone().contiguous() |
| ), |
| prompt_sha256_t=( |
| packet.prompt_sha256_t.detach().cpu().clone().contiguous() |
| ), |
| window_sha256_t=( |
| packet.window_sha256_t.detach().cpu().clone().contiguous() |
| ), |
| row_authority_sha256_t=( |
| packet.row_authority_sha256_t.detach() |
| .cpu() |
| .clone() |
| .contiguous() |
| ), |
| snapshot_file_sha256_t=( |
| packet.snapshot_file_sha256_t.detach() |
| .cpu() |
| .clone() |
| .contiguous() |
| ), |
| collection_authority_sha256_t=( |
| packet.collection_authority_sha256_t.detach() |
| .cpu() |
| .clone() |
| .contiguous() |
| ), |
| federation_authority_sha256_t=( |
| packet.federation_authority_sha256_t.detach() |
| .cpu() |
| .clone() |
| .contiguous() |
| ), |
| exclusion_collection_file_sha256_t=( |
| packet.exclusion_collection_file_sha256_t.detach() |
| .cpu() |
| .clone() |
| .contiguous() |
| ), |
| excluded_prior_work_ids_sha256_t=( |
| packet.excluded_prior_work_ids_sha256_t.detach() |
| .cpu() |
| .clone() |
| .contiguous() |
| ), |
| input_ids_sha256_t=( |
| packet.input_ids_sha256_t.detach().cpu().clone().contiguous() |
| ), |
| input_mask_sha256_t=( |
| packet.input_mask_sha256_t.detach().cpu().clone().contiguous() |
| ), |
| physical_source_ranges_sha256_t=( |
| packet.physical_source_ranges_sha256_t.detach() |
| .cpu() |
| .clone() |
| .contiguous() |
| ), |
| authority_sha256_t=( |
| packet.authority_sha256_t.detach().cpu().clone().contiguous() |
| ), |
| target_entered_forward_t=( |
| packet.target_entered_forward_t.detach() |
| .cpu() |
| .clone() |
| .contiguous() |
| ), |
| ) |
|
|
|
|
| def _target_free_calibration_prompt_fields_authority_sha256_t_boundary( |
| fields_t: tuple[torch.Tensor, ...], |
| ) -> torch.Tensor: |
| """Hash the exact ordered tensor fields that form prompt authority.""" |
|
|
| digest = hashlib.sha256() |
| digest.update(_TARGET_FREE_CALIBRATION_PACKET_DOMAIN) |
| for value_t in fields_t: |
| tensor_digest_t = _domain_separated_tensor_sha256_t_boundary( |
| value_t, |
| domain=_TARGET_FREE_CALIBRATION_PACKET_DOMAIN, |
| ) |
| digest.update(tensor_digest_t.numpy().tobytes()) |
| return torch.frombuffer( |
| bytearray(digest.digest()), |
| dtype=torch.uint8, |
| ).clone() |
|
|
|
|
| def build_all_knowledge_target_free_calibration_prompt_packet_boundary( |
| *, |
| input_ids_t: torch.Tensor, |
| input_mask_t: torch.Tensor, |
| global_cursor_t: torch.Tensor, |
| prompt_sha256_t: torch.Tensor, |
| window_sha256_t: torch.Tensor, |
| row_authority_sha256_t: torch.Tensor, |
| snapshot_file_sha256_t: torch.Tensor, |
| collection_authority_sha256_t: torch.Tensor, |
| federation_authority_sha256_t: torch.Tensor, |
| exclusion_collection_file_sha256_t: torch.Tensor, |
| excluded_prior_work_ids_sha256_t: torch.Tensor, |
| physical_source_ranges_sha256_t: torch.Tensor, |
| target_entered_forward_t: torch.Tensor, |
| ) -> AllKnowledgeTargetFreeCalibrationPromptPacket: |
| """Seal one already validated prompt selection into its tensor packet.""" |
|
|
| input_ids_t = input_ids_t.detach().cpu().clone().contiguous() |
| input_mask_t = input_mask_t.detach().cpu().clone().contiguous() |
| global_cursor_t = global_cursor_t.detach().cpu().clone().contiguous() |
| prompt_sha256_t = prompt_sha256_t.detach().cpu().clone().contiguous() |
| window_sha256_t = window_sha256_t.detach().cpu().clone().contiguous() |
| row_authority_sha256_t = ( |
| row_authority_sha256_t.detach().cpu().clone().contiguous() |
| ) |
| snapshot_file_sha256_t = ( |
| snapshot_file_sha256_t.detach().cpu().clone().contiguous() |
| ) |
| collection_authority_sha256_t = ( |
| collection_authority_sha256_t.detach().cpu().clone().contiguous() |
| ) |
| federation_authority_sha256_t = ( |
| federation_authority_sha256_t.detach().cpu().clone().contiguous() |
| ) |
| exclusion_collection_file_sha256_t = ( |
| exclusion_collection_file_sha256_t.detach() |
| .cpu() |
| .clone() |
| .contiguous() |
| ) |
| excluded_prior_work_ids_sha256_t = ( |
| excluded_prior_work_ids_sha256_t.detach() |
| .cpu() |
| .clone() |
| .contiguous() |
| ) |
| physical_source_ranges_sha256_t = ( |
| physical_source_ranges_sha256_t.detach() |
| .cpu() |
| .clone() |
| .contiguous() |
| ) |
| target_entered_forward_t = ( |
| target_entered_forward_t.detach().cpu().clone().contiguous() |
| ) |
| input_ids_sha256_t = ( |
| target_free_calibration_input_ids_sha256_t_boundary(input_ids_t) |
| ) |
| input_mask_sha256_t = ( |
| target_free_calibration_input_mask_sha256_t_boundary(input_mask_t) |
| ) |
| fields_t = ( |
| input_ids_t, |
| input_mask_t, |
| global_cursor_t, |
| prompt_sha256_t, |
| window_sha256_t, |
| row_authority_sha256_t, |
| snapshot_file_sha256_t, |
| collection_authority_sha256_t, |
| federation_authority_sha256_t, |
| exclusion_collection_file_sha256_t, |
| excluded_prior_work_ids_sha256_t, |
| input_ids_sha256_t, |
| input_mask_sha256_t, |
| physical_source_ranges_sha256_t, |
| target_entered_forward_t, |
| ) |
| return AllKnowledgeTargetFreeCalibrationPromptPacket( |
| input_ids_t=input_ids_t, |
| input_mask_t=input_mask_t, |
| global_cursor_t=global_cursor_t, |
| prompt_sha256_t=prompt_sha256_t, |
| window_sha256_t=window_sha256_t, |
| row_authority_sha256_t=row_authority_sha256_t, |
| snapshot_file_sha256_t=snapshot_file_sha256_t, |
| collection_authority_sha256_t=collection_authority_sha256_t, |
| federation_authority_sha256_t=federation_authority_sha256_t, |
| exclusion_collection_file_sha256_t=( |
| exclusion_collection_file_sha256_t |
| ), |
| excluded_prior_work_ids_sha256_t=( |
| excluded_prior_work_ids_sha256_t |
| ), |
| input_ids_sha256_t=input_ids_sha256_t, |
| input_mask_sha256_t=input_mask_sha256_t, |
| physical_source_ranges_sha256_t=( |
| physical_source_ranges_sha256_t |
| ), |
| authority_sha256_t=( |
| _target_free_calibration_prompt_fields_authority_sha256_t_boundary( |
| fields_t |
| ) |
| ), |
| target_entered_forward_t=target_entered_forward_t, |
| ) |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class PackedTokenPhysicalSourceRange: |
| """Cold provenance for bytes copied from one immutable packed object. |
| |
| ``object_compressed_*`` identifies the physical zstd frame range that had |
| to be read. ``source_uncompressed_*`` identifies the exact bytes selected |
| from that frame's logical token stream, and ``packet_*`` binds those bytes |
| to their position in the coalesced ``PackedTokenBatchPacket`` arena. This |
| is external-I/O evidence only; it never enters model forward or routing. |
| """ |
|
|
| component_sha256: str |
| payload_work_id: str |
| object_path: str |
| object_sha256: str |
| object_byte_count: int |
| object_compressed_byte_start: int |
| object_compressed_byte_end: int |
| source_uncompressed_byte_start: int |
| source_uncompressed_byte_end: int |
| packet_byte_start: int |
| packet_byte_end: int |
| global_cursor: int |
|
|
| def __post_init__(self) -> None: |
| for digest, label in ( |
| (self.component_sha256, "component"), |
| (self.payload_work_id, "payload WorkID"), |
| (self.object_sha256, "object"), |
| ): |
| try: |
| decoded = bytes.fromhex(digest) |
| except ValueError as error: |
| raise ValueError( |
| f"packed token physical {label} digest differs" |
| ) from error |
| if len(decoded) != 32: |
| raise ValueError( |
| f"packed token physical {label} digest differs" |
| ) |
| if not self.object_path or not Path(self.object_path).is_absolute(): |
| raise ValueError("packed token physical object path differs") |
| integer_values = ( |
| self.object_byte_count, |
| self.object_compressed_byte_start, |
| self.object_compressed_byte_end, |
| self.source_uncompressed_byte_start, |
| self.source_uncompressed_byte_end, |
| self.packet_byte_start, |
| self.packet_byte_end, |
| self.global_cursor, |
| ) |
| if any( |
| not isinstance(value, int) |
| or isinstance(value, bool) |
| or value < 0 |
| for value in integer_values |
| ): |
| raise ValueError("packed token physical byte geometry differs") |
| if ( |
| self.object_byte_count < 1 |
| or self.object_compressed_byte_end |
| <= self.object_compressed_byte_start |
| or self.object_compressed_byte_end > self.object_byte_count |
| or self.source_uncompressed_byte_end |
| <= self.source_uncompressed_byte_start |
| or self.packet_byte_end <= self.packet_byte_start |
| or ( |
| self.source_uncompressed_byte_end |
| - self.source_uncompressed_byte_start |
| ) |
| != self.packet_byte_end - self.packet_byte_start |
| or self.source_uncompressed_byte_start % 4 |
| or self.source_uncompressed_byte_end % 4 |
| or self.packet_byte_start % 4 |
| or self.packet_byte_end % 4 |
| ): |
| raise ValueError("packed token physical byte ranges differ") |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class PackedTokenShardGeometryPacket: |
| """Tensor-native single-window geometry cached once per sealed shard.""" |
|
|
| token_offsets_t: torch.Tensor |
| prompt_lengths_t: torch.Tensor |
| prompt_sha256_t: torch.Tensor |
| window_sha256_t: torch.Tensor |
|
|
| def __post_init__(self) -> None: |
| row_count = self.prompt_lengths_t.numel() |
| if ( |
| self.token_offsets_t.device.type != "cpu" |
| or self.token_offsets_t.dtype != torch.int64 |
| or self.token_offsets_t.ndim != 1 |
| or self.token_offsets_t.numel() != row_count + 1 |
| or not self.token_offsets_t.is_contiguous() |
| or not torch.equal( |
| self.token_offsets_t[:1], |
| torch.zeros((1,), dtype=torch.int64), |
| ) |
| or not bool( |
| torch.all( |
| self.token_offsets_t[1:] |
| > self.token_offsets_t[:-1] |
| ) |
| ) |
| or self.prompt_lengths_t.device.type != "cpu" |
| or self.prompt_lengths_t.dtype != torch.int32 |
| or self.prompt_lengths_t.ndim != 1 |
| or not self.prompt_lengths_t.is_contiguous() |
| or not bool(torch.all(self.prompt_lengths_t > 0)) |
| or not bool( |
| torch.all( |
| self.prompt_lengths_t.to(dtype=torch.int64) |
| < ( |
| self.token_offsets_t[1:] |
| - self.token_offsets_t[:-1] |
| ) |
| ) |
| ) |
| ): |
| raise ValueError( |
| "packed token shard geometry differs" |
| ) |
| for digest_t, label in ( |
| (self.prompt_sha256_t, "prompt"), |
| (self.window_sha256_t, "window"), |
| ): |
| if ( |
| digest_t.device.type != "cpu" |
| or digest_t.dtype != torch.uint8 |
| or digest_t.shape != (row_count, 32) |
| or not digest_t.is_contiguous() |
| ): |
| raise ValueError( |
| f"packed token shard {label} authority differs" |
| ) |
|
|
| def __len__(self) -> int: |
| return self.prompt_lengths_t.numel() |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class PackedTokenBatchPacket: |
| """Tensor-native packed rows at the immutable corpus I/O boundary. |
| |
| ``token_ids_t`` stores each row as its prompt immediately followed by its |
| answer. ``row_offsets_t`` and ``prompt_lengths_t`` recover those two spans |
| without per-row Python objects. Authority rows store the component SHA-256 |
| followed by the payload WorkID SHA-256; ``row_authority_index_t`` binds |
| each token row to one such pair. |
| """ |
|
|
| token_ids_t: torch.Tensor |
| row_offsets_t: torch.Tensor |
| prompt_lengths_t: torch.Tensor |
| prompt_sha256_t: torch.Tensor |
| window_sha256_t: torch.Tensor |
| authority_sha256_t: torch.Tensor |
| row_authority_index_t: torch.Tensor |
| global_cursor_start_t: torch.Tensor |
| task_intent_targets_t: torch.Tensor | None = None |
|
|
| def __post_init__(self) -> None: |
| tensors = ( |
| self.token_ids_t, |
| self.row_offsets_t, |
| self.prompt_lengths_t, |
| self.prompt_sha256_t, |
| self.window_sha256_t, |
| self.authority_sha256_t, |
| self.row_authority_index_t, |
| self.global_cursor_start_t, |
| ) |
| if any( |
| value.device.type != "cpu" or not value.is_contiguous() |
| for value in tensors |
| ): |
| raise ValueError( |
| "packed token batch packet tensors must be contiguous CPU tensors" |
| ) |
| if self.token_ids_t.dtype != torch.int32 or self.token_ids_t.ndim != 1: |
| raise ValueError("packed token batch token arena differs") |
| if ( |
| self.row_offsets_t.dtype != torch.int64 |
| or self.row_offsets_t.ndim != 1 |
| or self.row_offsets_t.numel() < 2 |
| ): |
| raise ValueError("packed token batch row offsets differ") |
| if not torch.equal( |
| self.row_offsets_t[:1], |
| torch.zeros((1,), dtype=torch.int64), |
| ) or not torch.equal( |
| self.row_offsets_t[-1:], |
| torch.tensor( |
| [self.token_ids_t.numel()], |
| dtype=torch.int64, |
| ), |
| ): |
| raise ValueError("packed token batch token frontier differs") |
| row_widths_t = self.row_offsets_t[1:] - self.row_offsets_t[:-1] |
| if not bool(torch.all(row_widths_t > 0)): |
| raise ValueError("packed token batch row widths differ") |
| row_count = self.row_offsets_t.numel() - 1 |
| if ( |
| self.prompt_lengths_t.dtype != torch.int32 |
| or self.prompt_lengths_t.shape != (row_count,) |
| or not bool(torch.all(self.prompt_lengths_t > 0)) |
| or not bool( |
| torch.all( |
| self.prompt_lengths_t.to(dtype=torch.int64) |
| < row_widths_t |
| ) |
| ) |
| ): |
| raise ValueError("packed token batch prompt geometry differs") |
| for digest_t, label in ( |
| (self.prompt_sha256_t, "prompt"), |
| (self.window_sha256_t, "window"), |
| ): |
| if digest_t.dtype != torch.uint8 or digest_t.shape != (row_count, 32): |
| raise ValueError(f"packed token batch {label} authority differs") |
| if ( |
| self.authority_sha256_t.dtype != torch.uint8 |
| or self.authority_sha256_t.ndim != 3 |
| or self.authority_sha256_t.shape[0] < 1 |
| or self.authority_sha256_t.shape[1:] != (2, 32) |
| ): |
| raise ValueError("packed token batch authority table differs") |
| if ( |
| self.row_authority_index_t.dtype != torch.int32 |
| or self.row_authority_index_t.shape != (row_count,) |
| or not bool(torch.all(self.row_authority_index_t >= 0)) |
| or not bool( |
| torch.all( |
| self.row_authority_index_t |
| < self.authority_sha256_t.shape[0] |
| ) |
| ) |
| ): |
| raise ValueError("packed token batch row authority indices differ") |
| if ( |
| self.global_cursor_start_t.dtype != torch.int64 |
| or self.global_cursor_start_t.shape != (1,) |
| or not bool(torch.all(self.global_cursor_start_t >= 0)) |
| ): |
| raise ValueError("packed token batch global cursor differs") |
| if self.task_intent_targets_t is not None and ( |
| self.task_intent_targets_t.device.type != "cpu" |
| or self.task_intent_targets_t.dtype != torch.float32 |
| or self.task_intent_targets_t.ndim != 2 |
| or self.task_intent_targets_t.shape[0] != row_count |
| or self.task_intent_targets_t.shape[1] < 1 |
| or not self.task_intent_targets_t.is_contiguous() |
| or not bool(torch.isfinite(self.task_intent_targets_t).all()) |
| ): |
| raise ValueError( |
| "packed token batch task-intent arena differs" |
| ) |
|
|
| def __len__(self) -> int: |
| return self.row_offsets_t.numel() - 1 |
|
|
| def prompt_ids_boundary(self, row_index: int) -> torch.Tensor: |
| """Return one prompt view at an explicit Python consumer boundary.""" |
|
|
| if ( |
| isinstance(row_index, bool) |
| or row_index < 0 |
| or row_index >= len(self) |
| ): |
| raise IndexError("packed token batch row is unavailable") |
| row_start = int(self.row_offsets_t[row_index]) |
| prompt_end = row_start + int(self.prompt_lengths_t[row_index]) |
| return self.token_ids_t[row_start:prompt_end] |
|
|
| def answer_ids_boundary(self, row_index: int) -> torch.Tensor: |
| """Return one answer view at an explicit Python consumer boundary.""" |
|
|
| if ( |
| isinstance(row_index, bool) |
| or row_index < 0 |
| or row_index >= len(self) |
| ): |
| raise IndexError("packed token batch row is unavailable") |
| row_start = int(self.row_offsets_t[row_index]) |
| answer_start = row_start + int(self.prompt_lengths_t[row_index]) |
| row_end = int(self.row_offsets_t[row_index + 1]) |
| return self.token_ids_t[answer_start:row_end] |
|
|
| def task_intent_targets_boundary( |
| self, |
| row_index: int, |
| ) -> torch.Tensor | None: |
| """Return one loss-side intent view at an explicit consumer boundary. |
| |
| Intent labels remain separate from ``token_ids_t`` and therefore |
| cannot enter the prompt forward. Keeping their rows in this packet |
| still gives the loss boundary the same cursor, authority, and reorder |
| semantics as the token arena. |
| """ |
|
|
| if ( |
| isinstance(row_index, bool) |
| or row_index < 0 |
| or row_index >= len(self) |
| ): |
| raise IndexError("packed token batch row is unavailable") |
| if self.task_intent_targets_t is None: |
| return None |
| return self.task_intent_targets_t[row_index] |
|
|
|
|
| class PackedTokenTrainingRow(dict[str, Any]): |
| """One metadata row backed by mmap token tensors at the I/O boundary.""" |
|
|
| batch_packet: PackedTokenBatchPacket | None |
| batch_packet_row_index: int | None |
|
|
| def __init__( |
| self, |
| metadata: dict[str, Any], |
| *, |
| prompt_ids: torch.Tensor, |
| answer_ids: torch.Tensor, |
| corrected_prompt_ids: torch.Tensor, |
| task_intent_targets: torch.Tensor | None, |
| source_row_sha256: str, |
| batch_packet: PackedTokenBatchPacket | None = None, |
| batch_packet_row_index: int | None = None, |
| ) -> None: |
| super().__init__(metadata) |
| if (batch_packet is None) != (batch_packet_row_index is None): |
| raise ValueError("packed token row batch binding is incomplete") |
| if batch_packet is not None and ( |
| not isinstance(batch_packet_row_index, int) |
| or isinstance(batch_packet_row_index, bool) |
| or batch_packet_row_index < 0 |
| or batch_packet_row_index >= len(batch_packet) |
| ): |
| raise ValueError("packed token row batch index differs") |
| self["prompt_ids"] = prompt_ids |
| self["answer_ids"] = answer_ids |
| if task_intent_targets is not None: |
| self["task_intent_targets"] = task_intent_targets |
| self.corrected_prompt_ids_t = corrected_prompt_ids |
| self.source_row_sha256 = source_row_sha256 |
| self.batch_packet = batch_packet |
| self.batch_packet_row_index = batch_packet_row_index |
|
|
| @staticmethod |
| def _integer_list(value: torch.Tensor) -> list[int]: |
| return [ |
| int(token) |
| for token in value.to(device="cpu", dtype=torch.long).tolist() |
| ] |
|
|
| def legacy_row_boundary(self) -> dict[str, Any]: |
| """Materialize Python sequences only for legacy single-row paths.""" |
|
|
| prompt_ids = self._integer_list(cast(torch.Tensor, self["prompt_ids"])) |
| answer_ids = self._integer_list(cast(torch.Tensor, self["answer_ids"])) |
| corrected_prompt_ids = self._integer_list(self.corrected_prompt_ids_t) |
| row = { |
| key: value |
| for key, value in self.items() |
| if key not in {"prompt_ids", "answer_ids", "task_intent_targets"} |
| } |
| row.update( |
| { |
| "prompt_ids": prompt_ids, |
| "answer_ids": answer_ids, |
| "input_ids": [*prompt_ids, *answer_ids], |
| "target_ids": [-100] * len(prompt_ids) + answer_ids, |
| "corrected_prompt_ids": corrected_prompt_ids, |
| "corrected_input_ids": [ |
| *corrected_prompt_ids, |
| *answer_ids, |
| ], |
| "corrected_target_ids": ( |
| [-100] * len(corrected_prompt_ids) + answer_ids |
| ), |
| } |
| ) |
| task_intent_targets = self.get("task_intent_targets") |
| if isinstance(task_intent_targets, torch.Tensor): |
| row["task_intent_targets"] = [ |
| float(value) |
| for value in task_intent_targets.to( |
| device="cpu", |
| dtype=torch.float32, |
| ).tolist() |
| ] |
| return row |
|
|
| def bmpas_tensor_row_boundary(self) -> dict[str, Any]: |
| """Tensor-native row for BMPAS train_fn — no Python list materialization.""" |
|
|
| prompt_t = cast(torch.Tensor, self["prompt_ids"]).to(dtype=torch.long) |
| answer_t = cast(torch.Tensor, self["answer_ids"]).to(dtype=torch.long) |
| corrected_t = self.corrected_prompt_ids_t.to(dtype=torch.long) |
| input_ids = torch.cat((prompt_t, answer_t), dim=0) |
| target_ids = torch.cat( |
| ( |
| torch.full((prompt_t.numel(),), -100, dtype=torch.long), |
| answer_t, |
| ), |
| dim=0, |
| ) |
| corrected_input = torch.cat((corrected_t, answer_t), dim=0) |
| corrected_target = torch.cat( |
| ( |
| torch.full((corrected_t.numel(),), -100, dtype=torch.long), |
| answer_t, |
| ), |
| dim=0, |
| ) |
| row: dict[str, Any] = { |
| key: value |
| for key, value in self.items() |
| if key not in {"prompt_ids", "answer_ids", "task_intent_targets"} |
| } |
| row.update( |
| { |
| "prompt_ids": prompt_t, |
| "answer_ids": answer_t, |
| "input_ids": input_ids, |
| "target_ids": target_ids, |
| "corrected_prompt_ids": corrected_t, |
| "corrected_input_ids": corrected_input, |
| "corrected_target_ids": corrected_target, |
| "_bmpas_tensor_native": True, |
| } |
| ) |
| task_intent_targets = self.get("task_intent_targets") |
| if isinstance(task_intent_targets, torch.Tensor): |
| row["task_intent_targets"] = task_intent_targets |
| return row |
|
|
| def __eq__(self, other: object) -> bool: |
| if isinstance(other, PackedTokenTrainingRow): |
| return self.source_row_sha256 == other.source_row_sha256 |
| if isinstance(other, dict): |
| return self.legacy_row_boundary() == other |
| return False |
|
|
