"""
Tokenizer wrapper around GPT-2 tiktoken BPE
"""

import tiktoken
from typing import List, Union


class Tokenizer:
    def __init__(self, model_name: str = "gpt2"):
        self.enc = tiktoken.get_encoding(model_name)
        self.vocab_size = self.enc.n_vocab  # 50,257
        self.eot_token = self.enc.eot_token  # 50,256 (<|endoftext|>)

    def encode(self, text: str) -> List[int]:
        """Encodes string to list of token IDs."""
        return self.enc.encode_ordinary(text)

    def decode(self, tokens: Union[List[int], List[List[int]]]) -> str:
        """Decodes list of token IDs to text string."""
        if isinstance(tokens, list) and len(tokens) > 0 and isinstance(tokens[0], list):
            tokens = tokens[0]
        return self.enc.decode(tokens)


def get_tokenizer():
    return Tokenizer()