summer
Publish experimental merged bilingual LoopedQwen + GraphEngram snapshot with evaluation and reload validation
5c270fa verified | """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) | |
| 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<len(self.nodes)-1:raise ValueError('Invalid extension bounds') | |
| def digest(tokens): | |
| return hashlib.sha256(b''.join(int(x).to_bytes(4,'little') for x in tokens)).digest() | |
| hashes={digest(tokens):node for node,tokens in enumerate(self.nodes[1:],1)} | |
| before=len(self.nodes)-1 | |
| for document in documents: | |
| previous=None | |
| for start in range(0,len(document),chunk_size): | |
| tokens=list(map(int,document[start:start+chunk_size])) | |
| if not tokens:continue | |
| sha=digest(tokens);node=hashes.get(sha) | |
| if node is None: | |
| if len(self.nodes)-1>=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 | |