| 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() |
