File size: 823 Bytes
1e4fb05 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | """
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()
|