File size: 3,373 Bytes
5d2c747 |
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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
"""
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 # Default value for unknown tokens
self.vocab = self.get_vocab()
self.padding_side = "right" # Default padding side
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."""
# check if the tokens are a tensor
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 |