File size: 2,024 Bytes
da49047 | 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 | """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")
|