| from pathlib import Path |
| import regex as re |
| import json |
| from collections import defaultdict |
| from collections.abc import Iterable, Iterator |
|
|
|
|
| def gpt2_bytes_to_unicode() -> dict[int, str]: |
| bs = (list(range(ord("!"), ord("~") + 1)) |
| + list(range(ord("¡"), ord("¬") + 1)) |
| + list(range(ord("®"), ord("ÿ") + 1))) |
| cs = bs[:] |
| n = 0 |
| for i in range(256): |
| if i not in bs: |
| bs.append(i) |
| cs.append(256+n) |
| n += 1 |
|
|
| return {b:chr(c) for b,c in zip(bs,cs)} |
| |
| def gpt2_unicode_to_bytes() -> dict[str,int]: |
| return {v:k for k,v in gpt2_bytes_to_unicode().items()} |
|
|
|
|
| PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""" |
|
|
| class Tokenizer: |
| def __init__(self, vocab: dict[int, bytes], merges: list[tuple[bytes,bytes]], special_tokens: list[str]| None = None): |
| self.vocab = vocab |
| self.merges = merges |
| self.special_tokens = special_tokens or [] |
| self.byte_to_id = {v: k for k, v in self.vocab.items()} |
| self.merge_rank = {pair: i for i, pair in enumerate(merges)} |
|
|
|
|
| |
| |
| @classmethod |
| def from_file(cls, vocab_filepath: str | Path, merges_filepath: str| Path, special_tokens: list[str] | None = None) -> "Tokenizer": |
| |
| byte_decoder = gpt2_unicode_to_bytes() |
| to_bytes = lambda s: bytes(byte_decoder[ch] for ch in s) |
|
|
| with open(vocab_filepath, encoding = "utf-8") as f: |
| raw = json.load(f) |
| vocab = {k:to_bytes(v) for v,k in raw.items()} |
|
|
| |
| merges = [] |
| with open(merges_filepath, encoding = "utf-8") as f: |
| for line in f: |
| line = line.rstrip("\n") |
| a , b = line.split(" ") |
| merges.append((to_bytes(a), to_bytes(b))) |
|
|
| return cls(vocab, merges, special_tokens) |
| |
|
|
| def encode(self, string: str) -> list[int]: |
| |
| if self.special_tokens: |
| specials = sorted(self.special_tokens, key = len, reverse=True) |
| segment_pattern = re.compile("(" + "|".join(re.escape(t) for t in specials) + ")") |
| segments = re.split(segment_pattern, string) |
| else: |
| segments = [string] |
|
|
| result = [] |
| chunk_pattern = re.compile(PAT) |
|
|
| for segment in segments: |
| if segment in self.special_tokens: |
| result.append(self.byte_to_id[segment.encode("utf-8")]) |
|
|
| else: |
| |
| for match in chunk_pattern.finditer(segment): |
| token = match.group() |
| byte_list = [bytes([x]) for x in token.encode("utf-8")] |
| new_list = byte_list.copy() |
|
|
| |
| while True: |
| |
| if len(new_list) < 2: |
| break |
|
|
| |
| pairs = [(i1,i2) for i1,i2 in zip(new_list, new_list[1:])] |
|
|
| |
| rank = {p: self.merge_rank[p] for p in pairs if p in self.merge_rank} |
| if not rank: |
| break |
| min_pair = min(rank, key=rank.get) |
|
|
| |
| merge_applied_list = [] |
| i = 0 |
| while i < len(new_list): |
| if i < len(new_list) - 1 and new_list[i] == min_pair[0] and new_list[i+1] == min_pair[1]: |
| merge_applied_list.append(min_pair[0] + min_pair[1]) |
| i += 2 |
| else: |
| merge_applied_list.append(new_list[i]) |
| i += 1 |
|
|
| new_list = merge_applied_list |
|
|
| for tok in new_list: |
| result.append(self.byte_to_id[tok]) |
|
|
|
|
| return result |
|
|
| |
| def encode_iterable(self, iterable: Iterable[str]) -> Iterator[int]: |
| for chunk in iterable: |
| yield from self.encode(chunk) |
|
|
|
|
|
|
| def decode(self, byte_array: list[int]) -> str: |
| decoded: bytes = b"" |
| for i in byte_array: |
| decoded += self.vocab[i] |
| return decoded.decode("utf-8", errors="replace") |
|
|
|
|
|
|
| if __name__ == "__main__": |
| |
| |
| |
| ... |
