File size: 1,878 Bytes
8d47318 | 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 | import json, unicodedata
from pathlib import Path
BYTE_BASE=0; BYTE_COUNT=256
def normalize_text(s:str)->str:
s=unicodedata.normalize('NFC',s)
return ''.join((' ' if (ord(c)<32 and c not in '\n\t') else c) for c in s).strip()
class HybridTokenizer:
def __init__(self, chars):
self.chars=list(chars); self.char_to_id={c:256+i for i,c in enumerate(self.chars)}
s=256+len(self.chars); self.PAD=s; self.BOS=s+1; self.USER=s+2; self.ASSIST=s+3; self.EOS=s+4; self.vocab_size=s+5
@classmethod
def load(cls,path): return cls(json.loads(Path(path).read_text(encoding='utf-8'))['chars'])
def encode_text(self,s):
out=[]
for c in normalize_text(s):
if c in self.char_to_id: out.append(self.char_to_id[c])
else: out.extend(c.encode('utf-8','strict'))
return out
def encode_user(self,s): return [self.BOS,self.USER,*self.encode_text(s),self.ASSIST]
def encode_example(self,prompt,answer,max_seq_len=128):
p=self.encode_text(prompt); a=self.encode_text(answer)
seq=[self.BOS,self.USER,*p,self.ASSIST,*a,self.EOS]
if len(seq)>max_seq_len:
excess=len(seq)-max_seq_len; p=p[min(excess,len(p)):]
seq=[self.BOS,self.USER,*p,self.ASSIST,*a,self.EOS]
if len(seq)>max_seq_len:
a=a[:max(1,max_seq_len-(4+len(p)))]
seq=[self.BOS,self.USER,*p,self.ASSIST,*a,self.EOS]
return seq
def token_bytes(self,tok):
if 0<=tok<256: return bytes([tok])
i=tok-256
if 0<=i<len(self.chars): return self.chars[i].encode('utf-8')
return b''
def decode(self,toks):
b=bytearray()
for t in toks:
if t in (self.PAD,self.BOS,self.USER,self.ASSIST,self.EOS): continue
b.extend(self.token_bytes(t))
return bytes(b).decode('utf-8','strict')
|