gabfon's picture
Upload folder using huggingface_hub
3e363e2 verified
Raw History Blame Contribute Delete
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
@dataclass(frozen=True)
class EdgeRecord:
source_id: str
relation_id: str
target_id: str
source_type: str
target_type: str
relation_name: str
attributes: Mapping[str, str]
@dataclass(frozen=True)
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,
)
@classmethod
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,
)
@classmethod
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,
)
@classmethod
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}"