Spaces:
Runtime error
Runtime error
File size: 3,146 Bytes
cc15101 | 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 86 87 88 89 90 91 92 93 94 95 96 | """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 ----
@classmethod
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) |