from transformers import AutoTokenizer

# Qwen3's tokenizer: BPE, ~151936 vocab, already strong on English/code/math
QWEN_TOKENIZER = "Qwen/Qwen3-4B"


class Tokenizer:
    """Wraps Qwen3's pretrained tokenizer for the data pipeline."""

    def __init__(self, name: str = QWEN_TOKENIZER):
        self.tok = AutoTokenizer.from_pretrained(name)
        # end-of-text id marks document boundaries when we pack tokens together
        self.eot_id = self.tok.eos_token_id

    def encode(self, text: str) -> list[int]:
        # plain token ids, no special tokens added automatically
        return self.tok.encode(text, add_special_tokens=False)

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

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