| 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)) |
|
|
