"""Experimental learned graph-memory augmentation, not a proven Engram successor. A frozen lexical router selects shared token nodes and adjacent nodes; trainable node embeddings, graph message passing, attention and residual gates fuse memory inside the backbone. No semantic compression or guaranteed recall is claimed. """ import hashlib import math import types from collections import defaultdict import torch from torch import nn import torch.nn.functional as F class TokenGraph: def __init__(self,documents,chunk_size=64,max_nodes=4096,candidates=12): if chunk_size<3 or max_nodes<1 or candidates<1: raise ValueError('Invalid graph settings') self.nodes=[[]] # ID 0 is a null memory node. self.adjacency=defaultdict(set) self.postings=defaultdict(list) self.candidates=candidates self.source_tokens=0 hashes={} for document in documents: previous=None for start in range(0,len(document),chunk_size): tokens=tuple(int(x) for x in document[start:start+chunk_size]) if not tokens:continue raw=b''.join(x.to_bytes(4,'little') for x in tokens) sha=hashlib.sha256(raw).digest() node=hashes.get(sha) if node is None: if len(self.nodes)-1>=max_nodes: previous=None # Do not create false adjacency across skipped chunks. continue node=len(self.nodes);hashes[sha]=node;self.nodes.append(list(tokens)) for size in (2,3): for i in range(len(tokens)-size+1): key=tokens[i:i+size] if node not in self.postings[key] and len(self.postings[key])<4: self.postings[key].append(node) elif tuple(self.nodes[node])!=tokens: raise ValueError('Hash collision') if previous is not None and previous!=node: self.adjacency[previous].add(node) self.adjacency[node].add(previous) previous=node self.source_tokens+=len(tokens) @classmethod def from_export(cls,state): graph=cls([],candidates=state['candidates']) graph.nodes=[list(node) for node in state['nodes']] graph.adjacency=defaultdict(set,{int(k):set(v) for k,v in state['adjacency'].items()}) graph.source_tokens=state['source_tokens'] for node,tokens in enumerate(graph.nodes[1:],1): for size in (2,3): for i in range(len(tokens)-size+1): key=tuple(tokens[i:i+size]) if node not in graph.postings[key] and len(graph.postings[key])<4: graph.postings[key].append(node) return graph def extend(self,documents,chunk_size=64,max_nodes=65536): """Append nodes without changing any existing node IDs.""" if chunk_size<3 or max_nodes=max_nodes: previous=None;continue node=len(self.nodes);hashes[sha]=node;self.nodes.append(tokens) for size in (2,3): for i in range(len(tokens)-size+1): key=tuple(tokens[i:i+size]) if node not in self.postings[key] and len(self.postings[key])<4: self.postings[key].append(node) elif self.nodes[node]!=tokens:raise ValueError('Hash collision') if previous is not None and previous!=node: self.adjacency[previous].add(node);self.adjacency[node].add(previous) previous=node;self.source_tokens+=len(tokens) return len(self.nodes)-1-before def route(self,input_ids): # This CPU routing path is a correctness prototype, not a fast GPU kernel. rows=input_ids.detach().cpu().tolist() output=[] for row in rows: sequence=[] for position in range(len(row)): seeds=[] for size in (3,2): if position+1>=size: for node in self.postings.get(tuple(row[position-size+1:position+1]),[]): if node not in seeds:seeds.append(node) candidates=list(seeds) for node in seeds: for neighbor in sorted(self.adjacency[node])[:2]: if neighbor not in candidates:candidates.append(neighbor) candidates=candidates[:self.candidates] sequence.append(candidates+[0]*(self.candidates-len(candidates))) output.append(sequence) return torch.tensor(output,dtype=torch.long,device=input_ids.device) def export(self): return dict(nodes=self.nodes,adjacency={str(k):sorted(v) for k,v in self.adjacency.items()}, candidates=self.candidates,source_tokens=self.source_tokens, router='exact token bigram/trigram + one-hop adjacency; non-differentiable') class GraphMemoryAttention(nn.Module): def __init__(self,hidden_size,graph,width=64,max_neighbors=8): super().__init__() self.node_embeddings=nn.Embedding(len(graph.nodes),width,padding_idx=0) neighbors=torch.zeros(len(graph.nodes),max_neighbors,dtype=torch.long) for node in range(1,len(graph.nodes)): adjacent=sorted(graph.adjacency[node])[:max_neighbors] neighbors[node,:len(adjacent)]=torch.tensor(adjacent,dtype=torch.long) self.register_buffer('neighbors',neighbors) self.edge_query=nn.Linear(width,width,bias=False) self.edge_key=nn.Linear(width,width,bias=False) self.edge_value=nn.Linear(width,width,bias=False) self.graph_norm=nn.LayerNorm(width) self.query=nn.Linear(hidden_size,width,bias=False) self.key=nn.Linear(width,width,bias=False) self.value=nn.Linear(width,width,bias=False) self.output=nn.Linear(width,hidden_size,bias=False) self.scale=nn.Parameter(torch.zeros(())) self.width=width def forward(self,hidden,node_ids): valid=node_ids.ne(0) nodes=self.node_embeddings(node_ids) # A genuine one-hop GAT update along the stored graph adjacency. # Degree is capped deterministically; topology itself is not learned. adjacent_ids=self.neighbors[node_ids] adjacent_valid=adjacent_ids.ne(0) adjacent=self.node_embeddings(adjacent_ids) edge_scores=(self.edge_query(nodes).unsqueeze(-2)*self.edge_key(adjacent)).sum(-1)/math.sqrt(self.width) edge_scores=edge_scores.masked_fill(~adjacent_valid,-1e4) edge_weights=edge_scores.softmax(-1)*adjacent_valid edge_weights=edge_weights/edge_weights.sum(-1,keepdim=True).clamp_min(1e-8) neighbor_message=(self.edge_value(adjacent)*edge_weights.unsqueeze(-1)).sum(-2) messages=self.graph_norm(nodes+neighbor_message) query=self.query(hidden.float()) scores=(self.key(messages)*query.unsqueeze(-2)).sum(-1)/math.sqrt(self.width) scores=scores.masked_fill(~valid,-1e4) weights=scores.softmax(-1)*valid weights=weights/weights.sum(-1,keepdim=True).clamp_min(1e-8) read=(self.value(messages)*weights.unsqueeze(-1)).sum(-2) delta=self.output(read)*self.scale.tanh() return hidden+delta.to(hidden.dtype) class GraphAugmentedLooped(nn.Module): """Learned graph memory inserted before selected outer Looped groups. Cache is deliberately unsupported until graph-route history is implemented. """ def __init__(self,backbone,graph,layers=(0,3),width=64): super().__init__() self.backbone=backbone self.graph=graph self.memories=nn.ModuleDict({str(i):GraphMemoryAttention(backbone.cfg.hidden_size,graph,width) for i in layers}) self.enabled=True self._routes=None self._active=False self.original_run_group=backbone.run_group parent=self def run_group(_base,h,t,cos,sin,cache=None,attn_mask=None,gates=None,allow_halt=False): if parent.enabled and str(t) in parent.memories: if parent._routes is None:raise RuntimeError('Use GraphAugmentedLooped.forward') h=parent.memories[str(t)](h,parent._routes) return parent.original_run_group(h,t,cos,sin,cache,attn_mask,gates,allow_halt) backbone.run_group=types.MethodType(run_group,backbone) def forward(self,input_ids,use_cache=False,**kwargs): if use_cache or kwargs.get('past_key_values') is not None: raise NotImplementedError('Graph route history is not implemented for cached decoding') if self._active:raise RuntimeError('Concurrent forwards are unsupported') self._active=True try: self._routes=self.graph.route(input_ids) if self.enabled else None return self.backbone(input_ids=input_ids,use_cache=False,**kwargs) finally: self._routes=None self._active=False