| """ |
| A simple wrapper around the GPT2 Tokenizer to |
| standardize the interface for tokenization. |
| """ |
|
|
| import tiktoken |
| import torch |
| from transformers import GPT2Tokenizer, GPT2LMHeadModel |
|
|
| class BaseTokenizer: |
| """A simple wrapper around the GPT2 BPE Tokenizer.""" |
|
|
| def __init__(self): |
| gpt2_tokenizer = tiktoken.get_encoding("gpt2") |
| self.tokenizer = gpt2_tokenizer |
| self.eot_token = gpt2_tokenizer.eot_token |
| self.pad_token = gpt2_tokenizer.eot_token |
| self.pad_token_id = gpt2_tokenizer.eot_token |
| self.vocab_size = gpt2_tokenizer.max_token_value + 1 |
| self.name_or_path = "gpt2_fast" |
| self.eos_token_id = gpt2_tokenizer.eot_token |
| self.unk_token_id = -1 |
| self.vocab = self.get_vocab() |
| self.padding_side = "right" |
|
|
| def pad_batch(self, token_lists, direction="right"): |
| """Pad a list of token lists to the same length, |
| and return the padded tensor, and mask tensor. |
| |
| Direction can be 'right' or 'left' to specifysba the padding direction. |
| """ |
| max_len = max(len(tokens) for tokens in token_lists) |
| padded_tokens = [] |
| mask = [] |
| for tokens in token_lists: |
| if direction == "right": |
| padded_tokens.append(tokens + [self.pad_token] * (max_len - len(tokens))) |
| mask.append([1] * len(tokens) + [0] * (max_len - len(tokens))) |
| elif direction == "left": |
| padded_tokens.append([self.pad_token] * (max_len - len(tokens)) + tokens) |
| mask.append([0] * (max_len - len(tokens)) + [1] * len(tokens)) |
| return torch.tensor(padded_tokens), torch.tensor(mask) |
|
|
| def encode(self, text, add_special_tokens=False, **kwargs): |
| """Encode a string into tokens.""" |
| return self.tokenizer.encode_ordinary(text) |
|
|
| def encode_batch(self, texts): |
| """Encode a list of strings into tokens.""" |
| return self.tokenizer.encode_ordinary_batch(texts) |
|
|
| def decode(self, tokens, **kwargs): |
| """Decode a list of tokens into a string.""" |
| |
| if torch.is_tensor(tokens): |
| tokens = tokens.tolist() |
| return self.tokenizer.decode(tokens) |
|
|
| def decode_batch(self, token_lists): |
| """Decode a list of token lists into a list of strings.""" |
| if torch.is_tensor(token_lists): |
| token_lists = token_lists.tolist() |
| return self.tokenizer.decode_batch(token_lists) |
|
|
| def convert_ids_to_tokens(self, token_ids): |
| """Convert a list of token ids to tokens.""" |
| return [self.tokenizer.decode([tid]) for tid in token_ids] |
| |
| def convert_tokens_to_ids(self, tokens): |
| """Convert a list of tokens to token ids.""" |
| return [self.tokenizer.encode_ordinary(t)[0] for t in tokens] |
| |
| def convert_tokens_to_string(self, tokens): |
| """Convert a list of token strings into a single decoded string.""" |
| token_ids = self.convert_tokens_to_ids(tokens) |
| return self.decode(token_ids) |
|
|
| |
| def get_vocab(self): |
| """Get the vocabulary as a dict of token to id.""" |
| vocab = {} |
| for i in range(self.vocab_size): |
| token = self.tokenizer.decode([i]) |
| vocab[token] = i |
| return vocab |
