src/oxmini/tokenizer.py

"""A deliberately small character-level tokenizer."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Iterable


class CharTokenizer:
    """Maps individual Unicode characters to integer token IDs."""

    unk_token = "<unk>"

    def __init__(self, itos: Iterable[str]) -> None:
        values = list(itos)
        # Reserve ID 0 for out-of-vocabulary input. In a deterministic corpus
        # split this normally remains unused, but it makes interactive prompts
        # safe when they contain a character absent from training.
        if not values or values[0] != self.unk_token:
            values = [self.unk_token, *[item for item in values if item != self.unk_token]]
        if len(values) != len(set(values)):
            raise ValueError("token vocabulary contains duplicates")
        self.itos = values
        self.stoi = {token: index for index, token in enumerate(values)}

    @classmethod
    def build(cls, text: str) -> "CharTokenizer":
        if not text:
            raise ValueError("cannot build a tokenizer from empty text")
        # Sorting gives a stable vocabulary independent of set iteration order.
        return cls(sorted(set(text)))

    @property
    def vocab_size(self) -> int:
        return len(self.itos)

    def encode(self, text: str) -> list[int]:
        unk_id = self.stoi[self.unk_token]
        return [self.stoi.get(char, unk_id) for char in text]

    def decode(self, ids: Iterable[int]) -> str:
        pieces: list[str] = []
        for token_id in ids:
            if not 0 <= int(token_id) < self.vocab_size:
                raise ValueError(f"token ID {token_id} is outside the vocabulary")
            token = self.itos[int(token_id)]
            pieces.append("�" if token == self.unk_token else token)
        return "".join(pieces)

    def to_dict(self) -> dict[str, object]:
        return {
            "type": "character-level",
            "unk_token": self.unk_token,
            "vocab_size": self.vocab_size,
            "itos": self.itos,
            "stoi": self.stoi,
        }

    def save(self, path: str | Path) -> None:
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open("w", encoding="utf-8") as handle:
            json.dump(self.to_dict(), handle, ensure_ascii=False, indent=2, sort_keys=True)
            handle.write("\n")

    @classmethod
    def load(cls, path: str | Path) -> "CharTokenizer":
        with Path(path).open("r", encoding="utf-8") as handle:
            values = json.load(handle)
        if values.get("type") != "character-level" or not isinstance(values.get("itos"), list):
            raise ValueError(f"unsupported tokenizer file: {path}")
        return cls(values["itos"])