File size: 3,083 Bytes
53bd701 | 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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | # standalone inference for the public Hub repo (torch + safetensors + tokenizers only)
import json, torch
from safetensors.torch import load_file
from tokenizers import Tokenizer
class FlyBrainLM(torch.nn.Module):
def __init__(self, vocab=4096, n_fly=48000, e_in=160, leak=0.3):
super().__init__()
self.n, self.leak = n_fly, leak
self.register_buffer("Widx", torch.zeros(2, 1, dtype=torch.long))
self.Wval = torch.nn.Parameter(torch.zeros(1))
self.Win = torch.nn.Parameter(torch.zeros(n_fly, e_in), requires_grad=False)
self.emb = torch.nn.Embedding(vocab, e_in)
self.lens1 = torch.nn.Linear(n_fly, 1024)
self.lens2 = torch.nn.Linear(1024, vocab)
def forward(self, idx):
W = torch.sparse_coo_tensor(self.Widx, self.Wval, (self.n, self.n)).coalesce()
e = self.emb(idx)
s = torch.zeros(idx.shape[0], self.n, device=idx.device)
outs = []
for t in range(idx.shape[1]):
s = (1 - self.leak) * s + self.leak * torch.tanh(
e[:, t, :] @ self.Win.T + torch.sparse.mm(W, s.T).T)
outs.append(s)
return self.lens2(torch.tanh(self.lens1(torch.stack(outs, 1))))
def load_model(repo_or_dir, device="cuda"):
try:
from huggingface_hub import snapshot_download
import os
d = repo_or_dir if os.path.isdir(repo_or_dir) else snapshot_download(repo_or_dir)
except Exception:
d = repo_or_dir
with open(f"{d}/config.json") as f:
cfg = json.load(f)
tok = Tokenizer.from_file(f"{d}/tokenizer.json")
m = FlyBrainLM(cfg["vocab"], cfg["n_fly"], cfg["e_in"], cfg.get("leak", 0.3))
m.load_state_dict(load_file(f"{d}/model.safetensors", device="cpu"))
return m.to(device).eval(), tok, cfg
def generate(model, tok, prompt, max_new=120, temp=0.8, ctx_win=160):
dev = next(model.parameters()).device
ids = tok.encode(prompt).ids[-ctx_win:]
W = torch.sparse_coo_tensor(model.Widx, model.Wval, (model.n, model.n)).coalesce()
with torch.no_grad():
e = model.emb(torch.tensor([ids], device=dev))
s = torch.zeros(1, model.n, device=dev)
for t in range(len(ids)):
s = 0.7 * s + 0.3 * torch.tanh(e[:, t, :] @ model.Win.T + torch.sparse.mm(W, s.T).T)
last, out = ids[-1], []
for _ in range(max_new):
e1 = model.emb(torch.tensor([[last]], device=dev))
s = 0.7 * s + 0.3 * torch.tanh(e1[:, 0, :] @ model.Win.T + torch.sparse.mm(W, s.T).T)
p = torch.softmax(model.lens2(torch.tanh(model.lens1(s)))[0] / temp, -1)
nx = int(torch.multinomial(p, 1))
if nx == tok.encode("[EOS]").ids[0]:
break
out.append(nx)
last = nx
txt = tok.decode(out)
cut = txt.find("\nHuman:")
return txt[:cut].strip() if cut >= 0 else txt.strip()
if __name__ == "__main__":
import sys
repo = sys.argv[1] if len(sys.argv) > 1 else "."
m, tok, cfg = load_model(repo)
print(generate(m, tok, "Human: Hi Fly! Who are you?\nFly:"))
|