src/tokenizer/__init__.py

"""API publique du tokenizer Cortex-0.

Exemple ::

    from src.tokenizer import load_tokenizer

    tok = load_tokenizer()                 # tokenizer/ à la racine du dépôt
    ids = tok.encode("Bonjour le monde")
    texte = tok.decode(ids)
"""

from __future__ import annotations

from pathlib import Path

from tokenizers import Tokenizer

from src.tokenizer.special_tokens import (
    BOS_ID,
    BOS_TOKEN,
    EOS_ID,
    EOS_TOKEN,
    PAD_ID,
    PAD_TOKEN,
    SPECIAL_TOKEN_IDS,
    TOKENIZER_DIR,
    TOKENIZER_FILE,
    UNK_ID,
    UNK_TOKEN,
    VOCAB_SIZE,
)

ROOT = Path(__file__).resolve().parents[2]
DEFAULT_TOKENIZER_DIR = ROOT / TOKENIZER_DIR


class CortexTokenizer:
    """Enveloppe fine autour de `tokenizers.Tokenizer` (BPE ByteLevel)."""

    def __init__(self, tokenizer: Tokenizer):
        self._tok = tokenizer

    @classmethod
    def load(cls, path: str | Path | None = None) -> CortexTokenizer:
        """Charge depuis un dossier, un fichier tokenizer.json, ou le dépôt."""
        if path is None:
            target = DEFAULT_TOKENIZER_DIR / TOKENIZER_FILE
        else:
            target = Path(path)
            if target.is_dir():
                target = target / TOKENIZER_FILE
        if not target.is_file():
            raise FileNotFoundError(
                f"tokenizer introuvable : {target}. "
                "Entraîne-le d'abord : python -m src.tokenizer.train_tokenizer --help"
            )
        return cls(Tokenizer.from_file(str(target)))

    @property
    def vocab_size(self) -> int:
        return self._tok.get_vocab_size()

    @property
    def raw(self) -> Tokenizer:
        return self._tok

    def token_to_id(self, token: str) -> int | None:
        return self._tok.token_to_id(token)

    def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
        encoding = self._tok.encode(text, add_special_tokens=add_special_tokens)
        return list(encoding.ids)

    def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str:
        return self._tok.decode(ids, skip_special_tokens=skip_special_tokens)


def load_tokenizer(path: str | Path | None = None) -> CortexTokenizer:
    """Raccourci de chargement (voir `CortexTokenizer.load`)."""
    return CortexTokenizer.load(path)


__all__ = [
    "BOS_ID",
    "BOS_TOKEN",
    "EOS_ID",
    "EOS_TOKEN",
    "PAD_ID",
    "PAD_TOKEN",
    "SPECIAL_TOKEN_IDS",
    "UNK_ID",
    "UNK_TOKEN",
    "VOCAB_SIZE",
    "CortexTokenizer",
    "load_tokenizer",
]