| """Small text-generation helper for an exported SparkBET repository.""" |
| from pathlib import Path |
| import torch |
| from safetensors.torch import load_file |
| from bet_model import SparkBET,BETConfig,uniform_steps |
|
|
|
|
| class Cortex: |
| def __init__(self,model,device=None): |
| self.model=model |
| self.device=torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu")) |
| self.model.to(self.device).eval() |
|
|
| @classmethod |
| def from_export(cls,folder,device=None): |
| folder=Path(folder);model=SparkBET(BETConfig()) |
| state=load_file(str(folder/"model.safetensors"),device="cpu") |
| if state and all(k.startswith("core.") for k in state):state={k[5:]:v for k,v in state.items()} |
| model.load_state_dict(state,strict=True) |
| return cls(model,device) |
|
|
| def generate_ids(self,ids,max_new_tokens=128,loops=8,temperature=0.0,top_k=None): |
| out=list(map(int,ids)) |
| for _ in range(int(max_new_tokens)): |
| current=out[-self.model.c.max_seq_len:] |
| x=torch.tensor([current],device=self.device,dtype=torch.long) |
| with torch.inference_mode(),torch.autocast(self.device.type,dtype=torch.float16,enabled=self.device.type=="cuda"): |
| logits=self.model(x,uniform_steps(loops))[0,-1].float() |
| if temperature and temperature>0: |
| logits=logits/float(temperature) |
| if top_k: |
| values,_=torch.topk(logits,min(int(top_k),logits.numel()));logits[logits<values[-1]]=-float("inf") |
| nxt=int(torch.multinomial(torch.softmax(logits,-1),1)) |
| else:nxt=int(logits.argmax()) |
| out.append(nxt) |
| if nxt==258:break |
| return out |
|
|
| def generate(self,text,max_new_tokens=128,loops=8,temperature=0.0,top_k=None): |
| ids=[257]+list(text.encode("utf-8")) |
| out=self.generate_ids(ids,max_new_tokens,loops,temperature,top_k) |
| body=bytes(i for i in out[1:] if 0<=i<=255) |
| return body.decode("utf-8",errors="replace") |
|
|
