Upload folder using huggingface_hub

3e363e2 verified 13 days ago

9.11 kB

"""Relation-aware inductive GraphSAGE and optional GAT encoders."""

from __future__ import annotations

from dataclasses import dataclass

import torch
from torch import Tensor, nn
from torch.nn import functional

from graph.config import GraphConfig
from graph.types import GraphTensors


@dataclass(frozen=True)
class GraphEncoderOutput:
    node_states: Tensor
    relation_states: Tensor
    edge_states: Tensor
    attention_weights: Tensor | None


def _edge_query(
    relation_states: Tensor, query_state: Tensor | None, edge_query: Tensor | None
) -> Tensor:
    """The query vector each edge's message is conditioned on.
    ``edge_query`` gives one row per edge, which is how a batch of disjoint
    graphs, each with its own query, is encoded in a single pass.
    """
    if edge_query is not None:
        return edge_query
    if query_state is None:
        return torch.zeros_like(relation_states)
    return query_state.unsqueeze(0).expand_as(relation_states)


class RelationAwareSAGELayer(nn.Module):
    def __init__(
        self,
        hidden_dim: int,
        edge_dim: int,
        num_relations: int,
        dropout: float,
        residual: bool = False,
    ) -> None:
        super().__init__()
        self.residual = residual
        self.self_proj = nn.Linear(hidden_dim, hidden_dim)
        self.message_proj = nn.Linear(
            hidden_dim * 2 + hidden_dim + edge_dim, hidden_dim
        )
        self.relation_embedding = nn.Embedding(num_relations, hidden_dim)
        self.norm = nn.LayerNorm(hidden_dim)
        self.dropout = nn.Dropout(dropout)

    def _skip(self, node_states: Tensor) -> Tensor | float:
        """The identity path, when enabled; it keeps depth from oversmoothing."""
        return node_states if self.residual else 0.0

    def forward(
        self,
        node_states: Tensor,
        graph: GraphTensors,
        relation_ids: Tensor,
        query_state: Tensor | None = None,
        edge_query: Tensor | None = None,
    ) -> tuple[Tensor, Tensor, Tensor | None]:
        source, target = graph.edge_index
        source_states = node_states[source]
        target_states = node_states[target]
        relation_states = self.relation_embedding(relation_ids)
        query_for_edges = _edge_query(relation_states, query_state, edge_query)
        messages = self.message_proj(
            torch.cat(
                (source_states, target_states, relation_states, graph.edge_features),
                dim=-1,
            )
        )
        messages = messages + query_for_edges
        aggregate = torch.zeros_like(node_states)
        aggregate.index_add_(0, target, messages)
        degree = torch.zeros(
            (graph.num_nodes,), device=node_states.device, dtype=node_states.dtype
        )
        degree.index_add_(0, target, torch.ones_like(target, dtype=node_states.dtype))
        aggregate = aggregate / degree.clamp_min(1.0).unsqueeze(-1)
        updated = self.norm(
            self._skip(node_states)
            + self.self_proj(node_states)
            + self.dropout(functional.gelu(aggregate))
        )
        return updated, relation_states, None


class RelationAwareGATLayer(RelationAwareSAGELayer):
    """GAT ablation that retains per-edge attention for explanation."""

    def __init__(
        self,
        hidden_dim: int,
        edge_dim: int,
        num_relations: int,
        dropout: float,
        residual: bool = False,
    ) -> None:
        super().__init__(hidden_dim, edge_dim, num_relations, dropout, residual)
        self.attn = nn.Linear(hidden_dim * 2 + hidden_dim + edge_dim, 1)

    def forward(
        self,
        node_states: Tensor,
        graph: GraphTensors,
        relation_ids: Tensor,
        query_state: Tensor | None = None,
        edge_query: Tensor | None = None,
    ) -> tuple[Tensor, Tensor, Tensor | None]:
        source, target = graph.edge_index
        source_states = node_states[source]
        target_states = node_states[target]
        relation_states = self.relation_embedding(relation_ids)
        query_for_edges = _edge_query(relation_states, query_state, edge_query)
        edge_input = torch.cat(
            (source_states, target_states, relation_states, graph.edge_features), dim=-1
        )
        logits = self.attn(edge_input).squeeze(-1) + (
            relation_states * query_for_edges
        ).sum(-1) / (node_states.shape[-1] ** 0.5)
        weights = torch.zeros_like(logits)
        for node in range(graph.num_nodes):
            mask = target == node
            if mask.any():
                weights[mask] = torch.softmax(logits[mask], dim=0)
        messages = self.message_proj(edge_input) * weights.unsqueeze(-1)
        aggregate = torch.zeros_like(node_states)
        aggregate.index_add_(0, target, messages)
        updated = self.norm(
            self._skip(node_states)
            + self.self_proj(node_states)
            + self.dropout(functional.gelu(aggregate))
        )
        return updated, relation_states, weights


