Download code/graph/graph.py from frontal-labs/graph-reasoner-typed-compositional: direct link, hf CLI and curl.
- Browser
- Download file 19 kB
-
https://huggingface.co/frontal-labs/graph-reasoner-typed-compositional/resolve/main/code/graph/graph.py
- Command line
-
hf download hf://frontal-labs/graph-reasoner-typed-compositional/code/graph/graph.py
-
curl -L -o graph.py https://huggingface.co/frontal-labs/graph-reasoner-typed-compositional/resolve/main/code/graph/graph.py
19 kB
| """In-memory graph truth and tensor conversion used by GRM inference.""" | |
| from __future__ import annotations | |
| import csv | |
| import re | |
| from collections.abc import Iterable, Iterator, Mapping | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from typing import Any | |
| import torch | |
| from graph.text_features import TextFeaturiser | |
| from graph.types import GraphExample, GraphPath, GraphTensors, PathStep | |
| from graph.utils import stable_hash | |
| #: Modulus for stable global type identities. Large enough that distinct | |
| #: relation or type names rarely collide, while the encoder reduces further into | |
| #: its configured vocabulary size. | |
| GLOBAL_TYPE_MODULUS = 1_048_576 | |
| class EdgeRecord: | |
| source_id: str | |
| relation_id: str | |
| target_id: str | |
| source_type: str | |
| target_type: str | |
| relation_name: str | |
| attributes: Mapping[str, str] | |
| class NodeRecord: | |
| node_id: str | |
| node_type: str | |
| attributes: Mapping[str, str] | |
| class GraphStore: | |
| """Authoritative graph store with deterministic traversal primitives.""" | |
| def __init__( | |
| self, | |
| nodes: Mapping[str, str], | |
| edges: Iterable[EdgeRecord], | |
| *, | |
| graph_version: str = "graph@0", | |
| node_attributes: Mapping[str, Mapping[str, str]] | None = None, | |
| relation_attributes: Mapping[str, Mapping[str, str]] | None = None, | |
| ) -> None: | |
| self.nodes = dict(nodes) | |
| if not self.nodes: | |
| raise ValueError("graph must contain at least one node") | |
| self.edges = tuple(edges) | |
| self.graph_version = graph_version | |
| self.node_attributes = { | |
| key: dict(value) for key, value in (node_attributes or {}).items() | |
| } | |
| self.relation_attributes = { | |
| key: dict(value) for key, value in (relation_attributes or {}).items() | |
| } | |
| self._out: dict[str, tuple[EdgeRecord, ...]] = {} | |
| for edge in self.edges: | |
| if edge.source_id not in self.nodes or edge.target_id not in self.nodes: | |
| raise ValueError(f"edge references unknown node: {edge}") | |
| self._out.setdefault(edge.source_id, ()) | |
| self._out[edge.source_id] = (*self._out[edge.source_id], edge) | |
| def outgoing(self, source_id: str) -> tuple[EdgeRecord, ...]: | |
| return self._out.get(source_id, ()) | |
| def resolve_entities( | |
| self, mentions: Iterable[str], text: str | None = None | |
| ) -> tuple[str, ...]: | |
| """Resolve exact/case-insensitive IDs; never invents an entity.""" | |
| lower = {key.lower(): key for key in self.nodes} | |
| resolved: list[str] = [] | |
| for mention in mentions: | |
| if mention in self.nodes: | |
| resolved.append(mention) | |
| elif mention.lower() in lower: | |
| resolved.append(lower[mention.lower()]) | |
| if not resolved and text: | |
| # Whole-token match only: a substring match resolves "US" inside "bus". | |
| for key in self.nodes: | |
| if re.search(rf"(?<!\w){re.escape(key)}(?!\w)", text, re.IGNORECASE): | |
| resolved.append(key) | |
| return tuple(dict.fromkeys(resolved)) | |
| def has_edge(self, source_id: str, relation_id: str, target_id: str) -> bool: | |
| return any( | |
| edge.relation_id == relation_id and edge.target_id == target_id | |
| for edge in self.outgoing(source_id) | |
| ) | |
| def iter_paths( | |
| self, start_ids: Iterable[str], max_hops: int | |
| ) -> Iterator[GraphPath]: | |
| """Enumerate simple directed paths up to ``max_hops`` in stable order.""" | |
| for start_id in start_ids: | |
| if start_id not in self.nodes: | |
| continue | |
| yield from self._walk(start_id, (), max_hops) | |
| def _walk( | |
| self, current_id: str, steps: tuple[PathStep, ...], remaining: int | |
| ) -> Iterator[GraphPath]: | |
| if steps: | |
| yield GraphPath(steps=steps) | |
| if remaining == 0: | |
| return | |
| for edge in self.outgoing(current_id): | |
| visited = ( | |
| {steps[0].source_id, *(step.target_id for step in steps)} | |
| if steps | |
| else {current_id} | |
| ) | |
| if edge.target_id in visited: | |
| continue | |
| yield from self._walk( | |
| edge.target_id, | |
| ( | |
| *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, | |
| ), | |
| ), | |
| remaining - 1, | |
| ) | |
| def to_example( | |
| self, | |
| *, | |
| feature_dim: int = 8, | |
| edge_feature_dim: int | None = None, | |
| featuriser: TextFeaturiser | None = None, | |
| ) -> GraphExample: | |
| """Convert to typed tensors. | |
| ``edge_feature_dim`` defaults to ``feature_dim`` but must be passed | |
| separately whenever the encoder's ``edge_input_dim`` differs from its | |
| ``graph_input_dim``. Emitting both at one width silently worked only | |
| because every earlier config happened to set the two equal; the shipped | |
| 3B config does not. | |
| """ | |
| edge_dim = feature_dim if edge_feature_dim is None else edge_feature_dim | |
| node_ids = tuple(self.nodes) | |
| node_index = {node_id: i for i, node_id in enumerate(node_ids)} | |
| node_types = tuple(self.nodes[node_id] for node_id in node_ids) | |
| node_type_index = {name: i for i, name in enumerate(sorted(set(node_types)))} | |
| relation_ids = tuple(sorted({edge.relation_id for edge in self.edges})) | |
| relation_index = {name: i for i, name in enumerate(relation_ids)} | |
| edge_index = torch.tensor( | |
| [ | |
| [node_index[e.source_id] for e in self.edges], | |
| [node_index[e.target_id] for e in self.edges], | |
| ], | |
| dtype=torch.long, | |
| ) | |
| if not self.edges: | |
| edge_index = torch.empty((2, 0), dtype=torch.long) | |
| if featuriser is not None: | |
| # Dense semantic features: a name's vector reflects the words it is made | |
| # of, so an unseen relation inherits the geometry of its parts instead of | |
| # landing on an arbitrary unused slot. | |
| node_features = featuriser.encode_identifiers( | |
| [f"{node_id} {self.nodes[node_id]}" for node_id in node_ids] | |
| ) | |
| edge_features = ( | |
| featuriser.encode_identifiers( | |
| [ | |
| f"{edge.relation_id} {edge.relation_name or ''}" | |
| for edge in self.edges | |
| ] | |
| ) | |
| if self.edges | |
| else torch.zeros((0, featuriser.dimension), dtype=torch.float32) | |
| ) | |
| else: | |
| node_features = torch.zeros( | |
| (len(node_ids), feature_dim), dtype=torch.float32 | |
| ) | |
| for i, node_id in enumerate(node_ids): | |
| attributes = self.node_attributes.get(node_id, {}) | |
| node_features[ | |
| i, | |
| stable_hash( | |
| _feature_key(node_id, self.nodes[node_id], attributes), | |
| feature_dim, | |
| ), | |
| ] = 1.0 | |
| edge_features = torch.zeros( | |
| (len(self.edges), edge_dim), dtype=torch.float32 | |
| ) | |
| for i, edge in enumerate(self.edges): | |
| attributes = { | |
| **self.relation_attributes.get(edge.relation_id, {}), | |
| **dict(edge.attributes), | |
| } | |
| edge_features[ | |
| i, | |
| stable_hash( | |
| _feature_key(edge.relation_id, edge.relation_name, attributes), | |
| edge_dim, | |
| ), | |
| ] = 1.0 | |
| tensors = GraphTensors( | |
| node_ids=torch.arange(len(node_ids), dtype=torch.long), | |
| node_type_ids=torch.tensor( | |
| [node_type_index[t] for t in node_types], dtype=torch.long | |
| ), | |
| node_features=node_features, | |
| edge_index=edge_index, | |
| edge_type_ids=torch.tensor( | |
| [relation_index[e.relation_id] for e in self.edges], dtype=torch.long | |
| ), | |
| # Stable identities: a large modulus keeps collisions rare, and the | |
| # encoder reduces into its own table size. These, not the local | |
| # indices, are what the learned embeddings index. | |
| node_type_hashes=torch.tensor( | |
| [ | |
| stable_hash(node_type, GLOBAL_TYPE_MODULUS) | |
| for node_type in node_types | |
| ], | |
| dtype=torch.long, | |
| ), | |
| edge_type_hashes=torch.tensor( | |
| [stable_hash(e.relation_id, GLOBAL_TYPE_MODULUS) for e in self.edges], | |
| dtype=torch.long, | |
| ), | |
| edge_features=edge_features, | |
| graph_id=torch.zeros(len(node_ids), dtype=torch.long), | |
| ) | |
| return GraphExample( | |
| graph=tensors, | |
| node_external_ids=node_ids, | |
| node_types=node_types, | |
| relation_external_ids=relation_ids, | |
| relation_names=relation_ids, | |
| edge_attributes=tuple(e.attributes for e in self.edges), | |
| node_attributes=tuple( | |
| self.node_attributes.get(node_id, {}) for node_id in node_ids | |
| ), | |
| relation_attributes=tuple( | |
| self.relation_attributes.get(relation_id, {}) | |
| for relation_id in relation_ids | |
| ), | |
| graph_version=self.graph_version, | |
| ) | |
| def from_example(cls, example: GraphExample) -> GraphStore: | |
| # One ``tolist`` per tensor: indexing a tensor element by element costs | |
| # a dispatch per edge and dominated this conversion on large graphs. | |
| sources, targets = example.graph.edge_index.tolist() | |
| relations = example.graph.edge_type_ids.tolist() | |
| nodes, types = example.node_external_ids, example.node_types | |
| edges = tuple( | |
| EdgeRecord( | |
| source_id=nodes[source], | |
| relation_id=example.relation_external_ids[relation], | |
| target_id=nodes[target], | |
| source_type=types[source], | |
| target_type=types[target], | |
| relation_name=example.relation_names[relation], | |
| attributes=example.edge_attributes[i] | |
| if example.edge_attributes | |
| else {}, | |
| ) | |
| for i, (source, target, relation) in enumerate( | |
| zip(sources, targets, relations, strict=True) | |
| ) | |
| ) | |
| return cls( | |
| dict(zip(example.node_external_ids, example.node_types, strict=True)), | |
| edges, | |
| graph_version=example.graph_version, | |
| node_attributes=dict( | |
| zip(example.node_external_ids, example.node_attributes, strict=True) | |
| ) | |
| if example.node_attributes | |
| else None, | |
| relation_attributes=dict( | |
| zip( | |
| example.relation_external_ids, | |
| example.relation_attributes, | |
| strict=True, | |
| ) | |
| ) | |
| if example.relation_attributes | |
| else None, | |
| ) | |
| def from_dict( | |
| cls, payload: Mapping[str, Any], *, graph_version: str | None = None | |
| ) -> GraphStore: | |
| """Adapt the existing graph-service JSON shape without weakening core types.""" | |
| raw_nodes = payload.get("nodes", ()) | |
| nodes: dict[str, str] = {} | |
| node_attributes: dict[str, Mapping[str, str]] = {} | |
| if isinstance(raw_nodes, Mapping): | |
| for node_id, value in raw_nodes.items(): | |
| if isinstance(value, Mapping): | |
| nodes[str(node_id)] = str(value.get("type", "entity")) | |
| node_attributes[str(node_id)] = { | |
| str(key): str(item) | |
| for key, item in value.get("attributes", {}).items() | |
| } | |
| else: | |
| nodes[str(node_id)] = str(value) | |
| else: | |
| for value in raw_nodes: | |
| node_id = value.get("id") or value.get("node_id") | |
| if not node_id: | |
| raise ValueError("graph JSON node requires id or node_id") | |
| nodes[str(node_id)] = str(value.get("type", "entity")) | |
| node_attributes[str(node_id)] = { | |
| str(key): str(item) | |
| for key, item in value.get("attributes", {}).items() | |
| } | |
| relation_names: dict[str, str] = {} | |
| relation_attributes: dict[str, dict[str, str]] = {} | |
| raw_relations = payload.get("relations", ()) | |
| if isinstance(raw_relations, Mapping): | |
| relation_items = ( | |
| {"id": key, **value} if isinstance(value, Mapping) else {"id": key} | |
| for key, value in raw_relations.items() | |
| ) | |
| else: | |
| relation_items = iter(raw_relations) | |
| for value in relation_items: | |
| if not isinstance(value, Mapping): | |
| continue | |
| relation_id = ( | |
| value.get("id") or value.get("relation_id") or value.get("name") | |
| ) | |
| if not relation_id: | |
| continue | |
| relation_id = str(relation_id) | |
| relation_names[relation_id] = str(value.get("name", relation_id)) | |
| relation_attributes[relation_id] = { | |
| str(key): str(item) for key, item in value.get("attributes", {}).items() | |
| } | |
| edges: list[EdgeRecord] = [] | |
| for value in payload.get("edges", ()): | |
| source = value.get("source") or value.get("source_id") | |
| relation = value.get("relation") or value.get("relation_id") | |
| target = value.get("target") or value.get("target_id") | |
| if not source or not relation or not target: | |
| raise ValueError( | |
| "graph JSON edge requires source, relation, and target" | |
| ) | |
| edges.append( | |
| EdgeRecord( | |
| source_id=str(source), | |
| relation_id=str(relation), | |
| target_id=str(target), | |
| source_type=nodes.get(str(source), "entity"), | |
| target_type=nodes.get(str(target), "entity"), | |
| relation_name=relation_names.get( | |
| str(relation), str(value.get("relation_name", relation)) | |
| ), | |
| attributes={ | |
| str(key): str(item) | |
| for key, item in value.get("attributes", {}).items() | |
| }, | |
| ) | |
| ) | |
| return cls( | |
| nodes, | |
| edges, | |
| graph_version=graph_version or str(payload.get("graph_version", "json@0")), | |
| node_attributes=node_attributes, | |
| relation_attributes=relation_attributes, | |
| ) | |
| def from_csv( | |
| cls, directory: str | Path, *, graph_version: str = "csv@0" | |
| ) -> GraphStore: | |
| root = Path(directory) | |
| nodes: dict[str, str] = {} | |
| node_attributes: dict[str, dict[str, str]] = {} | |
| with (root / "nodes.csv").open(newline="", encoding="utf-8") as handle: | |
| for row in csv.DictReader(handle): | |
| node_id = row.get("id") or row.get("node_id") or row.get("name") | |
| if not node_id: | |
| raise ValueError("nodes.csv requires id, node_id, or name") | |
| nodes[node_id] = row.get("type", "entity") | |
| node_attributes[node_id] = { | |
| key: value | |
| for key, value in row.items() | |
| if key not in {"id", "node_id", "name", "type"} | |
| and value is not None | |
| } | |
| relation_names: dict[str, str] = {} | |
| relation_attributes: dict[str, dict[str, str]] = {} | |
| relation_path = root / "relations.csv" | |
| if relation_path.exists(): | |
| with relation_path.open(newline="", encoding="utf-8") as handle: | |
| for row in csv.DictReader(handle): | |
| rid = row.get("id") or row.get("relation_id") or row.get("name") | |
| if rid: | |
| relation_names[rid] = row.get("name", rid) | |
| relation_attributes[rid] = { | |
| key: value | |
| for key, value in row.items() | |
| if key not in {"id", "relation_id", "name"} | |
| and value is not None | |
| } | |
| edges: list[EdgeRecord] = [] | |
| with (root / "edges.csv").open(newline="", encoding="utf-8") as handle: | |
| for row in csv.DictReader(handle): | |
| source = row.get("source") or row.get("source_id") | |
| relation = row.get("relation") or row.get("relation_id") | |
| target = row.get("target") or row.get("target_id") | |
| if not source or not relation or not target: | |
| raise ValueError("edges.csv requires source, relation, and target") | |
| edges.append( | |
| EdgeRecord( | |
| source_id=source, | |
| relation_id=relation, | |
| target_id=target, | |
| source_type=nodes.get(source, "entity"), | |
| target_type=nodes.get(target, "entity"), | |
| relation_name=relation_names.get(relation, relation), | |
| attributes={ | |
| k: v | |
| for k, v in row.items() | |
| if k | |
| not in { | |
| "source", | |
| "source_id", | |
| "relation", | |
| "relation_id", | |
| "target", | |
| "target_id", | |
| } | |
| and v is not None | |
| }, | |
| ) | |
| ) | |
| return cls( | |
| nodes, | |
| edges, | |
| graph_version=graph_version, | |
| node_attributes=node_attributes, | |
| relation_attributes=relation_attributes, | |
| ) | |
| def _feature_key(identifier: str, kind: str, attributes: Mapping[str, str]) -> str: | |
| serialized = ";".join(f"{key}={attributes[key]}" for key in sorted(attributes)) | |
| return f"{identifier}|{kind}|{serialized}" | |