| """Learned hard adjacent-pair tokenization for a static embedding table.""" |
|
|
| from __future__ import annotations |
|
|
| from collections import deque |
| import math |
| from typing import TYPE_CHECKING |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
| from .collapse import _MIX_A, _MIX_B |
| from .structured_tokenizer import batched_matching, structured_straight_through |
|
|
| if TYPE_CHECKING: |
| from .collapse import CollapseTable |
|
|
|
|
| class OnlineMergeTable(nn.Module): |
| """Score every adjacent pair and emit a legal hard MAP token sequence.""" |
|
|
| def __init__( |
| self, |
| buckets: int, |
| dim: int, |
| n_languages: int, |
| *, |
| rank: int = 32, |
| bias_init: float = -2.0, |
| pair_key_base: int, |
| compiled_pair_keys: torch.Tensor | None = None, |
| compiled_pair_scores: torch.Tensor | None = None, |
| ) -> None: |
| super().__init__() |
| if buckets <= 0 or dim <= 0 or n_languages <= 0 or rank <= 0: |
| raise ValueError("online tokenizer dimensions must be positive") |
| if pair_key_base <= 0: |
| raise ValueError("pair_key_base must be positive") |
| if (compiled_pair_keys is None) != (compiled_pair_scores is None): |
| raise ValueError("compiled pair keys and scores must be provided together") |
|
|
| self.buckets = int(buckets) |
| self.dim = int(dim) |
| self.rank = int(rank) |
| self.pair_key_base = int(pair_key_base) |
| self.compiled = compiled_pair_keys is not None |
| keys = torch.as_tensor( |
| [] if compiled_pair_keys is None else compiled_pair_keys, dtype=torch.long |
| ) |
| scores = torch.as_tensor( |
| [] if compiled_pair_scores is None else compiled_pair_scores, |
| dtype=torch.float32, |
| ) |
| if keys.ndim != 1 or scores.ndim != 1 or keys.numel() != scores.numel(): |
| raise ValueError("compiled pair keys and scores must have the same length") |
| if keys.numel() and not bool((keys[1:] > keys[:-1]).all()): |
| raise ValueError("compiled pair keys must be strictly increasing") |
| if keys.numel() and (int(keys.min()) < 0 or int(keys.max()) >= pair_key_base**2): |
| raise ValueError("compiled pair key is outside pair_key_base") |
| if not bool(torch.isfinite(scores).all()): |
| raise ValueError("compiled pair scores must be finite") |
| self.register_buffer("compiled_pair_keys", keys, persistent=True) |
| self.register_buffer("compiled_pair_scores", scores, persistent=True) |
|
|
| rows = keys.numel() if self.compiled else self.buckets |
| self.merged = nn.Parameter( |
| torch.zeros(rows, self.dim), requires_grad=not self.compiled |
| ) |
| self.language_bias = nn.Parameter( |
| torch.full((n_languages,), float(bias_init)), requires_grad=not self.compiled |
| ) |
| if not self.compiled: |
| self.score_hash = nn.Parameter(torch.zeros(self.buckets)) |
| self.left_projection = nn.Parameter(torch.empty(self.rank, self.dim)) |
| self.right_projection = nn.Parameter(torch.empty(self.rank, self.dim)) |
| self.interaction_weight = nn.Parameter(torch.zeros(self.rank)) |
| nn.init.xavier_uniform_(self.left_projection) |
| nn.init.xavier_uniform_(self.right_projection) |
|
|
| self.last_token_count = torch.empty(0) |
| self.last_map_edges = torch.empty(0, 0, dtype=torch.bool) |
| self._vocabulary_telemetry = None |
| self._cache_audit_active = False |
| self._cache_phase: str | None = None |
| self._cache_records: deque[ |
| tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] |
| ] = deque() |
|
|
| def extra_repr(self) -> str: |
| return ( |
| f"compiled={self.compiled}, buckets={self.buckets}, " |
| f"exact_pairs={self.compiled_pair_keys.numel()}, dim={self.dim}, rank={self.rank}, " |
| f"languages={self.language_bias.numel()}" |
| ) |
|
|
| def set_vocabulary_telemetry(self, window) -> None: |
| """Attach detached run telemetry without registering module state.""" |
|
|
| self._vocabulary_telemetry = window |
|
|
| def _hash(self, left_ids: torch.Tensor, right_ids: torch.Tensor) -> torch.Tensor: |
| return ((left_ids * _MIX_A + right_ids * _MIX_B).abs()) % self.buckets |
|
|
| def start_cache_audit(self) -> None: |
| if self._cache_phase is not None or self._cache_records: |
| raise RuntimeError("previous online tokenizer cache audit is incomplete") |
| self._cache_audit_active = True |
|
|
| def begin_cache_record(self) -> None: |
| if self._cache_audit_active: |
| if self._cache_phase is not None: |
| raise RuntimeError("online tokenizer cache phase is already active") |
| self._cache_phase = "record" |
|
|
| def begin_cache_replay(self) -> None: |
| if self._cache_audit_active: |
| if self._cache_phase is not None: |
| raise RuntimeError("online tokenizer cache phase is already active") |
| self._cache_phase = "replay" |
|
|
| def end_cache_phase(self) -> None: |
| self._cache_phase = None |
|
|
| def abort_cache_audit(self) -> None: |
| self._cache_records.clear() |
| self._cache_phase = None |
| self._cache_audit_active = False |
|
|
| def assert_cache_audit_empty(self) -> None: |
| if self._cache_phase is not None or self._cache_records: |
| raise RuntimeError("online tokenizer cache audit did not replay every record") |
| self._cache_audit_active = False |
|
|
| def _audit_map( |
| self, |
| content: torch.Tensor, |
| segment: torch.Tensor, |
| language: torch.Tensor, |
| map_edges: torch.Tensor, |
| ) -> None: |
| if not self._cache_audit_active or self._cache_phase is None: |
| return |
| values = tuple( |
| tensor.detach().clone() |
| for tensor in (content, segment, language, map_edges) |
| ) |
| if self._cache_phase == "record": |
| self._cache_records.append(values) |
| return |
| if self._cache_phase != "replay" or not self._cache_records: |
| self.abort_cache_audit() |
| raise RuntimeError("cache/replay MAP mismatch") |
| expected = self._cache_records.popleft() |
| if not all( |
| observed.shape == wanted.shape and torch.equal(observed, wanted) |
| for observed, wanted in zip(values, expected, strict=True) |
| ): |
| self.abort_cache_audit() |
| raise RuntimeError("cache/replay MAP mismatch") |
|
|
| def seed_from_collapse(self, source: "CollapseTable") -> None: |
| """Copy the registered Q15 tensors without sharing parameter storage.""" |
|
|
| if self.compiled: |
| raise ValueError("cannot seed a compiled online tokenizer from Q15 hashes") |
| expected = { |
| "merged": tuple(self.merged.shape), |
| "score": tuple(self.score_hash.shape), |
| "language_bias": tuple(self.language_bias.shape), |
| } |
| observed = { |
| "merged": tuple(source.merged.shape), |
| "score": tuple(source.score.shape), |
| "language_bias": tuple(source.language_bias.shape), |
| } |
| if observed != expected: |
| raise ValueError( |
| f"Q15 source tensors do not fit online table: {observed} != {expected}" |
| ) |
| with torch.no_grad(): |
| self.merged.copy_(source.merged) |
| self.score_hash.copy_(source.score) |
| self.language_bias.copy_(source.language_bias) |
|
|
| def _edge_scores_and_index( |
| self, |
| left_ids: torch.Tensor, |
| right_ids: torch.Tensor, |
| language_ids: torch.Tensor, |
| embedding_weight: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| if not ( |
| left_ids.ndim == right_ids.ndim == language_ids.ndim == 1 |
| and left_ids.shape == right_ids.shape == language_ids.shape |
| ): |
| raise ValueError("edge scorer needs aligned one-dimensional tensors") |
| if embedding_weight.ndim != 2 or embedding_weight.shape[1] != self.dim: |
| raise ValueError("embedding_weight has the wrong shape") |
| if left_ids.numel() and ( |
| int(torch.minimum(left_ids.min(), right_ids.min())) < 0 |
| or int(torch.maximum(left_ids.max(), right_ids.max())) |
| >= embedding_weight.shape[0] |
| ): |
| raise ValueError("pair token id is outside the embedding table") |
| if language_ids.numel() and ( |
| int(language_ids.min()) < 0 |
| or int(language_ids.max()) >= self.language_bias.numel() |
| ): |
| raise ValueError("pair language id is outside the bias table") |
|
|
| if self.compiled: |
| keys = left_ids.to(torch.long) * self.pair_key_base + right_ids.to(torch.long) |
| if self.compiled_pair_keys.numel(): |
| index = torch.searchsorted(self.compiled_pair_keys, keys) |
| safe_index = index.clamp(max=self.compiled_pair_keys.numel() - 1) |
| found = (index < self.compiled_pair_keys.numel()) & ( |
| self.compiled_pair_keys.index_select(0, safe_index) == keys |
| ) |
| global_score = self.compiled_pair_scores.index_select(0, safe_index) |
| else: |
| safe_index = torch.zeros_like(keys) |
| found = torch.zeros_like(keys, dtype=torch.bool) |
| global_score = torch.zeros_like(keys, dtype=torch.float32) |
| scores = torch.where( |
| found, |
| global_score |
| + self.language_bias.index_select(0, language_ids).float(), |
| torch.zeros_like(global_score), |
| ) |
| index = safe_index |
| else: |
| index = self._hash(left_ids, right_ids) |
| device_type = embedding_weight.device.type |
| with torch.autocast(device_type=device_type, enabled=False): |
| left = embedding_weight.index_select(0, left_ids).float() |
| right = embedding_weight.index_select(0, right_ids).float() |
| left_code = torch.tanh(F.linear(left, self.left_projection.float())) |
| right_code = torch.tanh(F.linear(right, self.right_projection.float())) |
| amortized = ( |
| (left_code * right_code * self.interaction_weight.float()).sum(dim=1) |
| / math.sqrt(self.rank) |
| ) |
| scores = ( |
| self.score_hash.index_select(0, index).float() |
| + self.language_bias.index_select(0, language_ids).float() |
| + amortized |
| ) |
| found = torch.ones_like(index, dtype=torch.bool) |
| if not bool(torch.isfinite(scores).all()): |
| raise FloatingPointError("nonfinite online tokenizer edge score") |
| return scores, index, found |
|
|
| def edge_scores( |
| self, |
| left_ids: torch.Tensor, |
| right_ids: torch.Tensor, |
| language_ids: torch.Tensor, |
| embedding_weight: torch.Tensor, |
| ) -> torch.Tensor: |
| """Return float32 occurrence scores for aligned ids and languages.""" |
|
|
| scores, _, _ = self._edge_scores_and_index( |
| left_ids, right_ids, language_ids, embedding_weight |
| ) |
| return scores |
|
|
| @staticmethod |
| def _sentence_layout( |
| segment: torch.Tensor, n_sentences: int |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| lengths = torch.zeros( |
| n_sentences, dtype=torch.long, device=segment.device |
| ) |
| lengths.index_add_(0, segment, torch.ones_like(segment)) |
| if bool((lengths < 1).any()): |
| raise ValueError("online tokenizer requires at least one token per sentence") |
| starts = lengths.cumsum(0) - lengths |
| return lengths, starts |
|
|
| def pool( |
| self, |
| vectors: torch.Tensor, |
| content: torch.Tensor, |
| segment: torch.Tensor, |
| language: torch.Tensor, |
| n_sentences: int, |
| *, |
| embedding_weight: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Pool the literal tokens emitted by hard MAP with marginal gradients.""" |
|
|
| if vectors.ndim != 2 or vectors.shape != (content.numel(), self.dim): |
| raise ValueError("vectors must align with content and the table dimension") |
| if content.ndim != 1 or segment.ndim != 1 or content.shape != segment.shape: |
| raise ValueError("content and segment must be aligned one-dimensional tensors") |
| if language.shape != (n_sentences,): |
| raise ValueError("language must contain one id per sentence") |
| if segment.numel() and ( |
| int(segment.min()) < 0 or int(segment.max()) >= n_sentences |
| ): |
| raise ValueError("segment index is outside the sentence batch") |
| lengths, starts = self._sentence_layout(segment, n_sentences) |
|
|
| numerator = torch.zeros( |
| n_sentences, self.dim, dtype=vectors.dtype, device=vectors.device |
| ) |
| numerator.index_add_(0, segment, vectors) |
| token_count = lengths.to(dtype=vectors.dtype) |
| maximum_edges = int((lengths - 1).max().item()) |
| edge_mask = torch.zeros( |
| n_sentences, maximum_edges, dtype=torch.bool, device=vectors.device |
| ) |
| padded_scores = torch.zeros( |
| n_sentences, maximum_edges, dtype=torch.float32, device=vectors.device |
| ) |
|
|
| if maximum_edges: |
| inside = segment[:-1] == segment[1:] |
| positions = torch.nonzero(inside, as_tuple=False).squeeze(1) |
| edge_segment = segment.index_select(0, positions) |
| within = positions - starts.index_select(0, edge_segment) |
| left_ids = content.index_select(0, positions) |
| right_ids = content.index_select(0, positions + 1) |
| edge_language = language.index_select(0, edge_segment) |
| scores, pair_index, eligible = self._edge_scores_and_index( |
| left_ids, right_ids, edge_language, embedding_weight |
| ) |
| eligible_segment = edge_segment[eligible] |
| eligible_within = within[eligible] |
| edge_mask[eligible_segment, eligible_within] = True |
| padded_scores[eligible_segment, eligible_within] = scores[eligible] |
| matching = batched_matching(padded_scores, edge_mask) |
| if self.training: |
| padded_edges = structured_straight_through( |
| matching.marginals, matching.map_edges, dtype=vectors.dtype |
| ) |
| else: |
| padded_edges = matching.map_edges.to(dtype=vectors.dtype) |
| edges = padded_edges[eligible_segment, eligible_within] |
| pair_vectors = self.merged.index_select(0, pair_index[eligible]).to(vectors.dtype) |
| eligible_positions = positions[eligible] |
| delta = ( |
| pair_vectors |
| - vectors.index_select(0, eligible_positions) |
| - vectors.index_select(0, eligible_positions + 1) |
| ) |
| numerator.index_add_(0, eligible_segment, delta * edges.unsqueeze(1)) |
| token_count.index_add_(0, eligible_segment, -edges) |
| map_edges = matching.map_edges |
| telemetry_marginals = matching.marginals |
| telemetry_hash_indices = pair_index[eligible] |
| else: |
| map_edges = edge_mask |
| telemetry_marginals = torch.zeros_like(padded_scores) |
| telemetry_hash_indices = torch.empty( |
| 0, dtype=torch.long, device=content.device |
| ) |
|
|
| if not bool(torch.isfinite(token_count).all()) or bool((token_count < 1).any()): |
| raise FloatingPointError("invalid online tokenizer token count") |
| mean = numerator / token_count.unsqueeze(1) |
| embeddings = F.normalize(mean, dim=1) |
| if not bool(torch.isfinite(embeddings).all()): |
| raise FloatingPointError("nonfinite online tokenizer sentence embedding") |
| self.last_token_count = token_count.detach() |
| self.last_map_edges = map_edges.detach() |
| if self._vocabulary_telemetry is not None: |
| self._vocabulary_telemetry.observe_online( |
| base_tokens=int(content.numel()), |
| edge_scores=padded_scores, |
| marginals=telemetry_marginals, |
| map_edges=map_edges, |
| hash_indices=telemetry_hash_indices, |
| edge_mask=edge_mask, |
| ) |
| self._audit_map(content, segment, language, map_edges) |
| return embeddings, map_edges |
|
|
