source/src/strata/tokenization/lattice.py

"""Reversible lattice tokenizers with stable character and byte anchoring."""

from __future__ import annotations

import bisect
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Mapping, Sequence

from strata.tokenization.special_tokens import DEFAULT_SPECIAL_TOKENS
from strata.types import AnchoredSpan


class TokenizerError(ValueError):
    """Raised when tokenizer inputs or outputs violate STRATA invariants."""


@dataclass(frozen=True, slots=True)
class TokenSpan:
    """A token with reversible text offsets."""

    token_index: int
    token_id: int
    piece: str
    byte_start: int
    byte_end: int
    char_start: int
    char_end: int
    is_special: bool = False

    @property
    def byte_length(self) -> int:
        return self.byte_end - self.byte_start

    @property
    def char_length(self) -> int:
        return self.char_end - self.char_start

    @property
    def anchor(self) -> AnchoredSpan:
        return AnchoredSpan.from_offsets(
            char_start=self.char_start,
            char_end=self.char_end,
            byte_start=self.byte_start,
            byte_end=self.byte_end,
        )


@dataclass(frozen=True, slots=True)
class WordSpan:
    """Whitespace-delimited surface word span used for graph alignment."""

    word_index: int
    text: str
    byte_start: int
    byte_end: int
    char_start: int
    char_end: int

    @property
    def anchor(self) -> AnchoredSpan:
        return AnchoredSpan.from_offsets(
            char_start=self.char_start,
            char_end=self.char_end,
            byte_start=self.byte_start,
            byte_end=self.byte_end,
        )


@dataclass(frozen=True, slots=True)
class LatticeEncoding:
    """Encoded text plus the span lattice needed by graph supervision."""

    text: str
    input_ids: tuple[int, ...]
    token_spans: tuple[TokenSpan, ...]
    word_spans: tuple[WordSpan, ...]
    char_to_byte: tuple[int, ...]
    tokenizer_name: str

    @property
    def attention_mask(self) -> tuple[int, ...]:
        return tuple(1 for _ in self.input_ids)

    @property
    def byte_length(self) -> int:
        return len(self.text.encode("utf-8"))

    @property
    def char_length(self) -> int:
        return len(self.text)

    def visible_tokens(self, *, byte_end: int) -> tuple[TokenSpan, ...]:
        """Return non-special tokens fully visible at a prefix byte boundary."""

        if byte_end < 0 or byte_end > self.byte_length:
            raise TokenizerError(
                f"byte_end must be within [0, {self.byte_length}], got {byte_end}"
            )
        return tuple(
            span
            for span in self.token_spans
            if not span.is_special and span.byte_end <= byte_end
        )

    def token_indices_for_anchor(self, anchor: AnchoredSpan) -> tuple[int, ...]:
        """Return token indices whose byte spans overlap an anchored graph node."""

        return tuple(
            span.token_index
            for span in self.token_spans
            if not span.is_special
            and span.byte_start < anchor.byte.end
            and anchor.byte.start < span.byte_end
        )


def _validate_special_tokens(special_tokens: Sequence[str]) -> None:
    if len(set(special_tokens)) != len(special_tokens):
        raise TokenizerError("special tokens must be unique")
    for token in special_tokens:
        if not token or not token.startswith("<") or not token.endswith(">"):
            raise TokenizerError(
                f"special token {token!r} must use angle-bracket namespace"
            )


def _char_to_byte_offsets(text: str) -> tuple[int, ...]:
    offsets = [0]
    byte_position = 0
    for char in text:
        byte_position += len(char.encode("utf-8"))
        offsets.append(byte_position)
    return tuple(offsets)


def _char_span_to_byte_span(
    char_start: int,
    char_end: int,
    char_to_byte: Sequence[int],
) -> tuple[int, int]:
    text_length = len(char_to_byte) - 1
    if char_start < 0 or char_end < char_start or char_end > text_length:
        raise TokenizerError(
            f"invalid char span [{char_start}, {char_end}) for text length "
            f"{text_length}"
        )
    return char_to_byte[char_start], char_to_byte[char_end]


def _byte_span_to_char_span(
    byte_start: int,
    byte_end: int,
    char_to_byte: Sequence[int],
) -> tuple[int, int]:
    if byte_start < 0 or byte_end < byte_start or byte_end > char_to_byte[-1]:
        raise TokenizerError(
            f"invalid byte span [{byte_start}, {byte_end}) for byte length "
            f"{char_to_byte[-1]}"
        )
    if byte_start == byte_end:
        char = bisect.bisect_right(char_to_byte, byte_start) - 1
        char = max(0, min(char, len(char_to_byte) - 1))
        return char, char

    char_start = bisect.bisect_right(char_to_byte, byte_start) - 1
    char_end = bisect.bisect_left(char_to_byte, byte_end)
    if char_end <= char_start:
        char_end = char_start + 1
    return char_start, min(char_end, len(char_to_byte) - 1)


_WORD_PATTERN = re.compile(r"\S+")


