graph_model.py

"""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<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