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