class GraphEncoder(nn.Module):
    """Inductive graph encoder; no trainable closed-world node-ID table is used."""

    def __init__(self, config: GraphConfig) -> None:
        super().__init__()
        self.config = config
        self.node_type_embedding = nn.Embedding(
            config.num_node_types, config.hidden_dim
        )
        self.node_input = nn.Linear(config.graph_input_dim, config.hidden_dim)
        layer_cls = (
            RelationAwareGATLayer
            if config.graph_encoder == "gat" or config.use_gat
            else RelationAwareSAGELayer
        )
        self.layers = nn.ModuleList(
            [
                layer_cls(
                    config.hidden_dim,
                    config.edge_input_dim,
                    config.num_relations,
                    config.dropout,
                    residual=config.graph_residual,
                )
                for _ in range(config.graph_layers)
            ]
        )
        self.edge_input = nn.Linear(config.edge_input_dim, config.hidden_dim)
        self.final_norm = nn.LayerNorm(config.hidden_dim)
        self.start_embedding = (
            nn.Embedding(2, config.hidden_dim) if config.start_marking else None
        )

    def forward(
        self,
        graph: GraphTensors,
        query_state: Tensor | None = None,
        *,
        edge_query: Tensor | None = None,
        start_mask: Tensor | None = None,
        node_query: Tensor | None = None,
    ) -> GraphEncoderOutput:
        """Encode ``graph``, conditioned on the query where it is given.
        ``edge_query``/``node_query`` give one query row per edge/node, for a
        batch of graphs with different queries; otherwise ``query_state`` is
        shared by the whole graph.  ``start_mask`` flags the query's start
        entities and is used when the encoder was built with ``start_marking``.
        """
        # External graph services are open-world: a tenant may introduce more
        # node/relation labels than the configured embedding tables, so the stable
        # global identities are reduced into the learned vocabulary here.  The
        # graph itself is never rewritten: its *local* type indices are what the
        # caller uses to pool per-edge states into a per-graph relation bank and to
        # map a bank row back to a relation name, and overwriting them with reduced
        # hashes would corrupt that bookkeeping.
        node_type_ids = graph.node_type_identity.remainder(self.config.num_node_types)
        edge_type_ids = graph.edge_type_identity.remainder(self.config.num_relations)
        node_states = self.node_input(graph.node_features) + self.node_type_embedding(
            node_type_ids
        )
        if self.start_embedding is not None and start_mask is not None:
            marked = start_mask.to(torch.long)
            node_states = node_states + self.start_embedding(marked)
            if node_query is None and query_state is not None:
                node_query = query_state.reshape(1, -1).expand_as(node_states)
            if node_query is not None:
                node_states = node_states + marked.unsqueeze(-1).to(
                    node_states.dtype
                ) * node_query.to(node_states.dtype)
        all_relation_states: list[Tensor] = []
        attention: Tensor | None = None
        for layer in self.layers:
            node_states, relation_states, layer_attention = layer(
                node_states, graph, edge_type_ids, query_state, edge_query
            )
            if layer_attention is not None:
                attention = layer_attention
            all_relation_states.append(relation_states)
        node_states = self.final_norm(node_states)
        edge_states = self.edge_input(graph.edge_features)
        relation_states = (
            all_relation_states[-1]
            if all_relation_states
            else torch.empty((0, self.config.hidden_dim), device=node_states.device)
        )
        return GraphEncoderOutput(node_states, relation_states, edge_states, attention)

	"""Relation-aware inductive GraphSAGE and optional GAT encoders."""

	from __future__ import annotations

	from dataclasses import dataclass

	import torch
	from torch import Tensor, nn
	from torch.nn import functional

	from graph.config import GraphConfig
	from graph.types import GraphTensors


	@dataclass(frozen=True)
	class GraphEncoderOutput:
	node_states: Tensor
	relation_states: Tensor
	edge_states: Tensor
	attention_weights: Tensor \| None


	def _edge_query(
	relation_states: Tensor, query_state: Tensor \| None, edge_query: Tensor \| None
	) -> Tensor:
	"""The query vector each edge's message is conditioned on.

	``edge_query`` gives one row per edge, which is how a batch of disjoint
	graphs, each with its own query, is encoded in a single pass.
	"""
	if edge_query is not None:
	return edge_query
	if query_state is None:
	return torch.zeros_like(relation_states)
	return query_state.unsqueeze(0).expand_as(relation_states)


	class RelationAwareSAGELayer(nn.Module):
	def __init__(
	self,
	hidden_dim: int,
	edge_dim: int,
	num_relations: int,
	dropout: float,
	residual: bool = False,
	) -> None:
	super().__init__()
	self.residual = residual
	self.self_proj = nn.Linear(hidden_dim, hidden_dim)
	self.message_proj = nn.Linear(
	hidden_dim * 2 + hidden_dim + edge_dim, hidden_dim
	)
	self.relation_embedding = nn.Embedding(num_relations, hidden_dim)
	self.norm = nn.LayerNorm(hidden_dim)
	self.dropout = nn.Dropout(dropout)

	def _skip(self, node_states: Tensor) -> Tensor \| float:
	"""The identity path, when enabled; it keeps depth from oversmoothing."""
	return node_states if self.residual else 0.0

	def forward(
	self,
	node_states: Tensor,
	graph: GraphTensors,
	relation_ids: Tensor,
	query_state: Tensor \| None = None,
	edge_query: Tensor \| None = None,
	) -> tuple[Tensor, Tensor, Tensor \| None]:
	source, target = graph.edge_index
	source_states = node_states[source]
	target_states = node_states[target]
	relation_states = self.relation_embedding(relation_ids)
	query_for_edges = _edge_query(relation_states, query_state, edge_query)
	messages = self.message_proj(
	torch.cat(
	(source_states, target_states, relation_states, graph.edge_features),
	dim=-1,
	)
	)
	messages = messages + query_for_edges
	aggregate = torch.zeros_like(node_states)
	aggregate.index_add_(0, target, messages)
	degree = torch.zeros(
	(graph.num_nodes,), device=node_states.device, dtype=node_states.dtype
	)
	degree.index_add_(0, target, torch.ones_like(target, dtype=node_states.dtype))
	aggregate = aggregate / degree.clamp_min(1.0).unsqueeze(-1)
	updated = self.norm(
	self._skip(node_states)
	+ self.self_proj(node_states)
	+ self.dropout(functional.gelu(aggregate))
	)
	return updated, relation_states, None


	class RelationAwareGATLayer(RelationAwareSAGELayer):
	"""GAT ablation that retains per-edge attention for explanation."""

	def __init__(
	self,
	hidden_dim: int,
	edge_dim: int,
	num_relations: int,
	dropout: float,
	residual: bool = False,
	) -> None:
	super().__init__(hidden_dim, edge_dim, num_relations, dropout, residual)
	self.attn = nn.Linear(hidden_dim * 2 + hidden_dim + edge_dim, 1)

	def forward(
	self,
	node_states: Tensor,
	graph: GraphTensors,
	relation_ids: Tensor,
	query_state: Tensor \| None = None,
	edge_query: Tensor \| None = None,
	) -> tuple[Tensor, Tensor, Tensor \| None]:
	source, target = graph.edge_index
	source_states = node_states[source]
	target_states = node_states[target]
	relation_states = self.relation_embedding(relation_ids)
	query_for_edges = _edge_query(relation_states, query_state, edge_query)
	edge_input = torch.cat(
	(source_states, target_states, relation_states, graph.edge_features), dim=-1
	)
	logits = self.attn(edge_input).squeeze(-1) + (
	relation_states * query_for_edges
	).sum(-1) / (node_states.shape[-1] ** 0.5)
	weights = torch.zeros_like(logits)
	for node in range(graph.num_nodes):
	mask = target == node
	if mask.any():
	weights[mask] = torch.softmax(logits[mask], dim=0)
	messages = self.message_proj(edge_input) * weights.unsqueeze(-1)
	aggregate = torch.zeros_like(node_states)
	aggregate.index_add_(0, target, messages)
	updated = self.norm(
	self._skip(node_states)
	+ self.self_proj(node_states)
	+ self.dropout(functional.gelu(aggregate))
	)
	return updated, relation_states, weights


	class GraphEncoder(nn.Module):
	"""Inductive graph encoder; no trainable closed-world node-ID table is used."""

	def __init__(self, config: GraphConfig) -> None:
	super().__init__()
	self.config = config
	self.node_type_embedding = nn.Embedding(
	config.num_node_types, config.hidden_dim
	)
	self.node_input = nn.Linear(config.graph_input_dim, config.hidden_dim)
	layer_cls = (
	RelationAwareGATLayer
	if config.graph_encoder == "gat" or config.use_gat
	else RelationAwareSAGELayer
	)
	self.layers = nn.ModuleList(
	[
	layer_cls(
	config.hidden_dim,
	config.edge_input_dim,
	config.num_relations,
	config.dropout,
	residual=config.graph_residual,
	)
	for _ in range(config.graph_layers)
	]
	)
	self.edge_input = nn.Linear(config.edge_input_dim, config.hidden_dim)
	self.final_norm = nn.LayerNorm(config.hidden_dim)
	self.start_embedding = (
	nn.Embedding(2, config.hidden_dim) if config.start_marking else None
	)

	def forward(
	self,
	graph: GraphTensors,
	query_state: Tensor \| None = None,
	*,
	edge_query: Tensor \| None = None,
	start_mask: Tensor \| None = None,
	node_query: Tensor \| None = None,
	) -> GraphEncoderOutput:
	"""Encode ``graph``, conditioned on the query where it is given.

	``edge_query``/``node_query`` give one query row per edge/node, for a
	batch of graphs with different queries; otherwise ``query_state`` is
	shared by the whole graph. ``start_mask`` flags the query's start
	entities and is used when the encoder was built with ``start_marking``.
	"""
	# External graph services are open-world: a tenant may introduce more
	# node/relation labels than the configured embedding tables, so the stable
	# global identities are reduced into the learned vocabulary here. The
	# graph itself is never rewritten: its local type indices are what the
	# caller uses to pool per-edge states into a per-graph relation bank and to
	# map a bank row back to a relation name, and overwriting them with reduced
	# hashes would corrupt that bookkeeping.
	node_type_ids = graph.node_type_identity.remainder(self.config.num_node_types)
	edge_type_ids = graph.edge_type_identity.remainder(self.config.num_relations)
	node_states = self.node_input(graph.node_features) + self.node_type_embedding(
	node_type_ids
	)
	if self.start_embedding is not None and start_mask is not None:
	marked = start_mask.to(torch.long)
	node_states = node_states + self.start_embedding(marked)
	if node_query is None and query_state is not None:
	node_query = query_state.reshape(1, -1).expand_as(node_states)
	if node_query is not None:
	node_states = node_states + marked.unsqueeze(-1).to(
	node_states.dtype
	) * node_query.to(node_states.dtype)
	all_relation_states: list[Tensor] = []
	attention: Tensor \| None = None
	for layer in self.layers:
	node_states, relation_states, layer_attention = layer(
	node_states, graph, edge_type_ids, query_state, edge_query
	)
	if layer_attention is not None:
	attention = layer_attention
	all_relation_states.append(relation_states)
	node_states = self.final_norm(node_states)
	edge_states = self.edge_input(graph.edge_features)
	relation_states = (
	all_relation_states[-1]
	if all_relation_states
	else torch.empty((0, self.config.hidden_dim), device=node_states.device)
	)
	return GraphEncoderOutput(node_states, relation_states, edge_states, attention)