tree.py

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))