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