Upload folder using huggingface_hub

README.md

@@ -0,0 +1,3 @@
+# FlyBrain LM — Drosophila whole-brain language model
+
+Recurrent core = real adult fruit-fly connectome (FlyWire v783, top-48k neurons, 6.65M synapses). See `flybrain.py` for standalone inference.

config.json

@@ -0,0 +1 @@
+{"architecture": "FlyBrainLM3", "vocab": 4096, "n_fly": 48000, "e_in": 160, "leak": 0.3, "connectome": "FlyWire v783 top-48k by out-degree", "tokenizer": "bpe_4k"}

flybrain.py

@@ -0,0 +1,73 @@
+# 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:"))

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:324e8ce9c796e8f4f33246c35c7ad14f345b409f53dd0bf2ed491cbb6b3c8c45
+size 243109934