from __future__ import annotations import hashlib import math import os from collections.abc import Callable, Mapping from dataclasses import dataclass from numbers import Integral, Real from pathlib import Path import torch import torch.nn.functional as F PINNED_ARTIFACT_VERSION = 1 PINNED_ARTIFACT_ALGORITHM = 'recursive_centered_pca_rank_median_v1' PINNED_ARTIFACT_METRIC = 'normalized_squared_l2_tree_ultrametric' PINNED_PRODUCTION_VOCAB_SIZE = 131072 PINNED_PRODUCTION_DEPTH = 17 DEFAULT_SEMANTIC_TREE_STEPS = 4 DEFAULT_SEMANTIC_TREE_SCHEDULE = 'shift' DEFAULT_SEMANTIC_TREE_SHIFT = 5.0 MAX_SEMANTIC_TREE_REFINEMENT_GAP = 2 SUPPORTED_SEMANTIC_TREE_INFERENCE_SCHEDULES = frozenset({'linear', 'shift', 'logit_normal', 'cosine', 'cosine2', 'exp'}) def _plain_positive_int(value: object, *, name: str) -> int: if isinstance(value, bool) or not isinstance(value, Integral): raise TypeError(f'{name} must be a positive integer') value = int(value) if value < 1: raise ValueError(f'{name} must be a positive integer') return value def _plain_nonnegative_int(value: object, *, name: str) -> int: if isinstance(value, bool) or not isinstance(value, Integral): raise TypeError(f'{name} must be a non-negative integer') value = int(value) if value < 0: raise ValueError(f'{name} must be a non-negative integer') return value def _plain_positive_float(value: object, *, name: str) -> float: if isinstance(value, bool) or not isinstance(value, Real): raise TypeError(f'{name} must be a finite positive number') value = float(value) if not math.isfinite(value) or value <= 0.0: raise ValueError(f'{name} must be a finite positive number') return value def _is_power_of_two(value: int) -> bool: return value > 0 and value & value - 1 == 0 def resolve_semantic_tree_steps(config, override=None): steps = 4 if override is None else override if isinstance(steps, bool) or steps not in (1, 4, 8, 16): raise ValueError('refinement_steps must be 1, 4, 8, or 16') return int(steps) @dataclass(frozen=True, slots=True) class SemanticTreeArtifactContract: path: str | os.PathLike[str] bytes: int | None = None sha256: str | None = None algorithm: str | None = PINNED_ARTIFACT_ALGORITHM def __post_init__(self) -> None: if not isinstance(self.path, (str, os.PathLike)): raise TypeError('semantic-tree artifact path must be path-like') raw_path = os.fspath(self.path) if not isinstance(raw_path, str) or not raw_path.strip(): raise ValueError('semantic-tree artifact path must be nonempty') object.__setattr__(self, 'path', Path(raw_path).expanduser()) if self.bytes is not None: object.__setattr__(self, 'bytes', _plain_positive_int(self.bytes, name='semantic-tree artifact bytes')) if self.sha256 is not None: if not isinstance(self.sha256, str): raise TypeError('semantic-tree artifact sha256 must be a string') digest = self.sha256.strip().lower() if len(digest) != 64 or any((character not in '0123456789abcdef' for character in digest)): raise ValueError('semantic-tree artifact sha256 must be 64 hexadecimal characters') object.__setattr__(self, 'sha256', digest) if self.algorithm not in (None, PINNED_ARTIFACT_ALGORITHM): raise ValueError(f'Unsupported semantic-tree algorithm: {self.algorithm!r}') def production_semantic_tree_contract(path: str | os.PathLike[str]) -> SemanticTreeArtifactContract: return SemanticTreeArtifactContract(path=path, sha256='4b7ee94fbb10fd4146a4ad167b595b037e3f591eaa12ecf3f0b411c82384a714', algorithm=PINNED_ARTIFACT_ALGORITHM) def _sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open('rb') as artifact_file: for block in iter(lambda : artifact_file.read(1024 * 1024), b''): digest.update(block) return digest.hexdigest() class BalancedSemanticTree: def __init__(self, leaf_to_token: torch.Tensor) -> None: if not torch.is_tensor(leaf_to_token): raise TypeError('leaf_to_token must be a tensor') if leaf_to_token.ndim != 1: raise ValueError('leaf_to_token must be one-dimensional') if leaf_to_token.dtype == torch.bool or leaf_to_token.is_floating_point() or leaf_to_token.is_complex(): raise TypeError('leaf_to_token must contain integer token IDs') vocab_size = int(leaf_to_token.numel()) if not _is_power_of_two(vocab_size): raise ValueError('semantic-tree vocabulary size must be a power of two') permutation = leaf_to_token.to(dtype=torch.long).contiguous() expected = torch.arange(vocab_size, device=permutation.device) if not torch.equal(torch.sort(permutation).values, expected): raise ValueError('leaf_to_token must be a permutation of [0, V)') token_to_leaf = torch.empty_like(permutation) token_to_leaf.scatter_(0, permutation, expected) self.vocab_size = vocab_size self.depth = int(math.log2(vocab_size)) self.leaf_to_token = permutation self.token_to_leaf = token_to_leaf @classmethod def _from_validated(cls, leaf_to_token: torch.Tensor, token_to_leaf: torch.Tensor) -> BalancedSemanticTree: tree = cls.__new__(cls) tree.vocab_size = int(leaf_to_token.numel()) tree.depth = int(math.log2(tree.vocab_size)) tree.leaf_to_token = leaf_to_token tree.token_to_leaf = token_to_leaf return tree @classmethod def from_artifact(cls, contract: SemanticTreeArtifactContract, expected_vocab_size: int=PINNED_PRODUCTION_VOCAB_SIZE) -> BalancedSemanticTree: if not isinstance(contract, SemanticTreeArtifactContract): raise TypeError('contract must be a SemanticTreeArtifactContract') expected_vocab_size = _plain_positive_int(expected_vocab_size, name='expected_vocab_size') if not _is_power_of_two(expected_vocab_size): raise ValueError('expected_vocab_size must be a power of two') path = contract.path assert isinstance(path, Path) if not path.is_file(): raise FileNotFoundError(f'semantic-tree artifact is not a file: {path}') actual_bytes = path.stat().st_size if contract.bytes is not None and actual_bytes != contract.bytes: raise ValueError(f'semantic-tree artifact byte count mismatch: expected {contract.bytes}, got {actual_bytes}') if contract.sha256 is not None: actual_sha256 = _sha256_file(path) if actual_sha256 != contract.sha256: raise ValueError(f'semantic-tree artifact sha256 mismatch: expected {contract.sha256}, got {actual_sha256}') artifact = torch.load(path, map_location='cpu', weights_only=True) if not isinstance(artifact, Mapping): raise TypeError('semantic-tree artifact must contain a mapping') algorithm = artifact.get('algorithm') allowed_algorithms = (contract.algorithm,) if contract.algorithm is not None else (PINNED_ARTIFACT_ALGORITHM,) if algorithm not in allowed_algorithms: raise ValueError(f'semantic-tree artifact algorithm mismatch: expected one of {allowed_algorithms!r}, got {algorithm!r}') pinned_values = {'version': PINNED_ARTIFACT_VERSION, 'metric': PINNED_ARTIFACT_METRIC, 'vocab_size': expected_vocab_size, 'depth': int(math.log2(expected_vocab_size)), 'radii_monotone': True} for (field_name, expected_value) in pinned_values.items(): if artifact.get(field_name) != expected_value: raise ValueError(f'semantic-tree artifact {field_name} mismatch: expected {expected_value!r}, got {artifact.get(field_name)!r}') leaf_to_token = artifact.get('leaf_to_token') heap_radii = artifact.get('heap_radii') if not torch.is_tensor(leaf_to_token) or leaf_to_token.dtype != torch.int32: raise TypeError('semantic-tree leaf_to_token must be torch.int32') if tuple(leaf_to_token.shape) != (expected_vocab_size,): raise ValueError('semantic-tree leaf_to_token shape disagrees with vocabulary') if not torch.is_tensor(heap_radii) or heap_radii.dtype != torch.float32: raise TypeError('semantic-tree heap_radii must be torch.float32') if tuple(heap_radii.shape) != (expected_vocab_size - 1,): raise ValueError('semantic-tree heap_radii shape disagrees with vocabulary') if not bool(torch.isfinite(heap_radii).all()) or bool((heap_radii < 0).any()): raise ValueError('semantic-tree heap_radii must be finite and nonnegative') return cls(leaf_to_token) @property def device(self) -> torch.device: return self.leaf_to_token.device def to(self, device: torch.device | str) -> BalancedSemanticTree: resolved = torch.device(device) if resolved == self.device: return self return type(self)._from_validated(self.leaf_to_token.to(resolved), self.token_to_leaf.to(resolved)) def _check_tokens(self, tokens: torch.Tensor, *, name: str) -> None: if not torch.is_tensor(tokens): raise TypeError(f'{name} must be a tensor') if tokens.dtype == torch.bool or tokens.is_floating_point() or tokens.is_complex(): raise TypeError(f'{name} must contain integer code IDs') if tokens.device != self.device: raise ValueError(f'{name} and semantic-tree lookups must share a device') if tokens.numel() and bool(((tokens < 0) | (tokens >= self.vocab_size)).any()): raise ValueError(f'{name} contains code IDs outside [0, V)') def _levels(self, levels: int | torch.Tensor, shape: torch.Size, device: torch.device) -> torch.Tensor: if isinstance(levels, bool) or (not isinstance(levels, Integral) and (not torch.is_tensor(levels))): raise TypeError('levels must contain integers') level_tensor = torch.as_tensor(levels, device=device) if level_tensor.dtype == torch.bool or level_tensor.is_floating_point() or level_tensor.is_complex(): raise TypeError('levels must contain integers') if level_tensor.ndim < len(shape): level_tensor = level_tensor.reshape(level_tensor.shape + (1,) * (len(shape) - level_tensor.ndim)) try: level_tensor = torch.broadcast_to(level_tensor, shape).long() except RuntimeError as error: raise ValueError('levels are not broadcastable to endpoint shape') from error if bool((level_tensor < 0).any()) or bool((level_tensor > self.depth).any()): raise ValueError(f'levels must lie in [0, {self.depth}]') return level_tensor def node_ids(self, tokens: torch.Tensor, levels: int | torch.Tensor) -> torch.Tensor: self._check_tokens(tokens, name='tokens') level_tensor = self._levels(levels, tokens.shape, tokens.device) leaves = self.token_to_leaf[tokens.long()] return torch.bitwise_right_shift(leaves, self.depth - level_tensor) @torch.no_grad() def sample_representatives(self, endpoints: torch.Tensor, levels: int | torch.Tensor, generator: torch.Generator | None=None) -> torch.Tensor: self._check_tokens(endpoints, name='endpoints') level_tensor = self._levels(levels, endpoints.shape, endpoints.device) leaves = self.token_to_leaf[endpoints.long()] noisy = level_tensor < self.depth if not bool(noisy.any()): return endpoints.clone() noisy_leaves = leaves[noisy] noisy_levels = level_tensor[noisy] shifts = self.depth - noisy_levels subtree_starts = torch.bitwise_left_shift(torch.bitwise_right_shift(noisy_leaves, shifts), shifts) subtree_sizes = torch.bitwise_left_shift(torch.ones_like(shifts), shifts) random_bits = torch.randint(self.vocab_size, (int(noisy.sum().item()),), device=endpoints.device, dtype=torch.long, generator=generator) sampled_leaves = subtree_starts + torch.bitwise_and(random_bits, subtree_sizes - 1) result = endpoints.clone() result[noisy] = self.leaf_to_token[sampled_leaves].to(endpoints.dtype) return result def _validate_ages(ages: torch.Tensor, *, steps: int) -> tuple[torch.Tensor, int]: steps = _plain_positive_int(steps, name='steps') if not torch.is_tensor(ages): raise TypeError('ages must be a tensor') if ages.dtype == torch.bool or ages.is_floating_point() or ages.is_complex(): raise TypeError('ages must contain integers') if bool((ages < 0).any()) or bool((ages > steps).any()): raise ValueError(f'ages must lie in [0, {steps}]') return (ages.long(), steps) def semantic_tree_levels_for_ages(ages: torch.Tensor, steps: int=DEFAULT_SEMANTIC_TREE_STEPS, depth: int=PINNED_PRODUCTION_DEPTH) -> torch.Tensor: (ages, steps) = _validate_ages(ages, steps=steps) depth = _plain_positive_int(depth, name='depth') if steps > depth: raise ValueError('steps cannot exceed tree depth') return torch.div(ages * depth, steps, rounding_mode='floor') def semantic_tree_noise_for_ages(ages: torch.Tensor, steps: int=DEFAULT_SEMANTIC_TREE_STEPS) -> torch.Tensor: (ages, steps) = _validate_ages(ages, steps=steps) return 1.0 - ages.to(dtype=torch.float32) / float(steps) def semantic_tree_activation_counts(content_tokens: int, *, activation_steps: int, schedule: str=DEFAULT_SEMANTIC_TREE_SCHEDULE, shift: float=DEFAULT_SEMANTIC_TREE_SHIFT, device: torch.device | str | None=None) -> torch.Tensor: content_tokens = _plain_positive_int(content_tokens, name='content_tokens') activation_steps = min(_plain_positive_int(activation_steps, name='activation_steps'), content_tokens) if not isinstance(schedule, str): raise TypeError('schedule must be a string') schedule = schedule.strip().lower() if schedule not in SUPPORTED_SEMANTIC_TREE_INFERENCE_SCHEDULES: allowed = ', '.join(sorted(SUPPORTED_SEMANTIC_TREE_INFERENCE_SCHEDULES)) raise ValueError(f'unsupported semantic-tree inference schedule {schedule!r}; expected one of: {allowed}') shift = _plain_positive_float(shift, name='shift') sample_times = torch.linspace(0.0, 1.0, activation_steps + 1, dtype=torch.float32) if schedule == 'linear': sigmas = sample_times elif schedule == 'shift': sigmas = shift * sample_times / (1.0 + (shift - 1.0) * sample_times) elif schedule == 'logit_normal': sigmas = 0.5 * (1.0 + torch.erf(torch.logit(sample_times) / math.sqrt(2.0))) elif schedule == 'cosine': sigmas = 1.0 - 0.5 * (1.0 + torch.cos(math.pi * sample_times)) elif schedule == 'cosine2': sigmas = torch.sin(0.5 * math.pi * sample_times) else: sigmas = (1.0 - torch.exp(-5.0 * sample_times)) / (1.0 - math.exp(-5.0)) sigmas = sigmas.clamp(0.0001, 1.0) boundaries = (sigmas * content_tokens).to(dtype=torch.long) counts = (boundaries[1:] - boundaries[:-1]).clamp_min(1) excess = int(counts.sum().item()) - content_tokens cursor = 0 while excess > 0: cursor %= counts.numel() if int(counts[cursor].item()) == 1: cursor += 1 continue counts[cursor] -= 1 excess -= 1 cursor += 1 if int(counts.sum().item()) != content_tokens: raise RuntimeError('semantic-tree activation counts do not cover content') return counts.flip(0).to(device=device) def resolve_semantic_tree_atomic_block_size(value: object=1) -> int: block_size = _plain_positive_int(value, name='semantic_tree_atomic_block_size') if block_size not in {1, 2, 4, 8}: raise ValueError('semantic_tree_atomic_block_size currently supports 1, 2, 4, or 8') return block_size def resolve_semantic_tree_refinement_gap(value: object=0) -> int: gap = _plain_nonnegative_int(value, name='semantic_tree_refinement_gap') if gap > MAX_SEMANTIC_TREE_REFINEMENT_GAP: raise ValueError('semantic_tree_refinement_gap currently supports only 0, 1, or 2') return gap def semantic_tree_decode_schedule(activation_counts: torch.Tensor, *, refinement_steps: int, refinement_gap: int=0) -> torch.Tensor: if not torch.is_tensor(activation_counts): raise TypeError('activation_counts must be a tensor') if activation_counts.ndim != 2: raise ValueError('activation_counts must have shape [batch, rounds]') if activation_counts.dtype not in {torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64}: raise TypeError('activation_counts must use an integer dtype') if activation_counts.shape[1] == 0: raise ValueError('activation_counts must contain at least one round') if bool((activation_counts < 0).any()): raise ValueError('activation_counts must be non-negative') refinement_steps = _plain_positive_int(refinement_steps, name='refinement_steps') refinement_gap = resolve_semantic_tree_refinement_gap(refinement_gap) flush_rounds = refinement_steps - 1 + refinement_steps * refinement_gap if flush_rounds == 0: return activation_counts flush = torch.zeros(activation_counts.shape[0], flush_rounds, dtype=activation_counts.dtype, device=activation_counts.device) return torch.cat((activation_counts, flush), dim=1) def semantic_tree_refinement_due(ages: torch.Tensor, next_refinement_round: torch.Tensor, *, decode_round: int, refinement_steps: int) -> torch.Tensor: if not torch.is_tensor(ages) or not torch.is_tensor(next_refinement_round): raise TypeError('ages and next_refinement_round must be tensors') if ages.shape != next_refinement_round.shape: raise ValueError('ages and next_refinement_round must align') if ages.device != next_refinement_round.device: raise ValueError('ages and next_refinement_round must share a device') for (tensor, name) in ((ages, 'ages'), (next_refinement_round, 'next_refinement_round')): if tensor.dtype == torch.bool or tensor.is_floating_point(): raise TypeError(f'{name} must use an integer dtype') decode_round = _plain_nonnegative_int(decode_round, name='decode_round') refinement_steps = _plain_positive_int(refinement_steps, name='refinement_steps') return ages.ge(0) & ages.lt(refinement_steps) & next_refinement_round.le(decode_round) EmbeddingSource = torch.Tensor | Callable[[torch.Tensor], torch.Tensor] def _embed_ids(source: EmbeddingSource, token_ids: torch.Tensor) -> torch.Tensor: if torch.is_tensor(source): if source.ndim != 2 or not source.is_floating_point(): raise TypeError('embedding weight must be a floating-point matrix') if source.device != token_ids.device: raise ValueError('embedding weight and token IDs must share a device') if token_ids.numel() and bool(((token_ids < 0) | (token_ids >= source.shape[0])).any()): raise ValueError('token ID lies outside the embedding weight') return F.embedding(token_ids.long(), source) if not callable(source): raise TypeError('embedding source must be a weight tensor or callable') embedded = source(token_ids) if not torch.is_tensor(embedded): raise TypeError('embedding callback must return a tensor') return embedded def build_semantic_tree_embeddings(representatives: torch.Tensor, noise_levels: float | torch.Tensor, embedding_callback: EmbeddingSource, mask_token_id: int) -> torch.Tensor: if not torch.is_tensor(representatives): raise TypeError('representatives must be a tensor') if representatives.dtype == torch.bool or representatives.is_floating_point() or representatives.is_complex(): raise TypeError('representatives must contain integer code IDs') if isinstance(mask_token_id, bool) or not isinstance(mask_token_id, Integral): raise TypeError('mask_token_id must be an integer') mask_token_id = int(mask_token_id) if mask_token_id < 0: raise ValueError('mask_token_id must be nonnegative') representative_embeddings = _embed_ids(embedding_callback, representatives) expected_prefix = tuple(representatives.shape) if representative_embeddings.ndim != representatives.ndim + 1 or tuple(representative_embeddings.shape[:-1]) != expected_prefix or (not representative_embeddings.is_floating_point()): raise ValueError('representative embeddings must have shape representatives.shape + [H]') mask_ids = torch.tensor([mask_token_id], dtype=representatives.dtype, device=representatives.device) mask_embedding = _embed_ids(embedding_callback, mask_ids) if mask_embedding.shape != (1, representative_embeddings.shape[-1]): raise ValueError('MASK embedding must have shape [1, H]') if mask_embedding.device != representative_embeddings.device or mask_embedding.dtype != representative_embeddings.dtype: raise ValueError('representative and MASK embeddings must share dtype/device') noise = torch.as_tensor(noise_levels, device=representatives.device, dtype=torch.float32) if noise.ndim < representatives.ndim: noise = noise.reshape(noise.shape + (1,) * (representatives.ndim - noise.ndim)) try: noise = torch.broadcast_to(noise, representatives.shape) except RuntimeError as error: raise ValueError('noise_levels are not broadcastable to representatives') from error if not bool(torch.isfinite(noise).all()) or bool((noise < 0).any()) or bool((noise > 1).any()): raise ValueError('noise_levels must be finite and lie in [0, 1]') noise_in_embedding_dtype = noise.to(dtype=representative_embeddings.dtype) combined = representative_embeddings + noise_in_embedding_dtype.unsqueeze(-1) * mask_embedding return torch.where(noise.eq(0).unsqueeze(-1), representative_embeddings, combined) @dataclass(frozen=True, slots=True) class SemanticTreeAdvance: representatives: torch.Tensor ages: torch.Tensor levels: torch.Tensor noise_levels: torch.Tensor committed: torch.Tensor @torch.no_grad() def advance_revisable_semantic_tree(current_representatives: torch.Tensor, x0_proposals: torch.Tensor, ages: torch.Tensor, tree: BalancedSemanticTree, steps: int=DEFAULT_SEMANTIC_TREE_STEPS, generator: torch.Generator | None=None) -> SemanticTreeAdvance: if not isinstance(tree, BalancedSemanticTree): raise TypeError('tree must be a BalancedSemanticTree') tree._check_tokens(current_representatives, name='current_representatives') tree._check_tokens(x0_proposals, name='x0_proposals') if current_representatives.shape != x0_proposals.shape: raise ValueError('current_representatives and x0_proposals must align') if current_representatives.dtype != x0_proposals.dtype: raise ValueError('current_representatives and x0_proposals must share dtype') if not torch.is_tensor(ages) or ages.shape != current_representatives.shape: raise ValueError('ages must align with representative tokens') (validated_ages, steps) = _validate_ages(ages, steps=steps) if validated_ages.device != current_representatives.device: raise ValueError('ages and representative tokens must share a device') if steps > tree.depth: raise ValueError('steps cannot exceed tree depth') unfinished = validated_ages < steps next_ages = torch.where(unfinished, validated_ages + 1, validated_ages) next_levels = semantic_tree_levels_for_ages(next_ages, steps=steps, depth=tree.depth) proposal_representatives = tree.sample_representatives(x0_proposals, next_levels, generator=generator) next_representatives = torch.where(unfinished, proposal_representatives, current_representatives) return SemanticTreeAdvance(representatives=next_representatives, ages=next_ages.to(dtype=ages.dtype), levels=next_levels, noise_levels=semantic_tree_noise_for_ages(next_ages, steps=steps), committed=next_ages.eq(steps))