Spaces:
Runtime error
Runtime error
| """AbyssTokenizer: byte-level BPE encode/decode. | |
| Source of truth on disk is merges.txt (integer id pairs, in priority | |
| order). The vocab (id -> bytes) is rebuilt by replaying those merges | |
| over the 256 base bytes.""" | |
| import os | |
| import json | |
| import regex as re | |
| from .regex_split import split | |
| from .special_tokens import SPECIAL_TOKENS, special_token_ids | |
| def merge_seq(ids, pair, new_id): | |
| """Replace every adjacent occurence of 'pair' in 'ids' with new_id.""" | |
| out = [] | |
| i = 0 | |
| while i < len(ids): | |
| if i < len(ids) - 1 and ids[i] == pair[0] and ids[i + 1] == pair[1]: | |
| out.append(new_id) | |
| i += 2 | |
| else: | |
| out.append(ids[i]) | |
| i += 1 | |
| return out | |
| class AbyssTokenizer: | |
| def __init__(self, merges, vocab_size): | |
| # merges: ordered list of (a, b) integer id pairs | |
| self.vocab_size = vocab_size | |
| self.merges = merges | |
| # rank[(a, b)] = priority index (lower = merged earlier) | |
| self.ranks = {pair: i for i, pair in enumerate(merges)} | |
| # rebuild id -> bytes | |
| self.vocab = {i: bytes([i]) for i in range(256)} | |
| for i, (a, b) in enumerate(merges): | |
| self.vocab[256 + i] = self.vocab[a] + self.vocab[b] | |
| # special tokens | |
| self.special = special_token_ids(vocab_size) | |
| self.inv_special = {v: k for k, v in self.special.items()} | |
| self._special_pat = re.compile( | |
| "(" + "|".join(re.escape(t) for t in SPECIAL_TOKENS) + ")" | |
| ) | |
| # ---- Encode ---- | |
| def _encode_chunk(self, ids): | |
| while len(ids) >= 2: | |
| pairs = set(zip(ids, ids[1:])) | |
| best = min(pairs, key=lambda p: self.ranks.get(p, float("inf"))) | |
| if best not in self.ranks: | |
| break | |
| new_id = 256 + self.ranks[best] | |
| ids = merge_seq(ids, best, new_id) | |
| return ids | |
| def encode(self, text): | |
| out = [] | |
| for piece in self._special_pat.split(text): | |
| if piece == "": | |
| continue | |
| if piece in self.special: | |
| out.append(self.special[piece]) | |
| else: | |
| for chunk in split(piece): | |
| ids = list(chunk.encode("utf-8")) | |
| out.extend(self._encode_chunk(ids)) | |
| return out | |
| # ---- Decode ---- | |
| def decode(self, ids): | |
| parts = [] | |
| for i in ids: | |
| if i in self.inv_special: | |
| parts.append(self.inv_special[i].encode("utf-8")) | |
| else: | |
| parts.append(self.vocab[i]) | |
| return b"".join(parts).decode("utf-8", errors="replace") | |
| # ---- load ---- | |
| def load(cls, in_dir): | |
| with open(os.path.join(in_dir, "config.json")) as f: | |
| vocab_size = json.load(f)["vocab_size"] | |
| merges = [] | |
| with open(os.path.join(in_dir, "merges.txt")) as f: | |
| for line in f: | |
| a, b = line.split() | |
| merges.append((int(a), int(b))) | |
| return cls(merges, vocab_size) |