Download code/graph/evaluation/baselines.py from frontal-labs/graph-reasoner-typed-compositional: direct link, hf CLI and curl.
- Browser
- Download file 15.8 kB
-
https://huggingface.co/frontal-labs/graph-reasoner-typed-compositional/resolve/main/code/graph/evaluation/baselines.py
- Command line
-
hf download hf://frontal-labs/graph-reasoner-typed-compositional/code/graph/evaluation/baselines.py
-
curl -L -o baselines.py https://huggingface.co/frontal-labs/graph-reasoner-typed-compositional/resolve/main/code/graph/evaluation/baselines.py
15.8 kB
| """Deterministic and neural baseline adapters used by the evaluator.""" | |
| from __future__ import annotations | |
| import math | |
| from collections.abc import Callable, Mapping | |
| from dataclasses import dataclass, replace | |
| from itertools import islice | |
| import torch | |
| from graph.graph import GraphStore | |
| from graph.model import GraphModel | |
| from graph.types import ( | |
| EntityAnswer, | |
| GraphPath, | |
| GraphQuery, | |
| PathStep, | |
| ReasoningResult, | |
| ReasoningStatus, | |
| ) | |
| from graph.verification import detect_conflicting_paths, verify_path | |
| #: Every baseline is a function from a graph and a query to a typed result. | |
| Predictor = Callable[[GraphStore, GraphQuery], ReasoningResult] | |
| class BaselineUnavailableError(RuntimeError): | |
| """Raised when a named upstream baseline has no supplied checkpoint/adapter.""" | |
| #: Paths enumerated from each start entity. A cap on the whole search lets the | |
| #: first start's dense neighbourhood starve every later start of paths. | |
| PATHS_PER_START = 8192 | |
| def _bounded_paths(graph: GraphStore, starts, max_hops: int): | |
| for start in starts: | |
| yield from islice(graph.iter_paths((start,), max_hops), PATHS_PER_START) | |
| def heuristic_graph_search(graph: GraphStore, query: GraphQuery) -> ReasoningResult: | |
| paths = tuple( | |
| path | |
| for path in _bounded_paths( | |
| graph, query.start_entity_ids or query.start_entities, query.max_hops | |
| ) | |
| if query.min_hops <= path.hop_count <= query.max_hops | |
| ) | |
| paths = tuple(verify_path(path, graph, query) for path in paths) | |
| valid = tuple(path for path in paths if path.verified) | |
| answers = tuple( | |
| EntityAnswer(endpoint, 1.0) | |
| for endpoint in dict.fromkeys( | |
| path.endpoint_id for path in valid if path.endpoint_id | |
| ) | |
| ) | |
| relation_answers = tuple( | |
| dict.fromkeys(path.steps[-1].relation_id for path in valid if path.steps) | |
| ) | |
| boolean_answer = bool(answers) if query.answer_mode == "boolean" else None | |
| aggregate_value = ( | |
| len(answers) if query.answer_mode in {"count", "aggregate"} else None | |
| ) | |
| entity_answers = answers if query.answer_mode == "entity" else () | |
| conflicts = tuple( | |
| path for group in detect_conflicting_paths(valid, graph) for path in group | |
| ) | |
| status = ( | |
| ReasoningStatus.SUPPORTED | |
| if valid | |
| else ( | |
| ReasoningStatus.UNSUPPORTED | |
| if query.target_entity_ids | |
| else ReasoningStatus.INSUFFICIENT_EVIDENCE | |
| ) | |
| ) | |
| if conflicts and query.operation in {"verify", "boolean", "is_supported"}: | |
| status = ReasoningStatus.CONTRADICTED | |
| if ( | |
| query.target_entity_ids | |
| and valid | |
| and not (set(query.target_entity_ids) & {path.endpoint_id for path in valid}) | |
| ): | |
| status = ReasoningStatus.UNSUPPORTED | |
| if query.answer_mode == "boolean" and status != ReasoningStatus.SUPPORTED: | |
| boolean_answer = ( | |
| False | |
| if status in {ReasoningStatus.UNSUPPORTED, ReasoningStatus.CONTRADICTED} | |
| else None | |
| ) | |
| entity_answers = ( | |
| answers | |
| if query.answer_mode == "entity" and status == ReasoningStatus.SUPPORTED | |
| else () | |
| ) | |
| return ReasoningResult( | |
| status=status, | |
| answers=entity_answers, | |
| supporting_paths=valid | |
| if query.return_paths and status == ReasoningStatus.SUPPORTED | |
| else (), | |
| conflicting_paths=conflicts if query.return_paths else (), | |
| confidence=1.0 if valid else 0.0, | |
| calibrated=False, | |
| model_version="heuristic-graph-search", | |
| graph_version=graph.graph_version, | |
| reasoning_steps=max((path.hop_count for path in valid), default=0), | |
| relation_answers=relation_answers if query.answer_mode == "relation" else (), | |
| boolean_answer=boolean_answer, | |
| aggregate_value=aggregate_value, | |
| ) | |
| def adamic_adar_search(graph: GraphStore, query: GraphQuery) -> ReasoningResult: | |
| """Use local common-neighbor evidence to rank otherwise valid graph paths.""" | |
| starts = query.start_entity_ids or query.start_entities | |
| neighbor_sets = { | |
| node: {edge.target_id for edge in graph.outgoing(node)} for node in graph.nodes | |
| } | |
| paths = tuple( | |
| path | |
| for path in _bounded_paths(graph, starts, query.max_hops) | |
| if query.min_hops <= path.hop_count <= query.max_hops | |
| ) | |
| scored: list[GraphPath] = [] | |
| for path in paths: | |
| checked = verify_path(path, graph, query) | |
| if not checked.verified: | |
| continue | |
| score = 0.0 | |
| if path.start_id and path.endpoint_id: | |
| common = neighbor_sets.get(path.start_id, set()) & neighbor_sets.get( | |
| path.endpoint_id, set() | |
| ) | |
| score = sum( | |
| 1.0 / max(1.0, math.log1p(len(neighbor_sets.get(node, ())))) | |
| for node in common | |
| ) | |
| scored.append(replace(checked, score=score)) | |
| scored.sort(key=lambda path: path.score, reverse=True) | |
| answers = tuple( | |
| EntityAnswer(endpoint, path.score) | |
| for endpoint, path in ( | |
| (path.endpoint_id, path) for path in scored if path.endpoint_id | |
| ) | |
| ) | |
| relation_answers = tuple( | |
| dict.fromkeys(path.steps[-1].relation_id for path in scored if path.steps) | |
| ) | |
| boolean_answer = bool(answers) if query.answer_mode == "boolean" else None | |
| aggregate_value = ( | |
| len(answers) if query.answer_mode in {"count", "aggregate"} else None | |
| ) | |
| conflicts = tuple( | |
| path | |
| for group in detect_conflicting_paths(tuple(scored), graph) | |
| for path in group | |
| ) | |
| status = ( | |
| ReasoningStatus.SUPPORTED | |
| if scored | |
| else ( | |
| ReasoningStatus.UNSUPPORTED | |
| if query.target_entity_ids | |
| else ReasoningStatus.INSUFFICIENT_EVIDENCE | |
| ) | |
| ) | |
| if conflicts and query.operation in {"verify", "boolean", "is_supported"}: | |
| status = ReasoningStatus.CONTRADICTED | |
| if ( | |
| query.target_entity_ids | |
| and scored | |
| and not (set(query.target_entity_ids) & {path.endpoint_id for path in scored}) | |
| ): | |
| status = ReasoningStatus.UNSUPPORTED | |
| if query.answer_mode == "boolean" and status != ReasoningStatus.SUPPORTED: | |
| boolean_answer = ( | |
| False | |
| if status in {ReasoningStatus.UNSUPPORTED, ReasoningStatus.CONTRADICTED} | |
| else None | |
| ) | |
| entity_answers = ( | |
| answers | |
| if query.answer_mode == "entity" and status == ReasoningStatus.SUPPORTED | |
| else () | |
| ) | |
| return ReasoningResult( | |
| status=status, | |
| answers=entity_answers, | |
| supporting_paths=tuple(scored) | |
| if query.return_paths and status == ReasoningStatus.SUPPORTED | |
| else (), | |
| conflicting_paths=conflicts if query.return_paths else (), | |
| confidence=1.0 if scored else 0.0, | |
| calibrated=False, | |
| model_version="adamic-adar", | |
| graph_version=graph.graph_version, | |
| reasoning_steps=max((path.hop_count for path in scored), default=0), | |
| relation_answers=relation_answers if query.answer_mode == "relation" else (), | |
| boolean_answer=boolean_answer, | |
| aggregate_value=aggregate_value, | |
| ) | |
| def neural_link_predictor( | |
| graph: GraphStore, query: GraphQuery, model: GraphModel, *, gat: bool = False | |
| ) -> ReasoningResult: | |
| """One-hop GraphSAGE/GAT link-prediction baseline, separate from GRM search.""" | |
| expected_encoder = "gat" if gat else "sage" | |
| actual_encoder = ( | |
| "gat" if model.config.graph_encoder == "gat" or model.config.use_gat else "sage" | |
| ) | |
| if actual_encoder != expected_encoder: | |
| raise BaselineUnavailableError( | |
| f"{expected_encoder} baseline requires a model configured with graph_encoder={expected_encoder!r}; got {actual_encoder!r}" | |
| ) | |
| example = model.graph_to_example(graph) | |
| # ``graph_to_example`` builds CPU tensors; the encoder lives wherever the | |
| # model was placed. Without this move the baseline works on CPU and raises a | |
| # device mismatch on GPU, which is how it was actually found. | |
| device = next(model.parameters()).device | |
| encoded = model.graph_encoder(example.graph.to(device)) | |
| starts = set(query.start_entity_ids or query.start_entities) | |
| candidates: list[GraphPath] = [] | |
| for edge in graph.edges: | |
| if edge.source_id not in starts: | |
| continue | |
| source_index = example.node_external_ids.index(edge.source_id) | |
| target_index = example.node_external_ids.index(edge.target_id) | |
| score = float( | |
| (encoded.node_states[source_index] * encoded.node_states[target_index]) | |
| .sum() | |
| .detach() | |
| .cpu() | |
| ) | |
| candidate = GraphPath( | |
| ( | |
| PathStep( | |
| edge.source_id, | |
| edge.relation_id, | |
| edge.target_id, | |
| edge.source_type, | |
| edge.target_type, | |
| edge.attributes, | |
| ), | |
| ), | |
| score=score, | |
| ) | |
| checked = verify_path(candidate, graph, query) | |
| if checked.verified: | |
| candidates.append(checked) | |
| candidates.sort(key=lambda path: path.score, reverse=True) | |
| answers = tuple( | |
| EntityAnswer(path.endpoint_id, path.score) | |
| for path in candidates | |
| if path.endpoint_id | |
| ) | |
| return ReasoningResult( | |
| status=ReasoningStatus.SUPPORTED | |
| if candidates | |
| else ReasoningStatus.INSUFFICIENT_EVIDENCE, | |
| answers=answers if query.answer_mode == "entity" else (), | |
| supporting_paths=tuple(candidates[: model.config.answer_top_k]) | |
| if query.return_paths | |
| else (), | |
| conflicting_paths=(), | |
| confidence=float( | |
| torch.sigmoid( | |
| torch.tensor(candidates[0].score if candidates else -4.0) | |
| ).item() | |
| ), | |
| calibrated=False, | |
| model_version=f"{expected_encoder}-link-predictor", | |
| graph_version=graph.graph_version, | |
| reasoning_steps=1 if candidates else 0, | |
| ) | |
| class ReasonOverrides: | |
| """Typed keyword overrides for :meth:`GraphModel.reason`. | |
| A typed record rather than a ``dict[str, object]``: the untyped mapping could | |
| not be splatted into ``reason`` without defeating the type checker, and the | |
| core pipeline is meant to avoid ``Any`` (PLAN section 51). | |
| """ | |
| use_pat: bool | None = None | |
| verify: bool | None = None | |
| select: bool = True | |
| query_conditioned: bool = True | |
| use_types: bool = True | |
| beam_width: int | None = None | |
| def describe(self) -> dict[str, str]: | |
| """Human-readable summary of the fields this ablation changes.""" | |
| default = ReasonOverrides() | |
| return { | |
| field: str(getattr(self, field)) | |
| for field in ( | |
| "use_pat", | |
| "verify", | |
| "select", | |
| "query_conditioned", | |
| "use_types", | |
| "beam_width", | |
| ) | |
| if getattr(self, field) != getattr(default, field) | |
| } | |
| #: Runtime ablations change how one fixed set of weights is used. Ablations that | |
| #: require *different weights* (no path supervision, no PAT supervision) are not | |
| #: listed here: they are produced by training a variant and are dispatched | |
| #: through an explicit predictor instead. Returning an unmodified | |
| #: ``model.reason`` under an ablation name, as the previous implementation did | |
| #: for ``graph-without-query-conditioning`` and ``graph-without-path-supervision``, | |
| #: reports a comparison that never ran. | |
| RUNTIME_ABLATIONS: Mapping[str, ReasonOverrides] = { | |
| "graph": ReasonOverrides(), | |
| "graph-without-pat": ReasonOverrides(use_pat=False), | |
| "graph-without-verifier": ReasonOverrides(verify=False), | |
| "graph-without-selection": ReasonOverrides(select=False), | |
| "graph-without-query-conditioning": ReasonOverrides(query_conditioned=False), | |
| "graph-untyped": ReasonOverrides(use_types=False), | |
| "graph-greedy": ReasonOverrides(beam_width=1), | |
| } | |
| #: Ablations that can only be evaluated against separately trained weights. | |
| TRAINING_ABLATIONS: frozenset[str] = frozenset( | |
| {"graph-without-path-supervision", "graph-without-pat-supervision"} | |
| ) | |
| #: The model was renamed from "grm" to "graph". Evaluation artifacts produced before | |
| #: the rename key their rows by the old names, so those names keep resolving rather | |
| #: than becoming "unknown baseline" errors when an old report is re-read or compared. | |
| DEPRECATED_BASELINE_NAMES: Mapping[str, str] = { | |
| "grm": "graph", | |
| "grm-3b": "graph", | |
| "grm-full": "graph", | |
| **{ | |
| f"grm-{suffix}": f"graph-{suffix}" | |
| for suffix in ( | |
| "without-pat", | |
| "without-verifier", | |
| "without-selection", | |
| "without-query-conditioning", | |
| "untyped", | |
| "greedy", | |
| "without-path-supervision", | |
| "without-pat-supervision", | |
| ) | |
| }, | |
| } | |
| def apply_overrides( | |
| model: GraphModel, graph: GraphStore, query: GraphQuery, overrides: ReasonOverrides | |
| ) -> ReasoningResult: | |
| """Call ``reason`` with an ablation's overrides, fully typed.""" | |
| return model.reason( | |
| graph, | |
| query, | |
| use_pat=overrides.use_pat, | |
| verify=overrides.verify, | |
| beam_width=overrides.beam_width, | |
| use_types=overrides.use_types, | |
| query_conditioned=overrides.query_conditioned, | |
| select=overrides.select, | |
| ) | |
| def make_baseline( | |
| name: str, | |
| model: GraphModel | None = None, | |
| upstream: Predictor | None = None, | |
| ): | |
| """Resolve a baseline name to a predictor, or fail loudly. | |
| ``upstream`` supplies the released G-Reasoner-34M predictor. It is injected | |
| rather than constructed here because it needs the downloaded checkpoint, a text | |
| embedding model and ``torch_geometric``, none of which should become an import | |
| dependency of the evaluation package. | |
| """ | |
| if name == "heuristic": | |
| return heuristic_graph_search | |
| if name == "adamic_adar": | |
| return adamic_adar_search | |
| if name == "g-reasoner-34m": | |
| if upstream is None: | |
| raise BaselineUnavailableError( | |
| "g-reasoner-34m requires the released upstream checkpoint and a text " | |
| "embedder; build one with graph.evaluation.upstream_baseline and pass " | |
| "it as upstream=" | |
| ) | |
| return upstream | |
| if name == "g-reasoner-34m-finetuned": | |
| raise BaselineUnavailableError( | |
| "g-reasoner-34m-finetuned requires fine-tuning the released checkpoint on " | |
| "this task; the released weights alone are the zero-shot baseline" | |
| ) | |
| if name in TRAINING_ABLATIONS: | |
| raise BaselineUnavailableError( | |
| f"{name} is a training-time ablation: evaluate it against weights trained " | |
| f"with that objective disabled, then pass them as an explicit predictor" | |
| ) | |
| if name in {"graphsage", "gat"}: | |
| if model is None: | |
| raise ValueError(f"baseline {name} requires a configured model") | |
| return lambda graph, query: neural_link_predictor( | |
| graph, query, model, gat=name == "gat" | |
| ) | |
| if model is None: | |
| raise ValueError(f"baseline {name} requires a GRM model adapter") | |
| # Pre-rename names stay readable so old artifacts do not become errors. | |
| key = DEPRECATED_BASELINE_NAMES.get(name, name) | |
| if key in RUNTIME_ABLATIONS: | |
| overrides = RUNTIME_ABLATIONS[key] | |
| return lambda graph, query: apply_overrides(model, graph, query, overrides) | |
| raise ValueError(f"unknown baseline {name!r}; known: {sorted(RUNTIME_ABLATIONS)}") | |