palimseste/tokenizer.py

"""PALIMPSESTE — Character-level tokenizer with vocab save/load.

For the hypervectorial substrate, a character-level tokenizer is the natural
choice: each character gets a stable atomic HV (via the :class:`Encoder`),
and text is encoded as a *sequence* (order-sensitive bundle of role-bound
char HVs). This keeps the vocabulary tiny and the encoding lossless, while
letting the associative memory learn character n-gram transitions.

This is deliberately **not** a BPE/SentencePiece tokenizer — those introduce
an opaque compression layer that fights the substrate's symbolic transparency
goal. Char-level keeps everything inspectable: you can always decode a
hypervector back to its nearest known characters.

The tokenizer implements the minimal HF-like interface needed for
``save_pretrained`` / ``from_pretrained`` compatibility:
``encode``, ``decode``, ``save_vocabulary``, ``vocab_size``, ``__len__``.
"""

from __future__ import annotations

from dataclasses import dataclass, field
import json
from pathlib import Path
import numpy as np

from .hv import HV, random_hv
from .learner import Encoder

__all__ = ["CharTokenizer", "VOCAB_SPECIAL"]


# Special tokens
VOCAB_SPECIAL = ["<pad>", "<bos>", "<eos>", "<unk>"]
PAD, BOS, EOS, UNK = 0, 1, 2, 3


@dataclass
class CharTokenizer:
    """Character-level tokenizer with a fixed atomic HV per character.

    The HVs are drawn from the shared :class:`Encoder` so that the same
    character always maps to the same hypervector, and the LM can bind them
    with positional roles to represent sequences.
    """

    encoder: Encoder
    # id -> char string (including special tokens at 0..3)
    id2char: list[str] = field(default_factory=lambda: list(VOCAB_SPECIAL))
    # char string -> id
    char2id: dict[str, int] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if not self.char2id:
            self.char2id = {c: i for i, c in enumerate(self.id2char)}
        # pre-register atom HVs for all known chars (lazy otherwise)
        for ch in self.id2char:
            self._char_hv(ch)

    # ----------------------------------------------------------- properties
    @property
    def vocab_size(self) -> int:
        return len(self.id2char)

    def __len__(self) -> int:
        return self.vocab_size

    @property
    def pad_token_id(self) -> int:
        return PAD

    @property
    def bos_token_id(self) -> int:
        return BOS

    @property
    def eos_token_id(self) -> int:
        return EOS

    # ----------------------------------------------------------- vocab build
    def build_vocab(self, text: str) -> None:
        """Extend the vocabulary with every character in ``text``."""
        for ch in text:
            if ch not in self.char2id:
                self.char2id[ch] = len(self.id2char)
                self.id2char.append(ch)
                self._char_hv(ch)  # materialize the atom

    # ----------------------------------------------------------- HV lookup
    def _char_hv(self, ch: str) -> HV:
        """Get the atomic HV for a character (creates it lazily)."""
        return self.encoder.encode_str(f"__char__{ch}")

    def token_hv(self, token_id: int) -> HV:
        """Get the HV for a token id."""
        return self._char_hv(self.id2char[token_id])

    # ----------------------------------------------------------- encode/decode
    def encode(self, text: str, add_bos: bool = False, add_eos: bool = False) -> list[int]:
        ids = []
        if add_bos:
            ids.append(BOS)
        for ch in text:
            ids.append(self.char2id.get(ch, UNK))
        if add_eos:
            ids.append(EOS)
        return ids

    def decode(self, ids: list[int]) -> str:
        out = []
        for i in ids:
            if i in (PAD, BOS, EOS):
                continue
            if 0 <= i < len(self.id2char):
                out.append(self.id2char[i])
            # else skip unknown
        return "".join(out)

    # ----------------------------------------------------------- sequence HV
    def encode_text_to_hv(self, text: str) -> HV:
        """Encode a string into a single order-sensitive hypervector.

        Each character is bound with its positional role, then bundled.
        """
        ids = self.encode(text)
        hvs = [self.token_hv(i) for i in ids]
        if not hvs:
            return self.token_hv(PAD)
        return self.encoder.encode_sequence(hvs)

    def encode_context(self, ids: list[int], window: int | None = None,
                       start_pos: int = 0) -> HV:
        """Encode a token-id context (window of recent tokens) into one HV.

        This is the ``s_t`` state used by the LM: the last ``window`` tokens,
        each bound with its positional role, bundled.

        Roles are **absolute positions modulo window** (not relative indices).
        This means the token at absolute position ``p`` gets role ``p % window``,
        so the state HV is identical whether computed incrementally or from
        scratch — essential for associative retrieval consistency between
        training and inference.

        Parameters
        ----------
        ids : list of token ids
        window : context window size (default: all)
        start_pos : absolute position of the *first* id in ``ids`` (default 0).
            Used when encoding a suffix of a longer sequence so the role
            assignment matches the absolute positions.
        """
        if window is not None:
            # keep the last `window` ids, but track their absolute positions
            if len(ids) > window:
                ids = ids[-window:]
                start_pos = start_pos + (len(ids) - window) if start_pos else 0
        hvs = [self.token_hv(i) for i in ids]
        if not hvs:
            return self.token_hv(PAD)
        # bind each token with role(start_pos + i) % window
        from .hv import bind, bundle, random_hv
        bound = []
        for i, x in enumerate(hvs):
            role_idx = (start_pos + i) % window if window else i
            role_hv = self.encoder._role(role_idx)
            bound.append(bind(x, role_hv))
        return bundle(bound, rng=self.encoder.rng, deterministic=True)

    # ----------------------------------------------------------- save/load
    def save_vocabulary(self, path: str | Path) -> tuple[str]:
        """Save the vocab as ``vocab.json`` (HF-compatible name)."""
        p = Path(path)
        p.parent.mkdir(parents=True, exist_ok=True)
        with open(p, "w", encoding="utf-8") as f:
            json.dump({"id2char": self.id2char}, f, indent=2, ensure_ascii=False)
        return (str(p),)

    @classmethod
    def load_vocabulary(cls, path: str | Path, encoder: Encoder) -> "CharTokenizer":
        p = Path(path)
        with open(p, "r", encoding="utf-8") as f:
            d = json.load(f)
        id2char = d["id2char"]
        tok = cls(encoder=encoder, id2char=id2char)
        return tok