Tokenizing and encoding/decoding text using the Tiktoken tokenizer.
| 32 | |
| 33 | |
| 34 | class Tokenizer: |
| 35 | """ |
| 36 | Tokenizing and encoding/decoding text using the Tiktoken tokenizer. |
| 37 | """ |
| 38 | |
| 39 | special_tokens: Dict[str, int] |
| 40 | |
| 41 | num_reserved_special_tokens = 256 |
| 42 | |
| 43 | pat_str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+" # noqa: E501 |
| 44 | |
| 45 | def __init__(self, model_path: str): |
| 46 | """ |
| 47 | Initializes the Tokenizer with a Tiktoken model. |
| 48 | |
| 49 | Args: |
| 50 | model_path (str): The path to the Tiktoken model file. |
| 51 | """ |
| 52 | assert os.path.isfile(model_path), model_path |
| 53 | |
| 54 | mergeable_ranks = load_tiktoken_bpe(model_path) |
| 55 | num_base_tokens = len(mergeable_ranks) |
| 56 | special_tokens = [ |
| 57 | "<|begin_of_text|>", |
| 58 | "<|end_of_text|>", |
| 59 | "<|reserved_special_token_0|>", |
| 60 | "<|reserved_special_token_1|>", |
| 61 | "<|reserved_special_token_2|>", |
| 62 | "<|reserved_special_token_3|>", |
| 63 | "<|start_header_id|>", |
| 64 | "<|end_header_id|>", |
| 65 | "<|reserved_special_token_4|>", |
| 66 | "<|eot_id|>", # end of turn |
| 67 | ] + [ |
| 68 | f"<|reserved_special_token_{i}|>" |
| 69 | for i in range(5, self.num_reserved_special_tokens - 5) |
| 70 | ] |
| 71 | self.special_tokens = { |
| 72 | token: num_base_tokens + i for i, token in enumerate(special_tokens) |
| 73 | } |
| 74 | self.model = tiktoken.Encoding( |
| 75 | name=Path(model_path).name, |
| 76 | pat_str=self.pat_str, |
| 77 | mergeable_ranks=mergeable_ranks, |
| 78 | special_tokens=self.special_tokens, |
| 79 | ) |
| 80 | logger.info(f"Reloaded tiktoken model from {model_path}") |
| 81 | |
| 82 | self.n_words: int = self.model.n_vocab |
| 83 | # BOS / EOS token IDs |
| 84 | self.bos_id: int = self.special_tokens["<|begin_of_text|>"] |
| 85 | self.eos_id: int = self.special_tokens["<|end_of_text|>"] |
| 86 | self.pad_id: int = self.n_words - 1 |
| 87 | self.stop_tokens = { |
| 88 | self.special_tokens["<|end_of_text|>"], |
| 89 | self.special_tokens["<|eot_id|>"], |
| 90 | } |
| 91 | logger.info( |