def _word_spans(text: str, char_to_byte: Sequence[int]) -> tuple[WordSpan, ...]:
    spans: list[WordSpan] = []
    for word_index, match in enumerate(_WORD_PATTERN.finditer(text)):
        byte_start, byte_end = _char_span_to_byte_span(
            match.start(), match.end(), char_to_byte
        )
        spans.append(
            WordSpan(
                word_index=word_index,
                text=match.group(0),
                byte_start=byte_start,
                byte_end=byte_end,
                char_start=match.start(),
                char_end=match.end(),
            )
        )
    return tuple(spans)


class ByteLatticeTokenizer:
    """UTF-8 byte tokenizer that always roundtrips and preserves spans.

    This tokenizer is intentionally simple and production-safe. It is a reliable
    fallback before a trained SentencePiece model exists, and it is useful for
    debugging graph alignment because every byte position is represented.
    """

    name = "strata-byte-lattice"

    def __init__(self, special_tokens: Sequence[str] = DEFAULT_SPECIAL_TOKENS) -> None:
        _validate_special_tokens(special_tokens)
        self.special_tokens = tuple(special_tokens)
        self.special_to_id = {token: idx for idx, token in enumerate(special_tokens)}
        self.id_to_special = {idx: token for token, idx in self.special_to_id.items()}
        self.byte_offset = len(self.special_tokens)
        self.vocab_size = self.byte_offset + 256

    @property
    def pad_token_id(self) -> int:
        return self.special_to_id["<PAD>"]

    @property
    def bos_token_id(self) -> int:
        return self.special_to_id["<BOS>"]

    @property
    def eos_token_id(self) -> int:
        return self.special_to_id["<EOS>"]

    def token_to_id(self, token: str) -> int:
        if token in self.special_to_id:
            return self.special_to_id[token]
        if re.fullmatch(r"<0x[0-9A-Fa-f]{2}>", token):
            return self.byte_offset + int(token[3:5], 16)
        raise TokenizerError(f"unknown token {token!r}")

    def id_to_token(self, token_id: int) -> str:
        if token_id in self.id_to_special:
            return self.id_to_special[token_id]
        if self.byte_offset <= token_id < self.byte_offset + 256:
            return f"<0x{token_id - self.byte_offset:02X}>"
        raise TokenizerError(f"token id {token_id} is outside vocab size {self.vocab_size}")

    def encode(
        self,
        text: str,
        *,
        add_bos: bool = False,
        add_eos: bool = False,
    ) -> LatticeEncoding:
        if not isinstance(text, str):
            raise TypeError(f"text must be str, got {type(text).__name__}")

        char_to_byte = _char_to_byte_offsets(text)
        data = text.encode("utf-8")
        input_ids: list[int] = []
        token_spans: list[TokenSpan] = []

        def append_special(token: str) -> None:
            token_id = self.special_to_id[token]
            token_spans.append(
                TokenSpan(
                    token_index=len(input_ids),
                    token_id=token_id,
                    piece=token,
                    byte_start=0,
                    byte_end=0,
                    char_start=0,
                    char_end=0,
                    is_special=True,
                )
            )
            input_ids.append(token_id)

        if add_bos:
            append_special("<BOS>")

        for byte_index, byte_value in enumerate(data):
            char_start, char_end = _byte_span_to_char_span(
                byte_index, byte_index + 1, char_to_byte
            )
            token_id = self.byte_offset + byte_value
            token_spans.append(
                TokenSpan(
                    token_index=len(input_ids),
                    token_id=token_id,
                    piece=f"<0x{byte_value:02X}>",
                    byte_start=byte_index,
                    byte_end=byte_index + 1,
                    char_start=char_start,
                    char_end=char_end,
                )
            )
            input_ids.append(token_id)

        if add_eos:
            append_special("<EOS>")

        return LatticeEncoding(
            text=text,
            input_ids=tuple(input_ids),
            token_spans=tuple(token_spans),
            word_spans=_word_spans(text, char_to_byte),
            char_to_byte=char_to_byte,
            tokenizer_name=self.name,
        )

    def decode(
        self,
        input_ids: Iterable[int],
        *,
        skip_special_tokens: bool = True,
        errors: str = "strict",
    ) -> str:
        data = bytearray()
        parts: list[str] = []

        def flush_data() -> None:
            if data:
                parts.append(bytes(data).decode("utf-8", errors=errors))
                data.clear()

        for token_id in input_ids:
            if self.byte_offset <= token_id < self.byte_offset + 256:
                data.append(token_id - self.byte_offset)
            elif token_id in self.id_to_special:
                if not skip_special_tokens:
                    flush_data()
                    parts.append(self.id_to_special[token_id])
            else:
                raise TokenizerError(
                    f"token id {token_id} is outside vocab size {self.vocab_size}"
                )
        flush_data()
        return "".join(parts)


