| """Query-conditioned iterative expansion with neural transition scoring. |
| |
| Two changes matter here relative to the earlier implementation. |
| |
| First, candidate expansion no longer filters transitions by |
| ``required_relation_sequence``. That field used to carry the gold composition, |
| so the search was walking a path it had been handed; with the field gone the |
| search must consider every type-legal transition and rely on the model to rank |
| them. |
| |
| Second, ranking is done by the reasoning core and the next-hop head rather than |
| by an MLP over static embeddings, so the beam is steered by a state that has |
| attended over the graph. Hard constraints (PLAN section 17) still prune, because |
| they are genuine query semantics rather than leaked supervision. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Mapping, Sequence |
| from dataclasses import dataclass |
|
|
| import torch |
| from torch import Tensor |
|
|
| from graph.composition import compose_relation_sequence |
| from graph.constraints import path_satisfies_constraints |
| from graph.graph import EdgeRecord, GraphStore |
| from graph.heads import NextHopHead |
| from graph.pat import PathAwareTransformer, PATOutput |
| from graph.reasoning_core import ReasoningCore |
| from graph.types import ( |
| ConstraintState, |
| GraphPath, |
| GraphQuery, |
| PathStep, |
| ReasoningState, |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class SearchOutput: |
| """Completed candidate paths plus the state that produced them.""" |
|
|
| paths: tuple[GraphPath, ...] |
| path_states: Tensor |
| reasoning_states: tuple[ReasoningState, ...] = () |
| pat_output: PATOutput | None = None |
| |
| candidate_states: Tensor | None = None |
| pooled_state: Tensor | None = None |
|
|
|
|
| class BeamSearchReasoner: |
| """Expand graph-local transitions under neural next-hop scoring.""" |
|
|
| def __init__( |
| self, |
| next_hop_head: NextHopHead, |
| beam_width: int, |
| candidate_width: int, |
| max_hops: int, |
| relation_top_k: int = 16, |
| max_completed: int = 256, |
| ) -> None: |
| self.next_hop_head = next_hop_head |
| self.beam_width = beam_width |
| self.candidate_width = candidate_width |
| self.max_hops = max_hops |
| self.relation_top_k = relation_top_k |
| self.max_completed = max_completed |
|
|
| def search( |
| self, |
| graph: GraphStore, |
| query: GraphQuery, |
| query_state: Tensor, |
| node_states: Mapping[str, Tensor], |
| relation_states: Mapping[str, Tensor], |
| pat: PathAwareTransformer, |
| reasoning_core: ReasoningCore, |
| memory: Tensor, |
| memory_mask: Tensor | None, |
| *, |
| use_pat: bool = True, |
| beam_width: int | None = None, |
| query_conditioned: bool = True, |
| ) -> SearchOutput: |
| constraints = ConstraintState.from_query(query) |
| width = beam_width or self.beam_width |
| starts = query.start_entity_ids or query.start_entities |
| conditioning = ( |
| query_state if query_conditioned else torch.zeros_like(query_state) |
| ) |
|
|
| beam: list[GraphPath] = [] |
| for start in starts: |
| for edge in self._candidate_edges(graph, start, constraints, ()): |
| candidate = self._append((), edge) |
| if self._partial_allowed(candidate, constraints): |
| beam.append(candidate) |
|
|
| completed: list[GraphPath] = [] |
| reasoning_states: list[ReasoningState] = [] |
| hop = 0 |
| horizon = min(query.max_hops, self.max_hops) |
|
|
| while beam and hop < horizon: |
| hop += 1 |
| path_states, pat_output = self._encode_paths( |
| beam, |
| node_states, |
| relation_states, |
| conditioning, |
| constraints, |
| pat, |
| use_pat, |
| ) |
| core = reasoning_core( |
| conditioning, memory, memory_mask, path_states.unsqueeze(0), hop=hop |
| ) |
| candidate_states = core.candidate_states[0] |
| reasoning_states.append( |
| self._reasoning_state( |
| beam, |
| conditioning, |
| node_states, |
| relation_states, |
| path_states, |
| hop, |
| constraints, |
| ) |
| ) |
|
|
| for path in beam: |
| scored = GraphPath( |
| path.steps, path.score, path.verified, path.query_compatibility |
| ) |
| if ( |
| constraints.min_hops <= scored.hop_count |
| and path_satisfies_constraints(scored, constraints, graph) |
| ): |
| completed.append(scored) |
|
|
| if hop >= horizon or len(completed) >= self.max_completed: |
| break |
|
|
| beam = self._expand( |
| graph, |
| beam, |
| candidate_states, |
| node_states, |
| relation_states, |
| constraints, |
| width, |
| ) |
|
|
| if not completed: |
| empty = torch.empty((0, conditioning.shape[-1]), device=conditioning.device) |
| return SearchOutput((), empty, tuple(reasoning_states), None, None, None) |
|
|
| ranked = self._dedupe(completed)[: self.max_completed] |
| path_states, pat_output = self._encode_paths( |
| ranked, |
| node_states, |
| relation_states, |
| conditioning, |
| constraints, |
| pat, |
| use_pat, |
| ) |
| core = reasoning_core( |
| conditioning, memory, memory_mask, path_states.unsqueeze(0), hop=horizon |
| ) |
| return SearchOutput( |
| paths=ranked, |
| path_states=path_states, |
| reasoning_states=tuple(reasoning_states), |
| pat_output=pat_output, |
| candidate_states=core.candidate_states[0], |
| pooled_state=core.state[0], |
| ) |
|
|
| |
|
|
| def _expand( |
| self, |
| graph: GraphStore, |
| beam: Sequence[GraphPath], |
| candidate_states: Tensor, |
| node_states: Mapping[str, Tensor], |
| relation_states: Mapping[str, Tensor], |
| constraints: ConstraintState, |
| width: int, |
| ) -> list[GraphPath]: |
| """Score every legal extension and keep the best ``width`` of them.""" |
| extensions: list[GraphPath] = [] |
| rows: list[Tensor] = [] |
| sources: list[Tensor] = [] |
| relations: list[Tensor] = [] |
| targets: list[Tensor] = [] |
| zero = torch.zeros_like(candidate_states[0]) |
| for index, path in enumerate(beam): |
| endpoint = path.endpoint_id |
| if endpoint is None: |
| continue |
| for edge in self._candidate_edges( |
| graph, endpoint, constraints, path.relation_sequence |
| ): |
| extended = self._append(path.steps, edge) |
| if not self._partial_allowed(extended, constraints): |
| continue |
| extensions.append(extended) |
| rows.append(candidate_states[index]) |
| sources.append(node_states.get(edge.source_id, zero)) |
| relations.append(relation_states.get(edge.relation_id, zero)) |
| targets.append(node_states.get(edge.target_id, zero)) |
| if not extensions: |
| return [] |
| scores = self.next_hop_head( |
| torch.stack(rows), |
| torch.stack(sources), |
| torch.stack(relations), |
| torch.stack(targets), |
| ) |
| order = torch.argsort(scores, descending=True)[:width] |
| return [ |
| GraphPath(extensions[int(i)].steps, float(scores[int(i)].item())) |
| for i in order.tolist() |
| ] |
|
|
| def _encode_paths( |
| self, |
| paths: Sequence[GraphPath], |
| node_states: Mapping[str, Tensor], |
| relation_states: Mapping[str, Tensor], |
| conditioning: Tensor, |
| constraints: ConstraintState, |
| pat: PathAwareTransformer, |
| use_pat: bool, |
| ) -> tuple[Tensor, PATOutput | None]: |
| """Encode candidate paths with PAT, falling back to relation composition. |
| |
| The no-PAT ablation must still produce a path representation, otherwise |
| it would be an ablation of the whole path channel rather than of PAT. |
| Relation composition supplies that representation (PLAN section 34). |
| """ |
| composed = torch.stack( |
| [ |
| compose_relation_sequence( |
| path.relation_sequence, relation_states, conditioning |
| ) |
| for path in paths |
| ] |
| ) |
| if not use_pat: |
| return composed, None |
| pat_output = pat.encode_paths( |
| paths, node_states, relation_states, conditioning, constraints |
| ) |
| return pat_output.path_states + composed, pat_output |
|
|
| def _reasoning_state( |
| self, |
| beam: Sequence[GraphPath], |
| conditioning: Tensor, |
| node_states: Mapping[str, Tensor], |
| relation_states: Mapping[str, Tensor], |
| path_states: Tensor, |
| hop: int, |
| constraints: ConstraintState, |
| ) -> ReasoningState: |
| base = conditioning[0] if conditioning.ndim == 2 else conditioning |
| zero = torch.zeros_like(base) |
| return ReasoningState( |
| query_embedding=base, |
| entity_states=torch.stack( |
| [node_states.get(path.endpoint_id or "", zero) for path in beam] |
| ), |
| relation_states=torch.stack( |
| [ |
| relation_states.get(path.steps[-1].relation_id, zero) |
| if path.steps |
| else zero |
| for path in beam |
| ] |
| ), |
| candidate_nodes=torch.arange(len(beam), device=base.device), |
| candidate_edges=torch.arange(len(beam), device=base.device), |
| candidate_paths=tuple(beam), |
| path_representations=path_states, |
| hop=hop, |
| constraint_state=constraints, |
| evidence=tuple(beam), |
| confidence=path_states.norm(dim=-1), |
| ) |
|
|
| @staticmethod |
| def _dedupe(paths: Sequence[GraphPath]) -> tuple[GraphPath, ...]: |
| unique: dict[tuple[tuple[str, str, str], ...], GraphPath] = {} |
| for path in sorted(paths, key=lambda item: item.score, reverse=True): |
| key = tuple((s.source_id, s.relation_id, s.target_id) for s in path.steps) |
| unique.setdefault(key, path) |
| return tuple(unique.values()) |
|
|
| @staticmethod |
| def _append(steps: Sequence[PathStep], edge: EdgeRecord) -> GraphPath: |
| return GraphPath( |
| steps=( |
| *steps, |
| PathStep( |
| source_id=edge.source_id, |
| relation_id=edge.relation_id, |
| target_id=edge.target_id, |
| source_type=edge.source_type, |
| target_type=edge.target_type, |
| edge_attributes=edge.attributes, |
| ), |
| ) |
| ) |
|
|
| def _candidate_edges( |
| self, |
| graph: GraphStore, |
| node_id: str, |
| constraints: ConstraintState, |
| prefix: tuple[str, ...], |
| ) -> tuple[EdgeRecord, ...]: |
| """Legal outgoing transitions, bounded by ``candidate_width``. |
| |
| Only hard query constraints prune here. ``required_relation_sequence`` |
| is honoured when a caller genuinely supplies one (the constraints stage |
| and the API do), but task generation no longer sets it, so on the main |
| task this enumerates every type-legal transition. |
| """ |
| selected: list[EdgeRecord] = [] |
| relations: set[str] = set() |
| required = constraints.required_relation_sequence |
| for edge in graph.outgoing(node_id): |
| if ( |
| required |
| and len(prefix) < len(required) |
| and edge.relation_id != required[len(prefix)] |
| ): |
| continue |
| if edge.relation_id in constraints.forbidden_relations: |
| continue |
| if ( |
| constraints.allowed_relations |
| and edge.relation_id not in constraints.allowed_relations |
| ): |
| continue |
| if ( |
| edge.relation_id not in relations |
| and len(relations) >= self.relation_top_k |
| ): |
| continue |
| relations.add(edge.relation_id) |
| selected.append(edge) |
| if len(selected) >= self.candidate_width: |
| break |
| return tuple(selected) |
|
|
| @staticmethod |
| def _partial_allowed(path: GraphPath, constraints: ConstraintState) -> bool: |
| if not path.steps: |
| return True |
| if constraints.forbidden_relations.intersection(path.relation_sequence): |
| return False |
| if constraints.allowed_relations and not set(path.relation_sequence).issubset( |
| constraints.allowed_relations |
| ): |
| return False |
| required = constraints.required_relation_sequence |
| if ( |
| required |
| and tuple(path.relation_sequence) != required[: len(path.relation_sequence)] |
| ): |
| return False |
| if constraints.forbidden_types.intersection( |
| {step.target_type for step in path.steps if step.target_type} |
| ): |
| return False |
| ids = [path.steps[0].source_id, *(step.target_id for step in path.steps)] |
| |
| return len(ids) == len(set(ids)) |
|
|
