src/eaimath/embedding_predictor.py

"""Tiny deterministic predictor for vocabulary rows omitted from an artifact.

The predictor borrows the idea of intra-frame prediction from codecs: a token's
string is side information already present in the tokenizer, so only a small
linear dictionary has to be stored. Exact retained rows are scattered over the
prediction during restoration.
"""

from __future__ import annotations

import math
from typing import Any

import torch

from .pack import pack_tensor, state_dict_bytes, unpack_tensor

FEATURE_DIM = 272


def _token_strings(tokenizer: Any, ids: list[int]) -> list[str]:
    tokens = tokenizer.convert_ids_to_tokens(ids)
    if isinstance(tokens, str):
        tokens = [tokens]
    return ["" if token is None else str(token) for token in tokens]


def token_string_features(tokenizer: Any, ids: torch.Tensor) -> torch.Tensor:
    """Return tokenizer-derived features without storing anything per token.

    Layout: 256 normalized UTF-8 byte counts, eight log-length buckets, four
    first-byte classes and four last-byte classes. The last two blocks retain a
    little order information while keeping the dictionary under one megabyte.
    """
    flat_ids = ids.detach().cpu().long().reshape(-1)
    strings = _token_strings(tokenizer, [int(value) for value in flat_ids.tolist()])
    features = torch.zeros((len(strings), FEATURE_DIM), dtype=torch.float32)
    for row, token in enumerate(strings):
        encoded = token.encode("utf-8", errors="replace") or b"\x00"
        byte_ids = torch.tensor(list(encoded), dtype=torch.int64)
        counts = torch.bincount(byte_ids, minlength=256).float()
        features[row, :256] = counts / counts.square().sum().sqrt().clamp_min(1.0)
        length_bucket = min(7, int(math.log2(max(1, len(encoded)))))
        features[row, 256 + length_bucket] = 1.0
        features[row, 264 + encoded[0] // 64] = 1.0
        features[row, 268 + encoded[-1] // 64] = 1.0
    return features


def _sample_ids(vocab_size: int, sample_size: int, offset: float = 0.0) -> torch.Tensor:
    count = min(vocab_size, sample_size)
    if count == vocab_size:
        return torch.arange(vocab_size, dtype=torch.int64)
    positions = (torch.arange(count, dtype=torch.float64) + offset) * vocab_size / count
    return positions.floor().clamp_max(vocab_size - 1).long().unique()


def fit_token_predictor(
    weight: torch.Tensor,
    tokenizer: Any,
    *,
    sample_size: int = 65_536,
    heldout_size: int = 4_096,
    ridge: float = 1e-2,
    device: str = "cpu",
) -> tuple[dict[str, Any], dict[str, float | int | str]]:
    """Fit and INT8-pack a ridge dictionary for a full embedding matrix."""
    if weight.ndim != 2 or weight.shape[0] < 2:
        raise ValueError(f"expected a vocabulary matrix, found {tuple(weight.shape)}")
    if ridge <= 0:
        raise ValueError("ridge must be positive")
    train_ids = _sample_ids(weight.shape[0], sample_size)
    x = token_string_features(tokenizer, train_ids).to(device)
    y = weight.detach().cpu().index_select(0, train_ids).float().to(device)
    gram = x.T @ x
    gram.diagonal().add_(ridge)
    basis = torch.linalg.solve(gram, x.T @ y).cpu()

    heldout_ids = _sample_ids(weight.shape[0], heldout_size, offset=0.5)
    # Exclude any collision with the evenly spaced training grid.
    train_set = set(int(value) for value in train_ids.tolist())
    heldout_list = [int(value) for value in heldout_ids.tolist() if int(value) not in train_set]
    if not heldout_list:
        heldout_list = [int(train_ids[-1])]
    heldout_ids = torch.tensor(heldout_list, dtype=torch.int64)
    hx = token_string_features(tokenizer, heldout_ids)
    target = weight.detach().cpu().index_select(0, heldout_ids).float()
    prediction = hx @ basis
    cosine = torch.nn.functional.cosine_similarity(prediction, target, dim=1).mean()
    denominator = target.square().mean().sqrt().clamp_min(1e-12)
    relative_rmse = (prediction - target).square().mean().sqrt() / denominator

    packed_basis = pack_tensor(basis, bits=8, group_size=256)
    entry: dict[str, Any] = {
        "kind": "token_string_ridge_v1",
        "feature_dim": FEATURE_DIM,
        "basis": packed_basis,
        "vocab_size": int(weight.shape[0]),
        "hidden_size": int(weight.shape[1]),
        "ridge": float(ridge),
        "sample_size": int(train_ids.numel()),
    }
    report: dict[str, float | int | str] = {
        "kind": entry["kind"],
        "stored_bytes": state_dict_bytes({"basis": packed_basis}),
        "sample_size": int(train_ids.numel()),
        "heldout_size": int(heldout_ids.numel()),
        "heldout_mean_cosine": float(cosine),
        "heldout_relative_rmse": float(relative_rmse),
    }
    return entry, report


def predict_token_rows(
    entry: dict[str, Any],
    tokenizer: Any,
    *,
    batch_size: int = 8_192,
) -> torch.Tensor:
    """Decode a predictor into the original dense vocabulary matrix."""
    if entry.get("kind") != "token_string_ridge_v1":
        raise ValueError(f"unsupported token predictor: {entry.get('kind')!r}")
    basis = unpack_tensor(entry["basis"]).float()
    rows = []
    for start in range(0, int(entry["vocab_size"]), batch_size):
        stop = min(int(entry["vocab_size"]), start + batch_size)
        ids = torch.arange(start, stop, dtype=torch.int64)
        rows.append((token_string_features(tokenizer, ids) @ basis).to(torch.bfloat16))
    return torch.cat(rows, dim=0)