class SentencePieceLatticeTokenizer:
    """SentencePiece tokenizer wrapper that preserves proto-provided offsets."""

    name = "strata-sentencepiece-lattice"

    def __init__(
        self,
        model_file: str | Path,
        *,
        special_tokens: Sequence[str] = DEFAULT_SPECIAL_TOKENS,
        require_exact_roundtrip: bool = True,
    ) -> None:
        try:
            import sentencepiece as spm
        except ImportError as exc:  # pragma: no cover - dependency declared.
            raise TokenizerError("sentencepiece is required for this tokenizer") from exc

        _validate_special_tokens(special_tokens)
        self.model_file = Path(model_file)
        if not self.model_file.exists():
            raise TokenizerError(f"SentencePiece model does not exist: {self.model_file}")

        self.processor = spm.SentencePieceProcessor(model_file=str(self.model_file))
        self.special_tokens = tuple(special_tokens)
        self.special_to_id = {token: idx for idx, token in enumerate(special_tokens)}
        self.id_to_special = {idx: token for token, idx in self.special_to_id.items()}
        self.sp_offset = len(self.special_tokens)
        self.sp_vocab_size = int(self.processor.vocab_size())
        self.vocab_size = self.sp_offset + self.sp_vocab_size
        self.require_exact_roundtrip = require_exact_roundtrip

    @property
    def pad_token_id(self) -> int:
        return self.special_to_id["<PAD>"]

    @property
    def bos_token_id(self) -> int:
        return self.special_to_id["<BOS>"]

    @property
    def eos_token_id(self) -> int:
        return self.special_to_id["<EOS>"]

    def encode(
        self,
        text: str,
        *,
        add_bos: bool = False,
        add_eos: bool = False,
    ) -> LatticeEncoding:
        if not isinstance(text, str):
            raise TypeError(f"text must be str, got {type(text).__name__}")

        char_to_byte = _char_to_byte_offsets(text)
        # SentencePiece 0.2.2 removed the legacy immutable-proto wrapper.  Its
        # offset mapping exposes the same IDs, pieces, and character spans
        # without adding a protobuf runtime dependency.
        try:
            encoded = self.processor.Encode(text, return_type="offset_mapping")
            proto_pieces = zip(
                encoded["ids"], encoded["pieces"], encoded["offsets"], strict=True
            )
        except (TypeError, ValueError):  # pragma: no cover - pre-0.2 compatibility.
            proto = self.processor.EncodeAsImmutableProto(text)
            proto_pieces = (
                (piece.id, piece.piece, (piece.begin, piece.end)) for piece in proto.pieces
            )
        input_ids: list[int] = []
        token_spans: list[TokenSpan] = []

        def append_special(token: str) -> None:
            token_id = self.special_to_id[token]
            token_spans.append(
                TokenSpan(
                    token_index=len(input_ids),
                    token_id=token_id,
                    piece=token,
                    byte_start=0,
                    byte_end=0,
                    char_start=0,
                    char_end=0,
                    is_special=True,
                )
            )
            input_ids.append(token_id)

        if add_bos:
            append_special("<BOS>")

        for piece_id, piece_text, offsets in proto_pieces:
            char_start, char_end = (int(value) for value in offsets)
            byte_start, byte_end = _char_span_to_byte_span(
                char_start, char_end, char_to_byte
            )
            token_id = self.sp_offset + int(piece_id)
            token_spans.append(
                TokenSpan(
                    token_index=len(input_ids),
                    token_id=token_id,
                    piece=str(piece_text),
                    byte_start=byte_start,
                    byte_end=byte_end,
                    char_start=char_start,
                    char_end=char_end,
                )
            )
            input_ids.append(token_id)

        if add_eos:
            append_special("<EOS>")

        if self.require_exact_roundtrip:
            decoded = self.decode(input_ids)
            if decoded != text:
                raise TokenizerError(
                    "SentencePiece model is not exact-roundtrip for this text; "
                    "train with identity normalization and preserved whitespace. "
                    f"decoded={decoded!r}, original={text!r}"
                )

        return LatticeEncoding(
            text=text,
            input_ids=tuple(input_ids),
            token_spans=tuple(token_spans),
            word_spans=_word_spans(text, char_to_byte),
            char_to_byte=char_to_byte,
            tokenizer_name=self.name,
        )

    def decode(
        self,
        input_ids: Iterable[int],
        *,
        skip_special_tokens: bool = True,
    ) -> str:
        sp_ids: list[int] = []
        parts: list[str] = []

        def flush_sp_ids() -> None:
            if sp_ids:
                parts.append(self.processor.DecodeIds(sp_ids))
                sp_ids.clear()

        for token_id in input_ids:
            if self.sp_offset <= token_id < self.sp_offset + self.sp_vocab_size:
                sp_ids.append(token_id - self.sp_offset)
            elif token_id in self.id_to_special:
                if not skip_special_tokens:
                    flush_sp_ids()
                    parts.append(self.id_to_special[token_id])
            else:
                raise TokenizerError(
                    f"token id {token_id} is outside vocab size {self.vocab_size}"
                )
        flush_sp_ids()
        return "".join(parts)


def special_token_ids(tokenizer: ByteLatticeTokenizer | SentencePieceLatticeTokenizer) -> Mapping[str, int]:
    """Return a copy of the tokenizer's reserved-token ID mapping."""

    return dict(tokenizer.special_to_id)