File size: 3,824 Bytes
8efdc48 dfd47b7 8efdc48 dfd47b7 8efdc48 | 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 | """Small, auditable byte-pair encoder trained only on supplied local text."""
import argparse, collections, json
from pathlib import Path
SPECIAL = ["<pad>", "<bos>", "<eos>", "<system>", "<user>", "<assistant>", "<tool>", "<plan>", "<memory>"]
class ByteBPETokenizer:
def __init__(self, merges=None, special_tokens=SPECIAL):
self.special_tokens = special_tokens
self.special_to_id = {s:i for i,s in enumerate(special_tokens)}
self.merges = [tuple(x) for x in (merges or [])] # child token IDs, in priority order
self.base = len(special_tokens)
self.pair_to_token = {p:self.base+256+i for i,p in enumerate(self.merges)}
self.pair_rank = {p:i for i,p in enumerate(self.merges)}
self.vocab_size = self.base + 256 + len(self.merges)
def encode_bytes(self, raw):
ids = [self.base+b for b in raw]
# Standard BPE: always apply the currently lowest-ranked available pair.
while len(ids) > 1:
choices=[(i, self.pair_rank[(ids[i],ids[i+1])]) for i in range(len(ids)-1) if (ids[i],ids[i+1]) in self.pair_rank]
if not choices: break
i,rank=min(choices,key=lambda x:x[1]); ids[i:i+2]=[self.base+256+rank]
return ids
def encode(self, text, add_bos=False, add_eos=False):
ids=[]; pos=0
while pos < len(text):
found=next(((s,i) for s,i in self.special_to_id.items() if text.startswith(s,pos)),None)
if found: ids.append(found[1]);pos+=len(found[0]);continue
ends=[text.find(s,pos) for s in self.special_to_id if text.find(s,pos)>=0]; end=min(ends) if ends else len(text)
ids.extend(self.encode_bytes(text[pos:end].encode('utf8')));pos=end
return ([self.special_to_id['<bos>']] if add_bos else [])+ids+([self.special_to_id['<eos>']] if add_eos else [])
def decode(self, ids):
def expand(x):
if x < self.base: return b''
if x < self.base+256: return bytes([x-self.base])
a,b=self.merges[x-self.base-256]; return expand(a)+expand(b)
out=[];data=bytearray()
for x in ids:
if x < self.base:
if data: out.append(data.decode('utf8',errors='replace'));data.clear()
out.append(self.special_tokens[x])
else:data.extend(expand(x))
if data:out.append(data.decode('utf8',errors='replace'))
return ''.join(out)
def save(self,path): Path(path).parent.mkdir(parents=True,exist_ok=True);Path(path).write_text(json.dumps({'merges':self.merges,'special_tokens':self.special_tokens}))
@classmethod
def load(cls,path):
d=json.loads(Path(path).read_text());return cls(d['merges'],d['special_tokens'])
def train(files,vocab_size):
symbols=[len(SPECIAL)+b for b in b''.join(Path(p).read_bytes() for p in files)];merges=[]
while len(SPECIAL)+256+len(merges)<vocab_size:
pairs=collections.Counter(zip(symbols,symbols[1:]));
if not pairs:break
pair,n=pairs.most_common(1)[0]
if n<2:break
token=len(SPECIAL)+256+len(merges);merges.append(pair);new=[];i=0
while i<len(symbols):
if i+1<len(symbols) and (symbols[i],symbols[i+1])==pair:new.append(token);i+=2
else:new.append(symbols[i]);i+=1
symbols=new
return ByteBPETokenizer(merges)
def main():
p=argparse.ArgumentParser(); sub=p.add_subparsers(dest='cmd',required=True); t=sub.add_parser('train'); t.add_argument('--input',required=True);t.add_argument('--out',required=True);t.add_argument('--vocab-size',type=int,default=16000);a=p.parse_args()
files=list(Path(a.input).rglob('*.txt')); assert files, 'No .txt files found'; tok=train(files,a.vocab_size);tok.save(a.out);print(f'Saved {tok.vocab_size} tokens to {a.out}')
if __name__=='__main__': main()
|