| 40 | |
| 41 | |
| 42 | class SimpleTokenizer: |
| 43 | def __init__(self, vocab_path, merges_path): |
| 44 | with open(vocab_path, "r", encoding="utf-8") as f: |
| 45 | self.encoder = json.load(f) |
| 46 | self.decoder = {v: k for k, v in self.encoder.items()} |
| 47 | |
| 48 | with open(merges_path, "r", encoding="utf-8") as f: |
| 49 | lines = f.read().strip().split("\n") |
| 50 | # Skip header if present |
| 51 | if lines and lines[0].startswith("#"): |
| 52 | lines = lines[1:] |
| 53 | merges = [tuple(line.split()) for line in lines if line.strip()] |
| 54 | self.bpe_ranks = dict(zip(merges, range(len(merges)))) |
| 55 | |
| 56 | self.byte_encoder = bytes_to_unicode() |
| 57 | self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} |
| 58 | self.cache = {} |
| 59 | |
| 60 | # Use 'regex' module if available (supports \p{L}), else fall back to ASCII |
| 61 | try: |
| 62 | import regex |
| 63 | self.pat = regex.compile( |
| 64 | r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", |
| 65 | regex.IGNORECASE, |
| 66 | ) |
| 67 | except ImportError: |
| 68 | # ASCII fallback — sufficient for English test strings |
| 69 | self.pat = re.compile( |
| 70 | r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[a-zA-Z]+|[0-9]|[^\s a-zA-Z0-9]+""", |
| 71 | re.IGNORECASE, |
| 72 | ) |
| 73 | |
| 74 | self.sot_token = self.encoder.get("<|startoftext|>", 49406) |
| 75 | self.eot_token = self.encoder.get("<|endoftext|>", 49407) |
| 76 | |
| 77 | def bpe(self, token): |
| 78 | if token in self.cache: |
| 79 | return self.cache[token] |
| 80 | |
| 81 | word = tuple(token[:-1]) + (token[-1] + "</w>",) |
| 82 | pairs = get_pairs(word) |
| 83 | |
| 84 | if not pairs: |
| 85 | return token + "</w>" |
| 86 | |
| 87 | while True: |
| 88 | bigram = min(pairs, key=lambda p: self.bpe_ranks.get(p, float("inf"))) |
| 89 | if bigram not in self.bpe_ranks: |
| 90 | break |
| 91 | |
| 92 | first, second = bigram |
| 93 | new_word = [] |
| 94 | i = 0 |
| 95 | while i < len(word): |
| 96 | try: |
| 97 | j = word.index(first, i) |
| 98 | new_word.extend(word[i:j]) |
| 99 | i = j |