from tokenizers import Tokenizer
import torch

class AobanTokenizer:

    def __init__(self, path="lightning_tokenizer.json"):

        self.tokenizer = Tokenizer.from_file(path)


        self.pad_id = self.tokenizer.token_to_id("<PAD>")

        self.eot_id = self.tokenizer.token_to_id("<|endoftext|>")


        if self.pad_id is None or self.eot_id is None:

            raise ValueError("Tokenizer missing special tokens")


    def encode(self, text, max_len=None):

        ids = self.tokenizer.encode(text).ids


        if max_len is not None:

            ids = ids[:max_len]

            ids += [self.pad_id] * (max_len - len(ids))


        return torch.tensor(ids, dtype=torch.long)


    def decode(self, indices):

        if isinstance(indices, torch.Tensor):

            indices = indices.tolist()

        return self.tokenizer.decode(indices)


    @property

    def vocab_size(self):

        return self.tokenizer.get_vocab_size()