gabfon's picture
Upload folder using huggingface_hub
3e363e2 verified
Raw History Blame Contribute Delete
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,
)
@dataclass(frozen=True)
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)}")