File size: 3,626 Bytes
76d09fe 02afb58 18fe584 02afb58 76d09fe 02afb58 76d09fe 02afb58 18fe584 | 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 97 98 99 100 101 102 103 104 105 | """The tokenizer of bard-small."""
import os
import re
import numpy as np
import tokenmonster
VOCAB_FILE = "bard-16k.vocab"
FORMAT_FILE = "format.txt"
EOS = "<|end|>"
NUM_MARKERS = ["<|num|>", "<|snum|>"]
MAX_DIGITS = 8
LEN_SHIFT = 27
VALUE_MASK = (1 << LEN_SHIFT) - 1
NUMBER_RE = re.compile(r"( ?)([0-9]+)")
def pack(digits: str) -> int:
assert 0 < len(digits) <= MAX_DIGITS
return int(digits) | (len(digits) << LEN_SHIFT)
def unpack(packed: int) -> str:
return str(packed & VALUE_MASK).zfill(packed >> LEN_SHIFT)
def split_numbers(text):
out, pos = [], 0
for m in NUMBER_RE.finditer(text):
if m.start() > pos:
out.append(text[pos:m.start()])
spaced, digits = bool(m.group(1)), m.group(2)
for i in range(0, len(digits), MAX_DIGITS):
out.append((spaced and i == 0, pack(digits[i:i + MAX_DIGITS])))
pos = m.end()
if pos < len(text):
out.append(text[pos:])
return out
def parse_special_tokens(format_text):
markers = re.findall(r"<\|[a-z_]+\|>", format_text)
roles = re.findall(r"<\|start\|>([a-z_]+)<\|content\|>", format_text)
return list(dict.fromkeys(markers + [f"<|start|>{r}" for r in roles]))
class BardTokenizer:
def __init__(self, vocab_path, format_path):
self.vocab = tokenmonster.load(vocab_path)
self.specials = {t: self.vocab.token_to_id(t)
for t in parse_special_tokens(open(format_path, encoding="utf-8").read()) + NUM_MARKERS}
missing = [t for t, i in self.specials.items() if i is None]
assert not missing, f"the vocabulary lacks {missing}"
self.eos_id = self.specials[EOS]
self.num_id, self.snum_id = (self.specials[t] for t in NUM_MARKERS)
self._special_text = {i: t for t, i in self.specials.items()}
@property
def vocab_size(self):
return max(self.specials.values()) + 1
def encode(self, text):
plan = split_numbers(text)
segs = [item for item in plan if isinstance(item, str)]
seg_ids = iter(self.vocab.tokenize(segs) if segs else [])
ids, vals = [], []
for item in plan:
if isinstance(item, str):
t = next(seg_ids)
ids.append(t.astype(np.int64, copy=False))
vals.append(np.zeros(len(t), np.int64))
else:
spaced, packed = item
ids.append(np.array([self.snum_id if spaced else self.num_id], np.int64))
vals.append(np.array([packed], np.int64))
if not ids:
return np.zeros(0, np.int64), np.zeros(0, np.int64)
return np.concatenate(ids), np.concatenate(vals)
def decode(self, ids, vals=None):
ids = np.asarray(ids).tolist()
vals = np.asarray(vals).tolist() if vals is not None else [0] * len(ids)
parts, run = [], []
for i, v in zip(ids, vals):
if i in (self.num_id, self.snum_id):
if run:
parts.append(self.vocab.decode(run))
run = []
parts.append((" " if i == self.snum_id else "") + unpack(v))
elif i in self._special_text:
if run:
parts.append(self.vocab.decode(run))
run = []
parts.append(self._special_text[i])
else:
run.append(i)
if run:
parts.append(self.vocab.decode(run))
return "".join(parts)
@classmethod
def from_dir(cls, path):
return cls(os.path.join(path, VOCAB_FILE), os.path.join(path, FORMAT_FILE))
|