cortex/tokenizer.py

"""Thin wrapper around the custom CortexSym tokenizer (tokenizers library)."""

from __future__ import annotations

from pathlib import Path

from tokenizers import Tokenizer

from .config import BOS_ID, EOT_ID, PAD_ID, UNK_ID

_DEFAULT = Path(__file__).resolve().parent.parent / "tokenizer" / "tokenizer.json"


class CortexTokenizer:
    def __init__(self, path: str | Path = _DEFAULT):
        self.tok = Tokenizer.from_file(str(path))
        self.bos, self.eot, self.pad, self.unk = BOS_ID, EOT_ID, PAD_ID, UNK_ID

    @property
    def vocab_size(self) -> int:
        return self.tok.get_vocab_size(with_added_tokens=True)

    def encode(self, text: str, add_bos: bool = False, add_eot: bool = True) -> list[int]:
        ids = self.tok.encode(text).ids
        if add_bos:
            ids = [self.bos] + ids
        if add_eot:
            ids = ids + [self.eot]
        return ids

    def decode(self, ids) -> str:
        return self.tok.decode(list(ids))