File size: 2,489 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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | import torch
from .model import EXLLM,EXLLMConfig
from .tokenizer import HybridTokenizer
class UTF8State:
__slots__=('need','lo','hi')
def __init__(self): self.need=0; self.lo=0x80; self.hi=0xBF
def accepts_byte(self,b):
if self.need: return self.lo<=b<=self.hi
return b in (9,10,13) or 0x20<=b<=0x7E or 0xC2<=b<=0xF4
def push(self,b):
if not self.accepts_byte(b): return False
if self.need:
self.need-=1; self.lo=0x80; self.hi=0xBF; return True
if b<=0x7F:return True
if 0xC2<=b<=0xDF:self.need=1
elif b==0xE0:self.need=2;self.lo=0xA0
elif 0xE1<=b<=0xEC or 0xEE<=b<=0xEF:self.need=2
elif b==0xED:self.need=2;self.hi=0x9F
elif b==0xF0:self.need=3;self.lo=0x90
elif 0xF1<=b<=0xF3:self.need=3
elif b==0xF4:self.need=3;self.hi=0x8F
return True
@property
def complete(self):return self.need==0
def load_model(ckpt='weights/EXLLM-ONI5M-v1.0.pt'):
tok=HybridTokenizer.load('tokenizer.json'); c=torch.load(ckpt,map_location='cpu',weights_only=False)
m=EXLLM(EXLLMConfig(**c['config']));m.load_state_dict(c['model']);m.eval();return m,tok,c
@torch.inference_mode()
def generate(m,tok,prompt,max_new=56):
ids=tok.encode_user(prompt); reserve=min(max_new,m.cfg.max_seq_len//2)
if len(ids)>m.cfg.max_seq_len-reserve:
keep=m.cfg.max_seq_len-reserve-3; body=ids[2:-1][-max(1,keep):];ids=[tok.BOS,tok.USER,*body,tok.ASSIST]
out=[];state=UTF8State();lps=[]
for _ in range(min(max_new,m.cfg.max_seq_len-len(ids))):
x=torch.tensor([ids+out],dtype=torch.long);logits=m(x)[0,-1].clone()
logits[tok.PAD]=logits[tok.BOS]=logits[tok.USER]=logits[tok.ASSIST]=-1e30
valid=torch.zeros_like(logits,dtype=torch.bool)
if state.complete:
valid[tok.EOS]=True
for t in range(256):
if state.accepts_byte(t):valid[t]=True
valid[256:256+len(tok.chars)]=True
else:
for t in range(256):
if state.accepts_byte(t):valid[t]=True
logits[~valid]=-1e30;lp=torch.log_softmax(logits,dim=-1);t=int(torch.argmax(logits));lps.append(float(lp[t]))
if t==tok.EOS:break
if t<256:state.push(t)
elif not state.complete:continue
out.append(t)
while out:
try:text=tok.decode(out);break
except UnicodeDecodeError:out.pop()
else:text=''
return text.strip(),sum(lps)/max(1,len(lps))
|