neuron/tools/retrieval-eval/legs.py

import json,pickle,os,math,urllib.request,numpy as np
S='/private/tmp/claude-501/-Users-timlingo/82369039-a20e-4b5a-8a5e-28234a57b996/scratchpad/sim/'
C=pickle.load(open(S+'corpus.pkl','rb'))
NODES=C['nodes']; EDGES=C['edges']; N=len(NODES)
E=np.load(S+'emb.npy'); HAVE=np.load(S+'have.npy')
En=E/np.maximum(np.linalg.norm(E,axis=1,keepdims=True),1e-12)
LAYERS={int(l['layer_id']):l for l in (C['layers'] or [])} if C['layers'] else {}
TRANS=set(i for i,l in LAYERS.items() if l.get('transparent'))
def addressable(s):
    if not s: return False
    return all(0x20<=ord(ch)<=0x7e for ch in s)
ADDR=np.array([addressable(n['id']) for n in NODES])
OK=np.array([ (n['layer_id'] not in TRANS) and ADDR[i] for i,n in enumerate(NODES)])
SAL=np.array([n['salience'] for n in NODES])
LOW=[ (n['content']+'\x00'+n['label']+'\x00'+n['tags']).lower() for n in NODES]
DL=np.array([float(len(n['content'])+len(n['label'])+len(n['tags'])) for n in NODES])
IDX={}
for i,n in enumerate(NODES):
    IDX.setdefault(n['id'],i)
STRUCT={"identity","contains","superseded_by","references","embodies","demonstrated_by","canonical-self","depends_on","currently_holds","activates"}
ADJ_F=[[] for _ in range(N)]; ADJ_T=[[] for _ in range(N)]
for e in EDGES:
    a=IDX.get(e['from']); b=IDX.get(e['to'])
    if a is None or b is None: continue
    ADJ_F[a].append((e,b)); ADJ_T[b].append((e,a))
EXCL=np.array([n['node_type'] in ('Tag','InternalStateEvent') for n in NODES])
avgdl_all=None
def tokenize(q):
    out=[]
    for t in q.split():
        if not any(t.lower()==x.lower() for x in out): out.append(t)
    return out
_qcache={}
def qemb(q):
    if q in _qcache: return _qcache[q]
    body=json.dumps({"model":"nomic-embed-text","prompt":q}).encode()
    r=urllib.request.urlopen(urllib.request.Request("http://127.0.0.1:11434/api/embeddings",data=body,headers={"Content-Type":"application/json"}),timeout=30)
    v=np.array(json.loads(r.read())["embedding"],dtype=np.float32)
    v=v/np.linalg.norm(v); _qcache[q]=v; return v
K1,B=1.2,0.75
SEED_MIN=0.60; SEED_K=8; ASSOC_SEEDS=3; DEPTH=2; FIRE=0.02; AMAX=64; DECAY=0.7
def legs(query):
    toks=tokenize(query)
    masks=[]; 
    hit_idx=[]; hit_mask=[]
    df=[0]*len(toks)
    lt=[t.lower() for t in toks]
    for i in range(N):
        if not OK[i]: continue
        s=LOW[i]; m=0
        for t,tok in enumerate(lt):
            if tok in s: m|=(1<<t)
        if m:
            hit_idx.append(i); hit_mask.append(m)
            for t in range(len(toks)):
                if m>>t&1: df[t]+=1
    dl_n=int(OK.sum()); avgdl=float(DL[OK].sum()/max(dl_n,1))
    idf=[math.log(1.0+((dl_n-d+0.5)/(d+0.5))) for d in df]
    L=[]
    for j,i in enumerate(hit_idx):
        norm=1.0-B+B*(DL[i]/avgdl); w=0.0
        for t in range(len(toks)):
            if hit_mask[j]>>t&1: w+=idf[t]*(K1+1.0)/(1.0+K1*norm)
        L.append((i,w,SAL[i]))
    L.sort(key=lambda x:(-x[1],-x[2]))
    qv=qemb(query)
    cos=En@qv
    cos=np.where(HAVE&OK,cos,-2.0)
    order=np.argsort(-cos)
    semfull=[(int(i),float(cos[i])) for i in order[:400]]
    Sleg=[(i,(c-SEED_MIN)/(1-SEED_MIN)) for i,c in semfull if c>SEED_MIN]
    semseed=[i for i,c in semfull[:SEED_K] if c>0.0]
    # assoc
    act={}; seen={}; qq=[]
    for i,_,_ in L[:ASSOC_SEEDS]:
        act[i]=1.0; seen[i]=2; qq.append((i,0))
    for i in semseed:
        if i in seen: continue
        act[i]=1.0; seen[i]=2; qq.append((i,0))
    qh=0
    while qh<len(qq):
        cur,h=qq[qh]; qh+=1
        if h>=DEPTH: continue
        parent=act[cur]
        for e,oi in ADJ_F[cur]+ADJ_T[cur]:
            if e['rel'] not in STRUCT: continue
            if EXCL[oi]: continue
            na=parent*e['w']*DECAY*SAL[oi]
            if na<FIRE: continue
            if seen.get(oi) and na<=act.get(oi,0): continue
            act[oi]=na
            if not seen.get(oi): seen[oi]=1
            if len(qq)<AMAX*4: qq.append((oi,h+1))
    A=[]
    for i,st in seen.items():
        if st!=1: continue
        if not OK[i] or not HAVE[i]: continue
        c=float(cos[i])
        if c<=0.0: continue
        A.append((i,c))
    A.sort(key=lambda x:-x[1]); A=A[:AMAX]
    return L,Sleg,A,cos
def interleave3(L,Sl,A,lim=10):
    out=[]; li=si=ai=0
    while len(out)<lim and (li<len(L) or si<len(Sl) or ai<len(A)):
        if li<len(L):
            if L[li][0] not in out: out.append(L[li][0])
            li+=1
        if len(out)>=lim: break
        if si<len(Sl):
            if Sl[si][0] not in out: out.append(Sl[si][0])
            si+=1
        if len(out)>=lim: break
        if ai<len(A):
            if A[ai][0] not in out: out.append(A[ai][0])
            ai+=1
    return out