| |
| """gary-neuron-chat: a tiny dialogue model that LEARNS FROM USE on three timescales. |
| FAST (within a conversation): a plastic hippocampal memory (softmax attention over |
| the conversation) recalls facts you stated earlier. No weights move. |
| SLOW (across sessions): `sleep` replays buffered conversations mixed with base |
| corpus (experience replay) and fine-tunes the cortex -- learning without forgetting. |
| The brain (cortex weights + replay buffer + age) persists in brain.npz and evolves. |
| Cortex = gary-4-petite fine-tuned on SODA dialogue. Plastic memory meta-trained |
| to retrieve. Pure numpy, no torch. Usage: python brain.py [chat|demo|sleep]""" |
| import os, sys, json, numpy as np |
| import gpt_numpy as G |
| from tokenizers import ByteLevelBPETokenizer |
| D=os.path.dirname(os.path.abspath(__file__)) |
| CFG=dict(E=96,H=4,L=4,BLK=128); EOT=0 |
| tok=ByteLevelBPETokenizer(f"{D}/petite_vocab.json", f"{D}/petite_merges.txt") |
| def _f(x): return float(np.ravel(x)[0]) |
|
|
| def load_brain(path=f"{D}/brain.npz"): |
| if os.path.exists(path): |
| z=np.load(path,allow_pickle=True) |
| P={k[2:]:z[k].astype(np.float32) for k in z.files if k.startswith("P/")} |
| Hm={k[3:]:z[k] for k in z.files if k.startswith("hm/")} |
| buf=list(z["buffer"]) if "buffer" in z.files else []; age=int(z["age"]) if "age" in z.files else 0 |
| else: |
| zc=np.load(f"{D}/cortex.npz",allow_pickle=True); P={k[2:]:zc[k].astype(np.float32) for k in zc.files if k.startswith("P/")} |
| zh=np.load(f"{D}/hippo.npz"); Hm={k:zh[k] for k in zh.files}; buf=[]; age=0 |
| return P,Hm,buf,age |
|
|
| def save_brain(P,Hm,buf,age,path=f"{D}/brain.npz"): |
| d={"P/"+k:v for k,v in P.items()}; d.update({"hm/"+k:np.atleast_1d(v) for k,v in Hm.items()}) |
| d["buffer"]=np.array(buf[-400:],dtype=object); d["age"]=age; np.savez(path,**d) |
|
|
| def cortex(P, ids): |
| x=np.array([ids[-CFG["BLK"]:]]); logits,cache=G.forward(P,x,CFG); return cache["xf"][0], logits[0] |
|
|
| def hippo_bias(Hm, xf, logits, win_ids): |
| Wq=Hm["Wq"].astype(np.float64); bN=_f(Hm["bN"]); rdec=_f(Hm["rdec"]); g=_f(Hm["g"]); T=xf.shape[0] |
| if T<2: return np.zeros(logits.shape[1]) |
| Q=xf.astype(np.float64)@Wq; qa=Q[-1]; sim=Q@qa |
| lp=logits-logits.max(1,keepdims=True); lp=lp-np.log(np.exp(lp).sum(1,keepdims=True)) |
| nxt=np.array(list(win_ids[1:])+[0]); nll=np.zeros(T); nll[:-1]=-lp[np.arange(T-1),nxt[:-1]] |
| t=np.arange(T); dist=(T-1)-t; valid=t<(T-1) |
| score=np.where(valid, sim+bN*nll+rdec*dist, -1e9); score-=score.max(); a=np.exp(score)*valid; a/=a.sum()+1e-12 |
| S=np.zeros(logits.shape[1]); np.add.at(S, nxt, a); return g*S |
|
|
| def gen_reply(P,Hm,ids,max_new=20,temp=0.0,use_hippo=True,seed=0): |
| rng=np.random.default_rng(seed); out=[] |
| for _ in range(max_new): |
| win=ids[-CFG["BLK"]:]; xf,logits=cortex(P,win); ll=logits[-1].astype(np.float64) |
| if use_hippo: ll=ll+hippo_bias(Hm,xf,logits,win) |
| if temp<=0: nx=int(ll.argmax()) |
| else: |
| q=np.exp(ll/temp-(ll/temp).max()); q/=q.sum(); nx=int(rng.choice(len(q),p=q)) |
| if nx==EOT: break |
| ids=ids+[nx]; out.append(nx) |
| if tok.decode(out).endswith("\n"): break |
| return tok.decode(out).strip(), ids |
|
|
| def sleep(P, buf, base_path=f"{D}/train.bin", steps=30, lr=1.2e-4, mix=0.12): |
| """Consolidate: replay buffered convos + base corpus (mostly base -> no forgetting).""" |
| if not buf: return P, None |
| base=np.memmap(base_path,dtype=np.uint16,mode="r"); T=CFG["BLK"] |
| ci=[]; [ci.extend(tok.encode(t).ids+[EOT]) for t in buf]; convo=np.array(ci,dtype=np.int64) |
| opt=G.Adam(P,lr=lr); rng=np.random.default_rng(0) |
| for _ in range(steps): |
| xs=[];ys=[] |
| for _ in range(16): |
| if rng.random()<mix and len(convo)>T+1: s=convo[(i:=rng.integers(0,len(convo)-T-1)):i+T+1] |
| else: j=rng.integers(0,len(base)-T-1); s=np.asarray(base[j:j+T+1],dtype=np.int64) |
| xs.append(s[:T]); ys.append(s[1:T+1]) |
| loss,cache=G.forward(P,np.stack(xs),CFG,np.stack(ys)); opt.step(P,G.backward(P,CFG,cache)) |
| return P, None |
|
|
| def chat(): |
| P,Hm,buf,age=load_brain(); print(f"gary-neuron-chat (age {age}). Tell me about yourself; I'll remember within our chat. '/sleep' to consolidate, '/exit' to leave.") |
| ids=tok.encode("G: hi").ids; sess=[] |
| while True: |
| try: u=input("\nyou: ").strip() |
| except (EOFError,KeyboardInterrupt): break |
| if u=="/exit": break |
| if u=="/sleep": |
| P,_=sleep(P,sess); age+=1; save_brain(P,Hm,buf+sess,age); sess=[]; print("...slept; cortex consolidated."); continue |
| ids=ids+tok.encode(f"\nU: {u}\nG:").ids; r,ids=gen_reply(P,Hm,ids); print("gary:",r); sess.append(f"U: {u}\nG: {r}") |
| save_brain(P,Hm,buf+sess,age); print("brain saved.") |
|
|
| if __name__=="__main__": |
| cmd=sys.argv[1] if len(sys.argv)>1 else "chat" |
| if cmd=="chat": chat() |
| elif cmd=="demo": |
| P,Hm,buf,age=load_brain() |
| ids=tok.encode("G: hello").ids |
| for u in ["my dog is named Buddy","how is the weather","i like to read books","what time is it"]: |
| ids=ids+tok.encode(f"\nU: {u}\nG: ok").ids |
| q="what is my dog 's name ?"; qids=ids+tok.encode(f"\nU: {q}\nG:").ids |
| on,_=gen_reply(P,Hm,qids,use_hippo=True); off,_=gen_reply(P,Hm,qids,use_hippo=False) |
| print(f"Q: {q}\n plastic-memory ON : {on!r}\n plastic-memory OFF: {off!r}") |
|
|
