| """Typed contracts shared by the GRM model, search, verifier, and API.""" |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Mapping |
| from dataclasses import dataclass, field |
| from enum import StrEnum |
| from typing import Any |
|
|
| import torch |
| from torch import Tensor |
|
|
|
|
| class ReasoningStatus(StrEnum): |
| SUPPORTED = "supported" |
| UNSUPPORTED = "unsupported" |
| CONTRADICTED = "contradicted" |
| INSUFFICIENT_EVIDENCE = "insufficient_evidence" |
| ERROR = "error" |
|
|
|
|
| @dataclass(frozen=True) |
| class GraphTensors: |
| """A validated, batch-agnostic graph tensor contract. |
| |
| ``edge_index`` is shaped ``[2, E]`` and stores source/target node indices. |
| IDs are indices into the graph tensors, never user-facing entity IDs. |
| |
| Two kinds of type id are carried, and the distinction is load-bearing: |
| |
| ``node_type_ids`` / ``edge_type_ids`` |
| *Local* indices -- positions within this graph's sorted type and relation |
| lists. They exist for bookkeeping: pooling per-edge states into a |
| per-graph relation bank, and mapping a bank row back to a relation name. |
| ``node_type_hashes`` / ``edge_type_hashes`` |
| *Global, graph-independent* identities, used for the learned embeddings. |
| |
| Using the local index for the embedding lookup was a significant bug. Local |
| indices come from the alphabetical rank of a name *within one graph*, so a |
| graph that simply lacked one relation shifted every later index: ``owns`` |
| could be embedding row 1 in one graph and row 0 in the next. The relation and |
| type embedding tables could therefore not learn a stable representation of any |
| relation, and renaming labels reshuffled everything, which is what collapsed |
| the unseen-ontology holdout. |
| """ |
|
|
| node_ids: Tensor |
| node_type_ids: Tensor |
| node_features: Tensor |
| edge_index: Tensor |
| edge_type_ids: Tensor |
| edge_features: Tensor |
| graph_id: Tensor |
| |
| |
| node_type_hashes: Tensor | None = None |
| edge_type_hashes: Tensor | None = None |
|
|
| def __post_init__(self) -> None: |
| _require_rank(self.node_ids, 1, "node_ids") |
| _require_rank(self.node_type_ids, 1, "node_type_ids") |
| _require_rank(self.node_features, 2, "node_features") |
| _require_rank(self.edge_index, 2, "edge_index") |
| _require_rank(self.edge_type_ids, 1, "edge_type_ids") |
| _require_rank(self.edge_features, 2, "edge_features") |
| _require_rank(self.graph_id, 1, "graph_id") |
| if self.edge_index.shape[0] != 2: |
| raise ValueError( |
| f"edge_index must have shape [2, E], got {tuple(self.edge_index.shape)}" |
| ) |
| n = self.node_ids.shape[0] |
| e = self.edge_index.shape[1] |
| if self.node_type_ids.shape[0] != n or self.node_features.shape[0] != n: |
| raise ValueError( |
| "node_ids, node_type_ids, and node_features must have the same first dimension" |
| ) |
| if self.edge_type_ids.shape[0] != e or self.edge_features.shape[0] != e: |
| raise ValueError( |
| "edge_index, edge_type_ids, and edge_features must have the same edge dimension" |
| ) |
| if self.graph_id.shape[0] != n: |
| raise ValueError("graph_id must contain one graph ID per node") |
| if self.edge_index.numel() and ( |
| self.edge_index.min() < 0 or self.edge_index.max() >= n |
| ): |
| raise ValueError("edge_index contains a node index outside node_ids") |
| if ( |
| self.node_ids.dtype not in (torch.int32, torch.int64) |
| or self.node_type_ids.dtype not in (torch.int32, torch.int64) |
| or self.edge_type_ids.dtype not in (torch.int32, torch.int64) |
| or self.graph_id.dtype not in (torch.int32, torch.int64) |
| ): |
| raise TypeError( |
| "node, type, edge-type, and graph IDs must be integer tensors" |
| ) |
| if self.edge_index.dtype not in (torch.int32, torch.int64): |
| raise TypeError("edge_index must be an integer tensor") |
| if self.node_type_hashes is not None: |
| _require_rank(self.node_type_hashes, 1, "node_type_hashes") |
| if self.node_type_hashes.shape[0] != n: |
| raise ValueError("node_type_hashes must have one entry per node") |
| if self.edge_type_hashes is not None: |
| _require_rank(self.edge_type_hashes, 1, "edge_type_hashes") |
| if self.edge_type_hashes.shape[0] != e: |
| raise ValueError("edge_type_hashes must have one entry per edge") |
| if ( |
| not self.node_features.is_floating_point() |
| or not self.edge_features.is_floating_point() |
| ): |
| raise TypeError( |
| "node_features and edge_features must be floating-point tensors" |
| ) |
|
|
| @property |
| def num_nodes(self) -> int: |
| return int(self.node_ids.shape[0]) |
|
|
| @property |
| def num_edges(self) -> int: |
| return int(self.edge_index.shape[1]) |
|
|
| @property |
| def node_type_identity(self) -> Tensor: |
| """Global node-type ids when present, else the local indices.""" |
| return ( |
| self.node_type_ids |
| if self.node_type_hashes is None |
| else self.node_type_hashes |
| ) |
|
|
| @property |
| def edge_type_identity(self) -> Tensor: |
| """Global relation ids when present, else the local indices.""" |
| return ( |
| self.edge_type_ids |
| if self.edge_type_hashes is None |
| else self.edge_type_hashes |
| ) |
|
|
| def to(self, device: torch.device | str) -> GraphTensors: |
| return GraphTensors( |
| node_ids=self.node_ids.to(device), |
| node_type_ids=self.node_type_ids.to(device), |
| node_features=self.node_features.to(device), |
| edge_index=self.edge_index.to(device), |
| edge_type_ids=self.edge_type_ids.to(device), |
| edge_features=self.edge_features.to(device), |
| graph_id=self.graph_id.to(device), |
| node_type_hashes=None |
| if self.node_type_hashes is None |
| else self.node_type_hashes.to(device), |
| edge_type_hashes=None |
| if self.edge_type_hashes is None |
| else self.edge_type_hashes.to(device), |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class GraphQuery: |
| text: str | None = None |
| start_entities: tuple[str, ...] = () |
| start_entity_ids: tuple[str, ...] = () |
| allowed_relations: tuple[str, ...] = () |
| required_relations: tuple[str, ...] = () |
| required_relation_sequence: tuple[str, ...] = () |
| forbidden_relations: tuple[str, ...] = () |
| required_types: tuple[str, ...] = () |
| forbidden_types: tuple[str, ...] = () |
| required_intermediate_types: tuple[str, ...] = () |
| max_hops: int = 3 |
| min_hops: int = 1 |
| exact_hops: int | None = None |
| return_paths: bool = True |
| verify: bool = True |
| operation: str = "find_entities" |
| answer_mode: str = "entity" |
| target_entity_ids: tuple[str, ...] = () |
|
|
| def __post_init__(self) -> None: |
| if self.max_hops < 0 or self.min_hops < 0: |
| raise ValueError("hop bounds must be non-negative") |
| if self.min_hops > self.max_hops: |
| raise ValueError("min_hops cannot exceed max_hops") |
| if self.exact_hops is not None and not ( |
| self.min_hops <= self.exact_hops <= self.max_hops |
| ): |
| raise ValueError("exact_hops must be within min_hops and max_hops") |
| if self.answer_mode not in { |
| "entity", |
| "relation", |
| "boolean", |
| "count", |
| "aggregate", |
| }: |
| raise ValueError(f"unsupported answer_mode: {self.answer_mode!r}") |
|
|
| def canonical_text(self) -> str: |
| return "|".join( |
| ( |
| self.operation, |
| self.text or "", |
| ",".join(self.start_entity_ids or self.start_entities), |
| ",".join(self.allowed_relations), |
| ",".join(self.required_relations), |
| ",".join(self.required_relation_sequence), |
| ",".join(self.forbidden_relations), |
| ",".join(self.required_types), |
| ",".join(self.forbidden_types), |
| ",".join(self.required_intermediate_types), |
| str(self.min_hops), |
| str(self.max_hops), |
| str(self.exact_hops), |
| self.answer_mode, |
| ",".join(self.target_entity_ids), |
| ) |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class PathStep: |
| source_id: str |
| relation_id: str |
| target_id: str |
| source_type: str | None = None |
| target_type: str | None = None |
| edge_attributes: Mapping[str, str] = field(default_factory=dict) |
|
|
|
|
| @dataclass(frozen=True) |
| class GraphPath: |
| steps: tuple[PathStep, ...] |
| score: float = 0.0 |
| verified: bool = False |
| query_compatibility: float | None = None |
|
|
| @property |
| def hop_count(self) -> int: |
| return len(self.steps) |
|
|
| @property |
| def start_id(self) -> str | None: |
| return self.steps[0].source_id if self.steps else None |
|
|
| @property |
| def endpoint_id(self) -> str | None: |
| return self.steps[-1].target_id if self.steps else None |
|
|
| @property |
| def relation_sequence(self) -> tuple[str, ...]: |
| return tuple(step.relation_id for step in self.steps) |
|
|
|
|
| @dataclass(frozen=True) |
| class ConstraintState: |
| allowed_relations: frozenset[str] = frozenset() |
| required_relations: frozenset[str] = frozenset() |
| required_relation_sequence: tuple[str, ...] = () |
| forbidden_relations: frozenset[str] = frozenset() |
| required_types: frozenset[str] = frozenset() |
| forbidden_types: frozenset[str] = frozenset() |
| required_intermediate_types: frozenset[str] = frozenset() |
| min_hops: int = 0 |
| max_hops: int = 8 |
| exact_hops: int | None = None |
| require_unique_nodes: bool = False |
|
|
| @classmethod |
| def from_query(cls, query: GraphQuery) -> ConstraintState: |
| return cls( |
| allowed_relations=frozenset(query.allowed_relations), |
| required_relations=frozenset(query.required_relations), |
| required_relation_sequence=query.required_relation_sequence, |
| forbidden_relations=frozenset(query.forbidden_relations), |
| required_types=frozenset(query.required_types), |
| forbidden_types=frozenset(query.forbidden_types), |
| required_intermediate_types=frozenset(query.required_intermediate_types), |
| min_hops=query.min_hops, |
| max_hops=query.max_hops, |
| exact_hops=query.exact_hops, |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class ReasoningState: |
| query_embedding: Tensor |
| entity_states: Tensor |
| relation_states: Tensor |
| candidate_nodes: Tensor |
| candidate_edges: Tensor |
| candidate_paths: tuple[GraphPath, ...] |
| path_representations: Tensor |
| hop: int |
| constraint_state: ConstraintState |
| evidence: tuple[GraphPath, ...] = () |
| confidence: Tensor | None = None |
|
|
|
|
| @dataclass(frozen=True) |
| class EntityAnswer: |
| entity_id: str |
| score: float |
|
|
|
|
| @dataclass(frozen=True) |
| class ReasoningResult: |
| status: ReasoningStatus |
| answers: tuple[EntityAnswer, ...] |
| supporting_paths: tuple[GraphPath, ...] |
| conflicting_paths: tuple[GraphPath, ...] |
| confidence: float |
| calibrated: bool |
| model_version: str |
| graph_version: str |
| reasoning_steps: int |
| metadata: Mapping[str, Any] = field(default_factory=dict) |
| relation_answers: tuple[str, ...] = () |
| boolean_answer: bool | None = None |
| aggregate_value: int | float | None = None |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return { |
| "status": self.status.value, |
| "answers": [ |
| {"entity_id": a.entity_id, "score": a.score} for a in self.answers |
| ], |
| "supporting_paths": [_path_to_dict(p) for p in self.supporting_paths], |
| "conflicting_paths": [_path_to_dict(p) for p in self.conflicting_paths], |
| "confidence": self.confidence, |
| "calibrated": self.calibrated, |
| "model_version": self.model_version, |
| "graph_version": self.graph_version, |
| "reasoning_steps": self.reasoning_steps, |
| "metadata": dict(self.metadata), |
| "relation_answers": list(self.relation_answers), |
| "boolean_answer": self.boolean_answer, |
| "aggregate_value": self.aggregate_value, |
| } |
|
|
|
|
| @dataclass(frozen=True) |
| class GraphExample: |
| graph: GraphTensors |
| node_external_ids: tuple[str, ...] |
| node_types: tuple[str, ...] |
| relation_external_ids: tuple[str, ...] |
| relation_names: tuple[str, ...] |
| edge_attributes: tuple[Mapping[str, str], ...] = () |
| node_attributes: tuple[Mapping[str, str], ...] = () |
| relation_attributes: tuple[Mapping[str, str], ...] = () |
| graph_version: str = "synthetic@0" |
|
|
| def __post_init__(self) -> None: |
| if len(self.node_external_ids) != self.graph.num_nodes: |
| raise ValueError("node_external_ids must align with graph nodes") |
| if len(self.node_types) != self.graph.num_nodes: |
| raise ValueError("node_types must align with graph nodes") |
| if len(self.edge_attributes) not in (0, self.graph.num_edges): |
| raise ValueError("edge_attributes must be empty or align with graph edges") |
| if len(self.node_attributes) not in (0, self.graph.num_nodes): |
| raise ValueError("node_attributes must be empty or align with graph nodes") |
| if len(self.relation_external_ids) != len(self.relation_names): |
| raise ValueError("relation IDs and names must align") |
| if len(self.relation_attributes) not in (0, len(self.relation_external_ids)): |
| raise ValueError( |
| "relation_attributes must be empty or align with relations" |
| ) |
| if len(set(self.node_external_ids)) != len(self.node_external_ids): |
| raise ValueError("node_external_ids must be unique") |
|
|
| @property |
| def node_index(self) -> Mapping[str, int]: |
| return {value: i for i, value in enumerate(self.node_external_ids)} |
|
|
|
|
| def _require_rank(value: Tensor, rank: int, name: str) -> None: |
| if value.ndim != rank: |
| raise ValueError( |
| f"{name} must have rank {rank}, got shape {tuple(value.shape)}" |
| ) |
|
|
|
|
| def _path_to_dict(path: GraphPath) -> dict[str, Any]: |
| return { |
| "steps": [ |
| { |
| "source_id": step.source_id, |
| "relation_id": step.relation_id, |
| "target_id": step.target_id, |
| "source_type": step.source_type, |
| "target_type": step.target_type, |
| "edge_attributes": dict(step.edge_attributes), |
| } |
| for step in path.steps |
| ], |
| "score": path.score, |
| "verified": path.verified, |
| "query_compatibility": path.query_compatibility, |
| } |
|
|
