diff --git a/README.md b/README.md index f0d5c64..2b37ac9 100644 --- a/README.md +++ b/README.md @@ -67,6 +67,7 @@ Each stage writes into `data/` and the next reads from there. ```sh uv run python -m graph.crawler --field gnn # snowball crawl from the seed papers +uv run python -m graph.semanticscholar --field gnn # recover citations OpenAlex lost uv run python -m graph.build --field gnn # clean it into a dated snapshot uv run python -m gnn.dataset --field gnn # cut the temporal link-prediction split uv run python -m gnn.train --field gnn # fit, and score on citations it never saw @@ -75,6 +76,11 @@ uv run python -m gnn.lineages --field gnn # judge the rankings against fie uv run python -m llm.narrative --field gnn --target W2519887557 ``` +Recovering citations is optional. OpenAlex lists about one reference in eight +under a record it has since deleted, and Semantic Scholar still has most of them. +It takes a few minutes on the shared rate limit; an `S2_API_KEY` in the +environment raises it. + The last command is what the rest is for: the ancestors that carried the paper, why, and cited to the abstracts it read. `python -m gnn.originators` prints the same ranking without the prose, next to the unlearned baselines it has to beat. diff --git a/src/gnn/PLAN.md b/src/gnn/PLAN.md index 2348253..34cc25d 100644 --- a/src/gnn/PLAN.md +++ b/src/gnn/PLAN.md @@ -53,18 +53,19 @@ One PR each, named for what it does. - Hashed title words: -0.02, overfits. - `topic_id`: +0.06, but it leaks. OpenAlex assigns topics from citations as of the crawl, the reason `cited_by_count` is excluded. - Held-out share, learning rate, depth, width, dropout, weight decay, and MRR instead of AP for checkpoint selection: nothing beyond noise. -- [ ] **Originator ranking evaluation.** `gnn.lineages` is built: Recall@10 and nDCG@10 over reachable gold originators for the learned ranking, the same walk with equal weights, and every baseline; coverage apart; top-10 overlap across 3 training seeds and with 10% of citations dropped. Waiting on the gold set, `fields/gold/gnn.yaml` (not `fields/*.yaml`, which the API lists as fields): 18 targets drafted, being verified. On the draft, served snapshot: - - | Scorer | Recall@10 | nDCG@10 | Overlap, 10% dropped | - |---|---|---|---| - | learned | 0.495 | 0.484 | 0.844 | - | equal-weight walk | 0.444 | 0.395 | 0.856 | - | time_decayed_pagerank | 0.432 | 0.380 | 0.833 | - | in_degree | 0.127 | 0.076 | 0.833 | - | gateway | 0.074 | 0.097 | 0.822 | - | path_weight | 0.000 | 0.000 | 0.939 | - - Coverage 90/94; learned top-10 overlap across seeds 0.961. `path_weight` ranks 1980s foundations first, the failure time decay exists to prevent. +- [ ] **Originator ranking evaluation.** `gnn.lineages` is built: Recall@10 and nDCG@10 over reachable gold originators for the learned ranking, the same walk with equal weights, and every baseline; coverage apart; top-10 overlap across 3 training seeds and with 10% of citations dropped. Waiting on the gold set, `fields/gold/gnn.yaml` (not `fields/*.yaml`, which the API lists as fields): 18 targets drafted, being verified. On the draft, served snapshot, before and after recovering citations from Semantic Scholar: + + | Scorer | Recall@10 | nDCG@10 | Overlap, 10% dropped | Gold found | + |---|---|---|---|---| + | learned | 0.484 / 0.465 | 0.466 / 0.439 | 0.844 / 0.894 | 45 / 44 | + | equal-weight walk | 0.410 / 0.397 | 0.366 / 0.333 | 0.844 / 0.867 | 38 / 38 | + | time_decayed_pagerank | 0.432 / 0.401 | 0.380 / 0.352 | 0.833 / 0.872 | 40 / 38 | + | in_degree | 0.127 / 0.097 | 0.076 / 0.060 | 0.833 / 0.861 | | + | gateway | 0.063 / 0.093 | 0.081 / 0.117 | 0.811 / 0.778 | | + | path_weight | 0.000 / 0.000 | 0.000 / 0.000 | 0.900 / 0.883 | | + + Coverage 90/94, then 92/94; learned top-10 overlap across seeds 0.94 to 0.95. The learned ranking leads, and the model's weights are worth +0.07 recall over the same walk unweighted. Recovered citations make more gold originators reachable and every ranking steadier when citations go missing, but find no more of them: recall is over reachable originators, so it drops as the denominator grows. `path_weight` ranks 1980s foundations first, the failure time decay exists to prevent. +- [x] **Recovered citations.** `graph.semanticscholar`, an optional stage between crawl and build: 6811 MAG-era works' Semantic Scholar references matched back to crawled works by MAG id, else by title and year. 22582 citations matched, 6808 of them lost from OpenAlex; GCN regains Bruna, Henaff and Weston. Served. - [ ] **Citation influence labels.** Semantic Scholar `isInfluential` per citation; per-paper AUC of the decoder on influential vs incidental references. - [ ] **Influence-aligned objective.** Fine-tune the decoder listwise on influential references, link prediction as pretraining. Validate on the gold set. - [x] **Crawl expansion.** Uncapped backward crawl, plus `forward_depth: 1`: the in-topic works citing the seeds. 2789 works and 15k citations became 9249 and 104k. GAT is seeded by its primary record. Cleaning re-dates 262 works OpenAlex dates years after their citers, and a year heavier than its split's share no longer empties the split after it. Gold originators reachable from their targets: 88/94 to 90/94. The other four are references OpenAlex holds under dead ids (about 12% of all references), among them GCN's to Bruna, Henaff and Planetoid. @@ -73,6 +74,5 @@ Optional: blinded LLM pairwise judge, trusted only after it agrees with the gold ## Open decisions -- Semantic Scholar as a second data source: citation influence labels, and the references OpenAlex holds under dead ids. - Gold set verification, then commit `fields/gold/gnn.yaml` and rerun `gnn.lineages`. - Abstract embeddings need a new dependency; deferred until originator ranking evaluation shows text is the bottleneck. diff --git a/src/graph/build.py b/src/graph/build.py index 785d2cd..09016df 100644 --- a/src/graph/build.py +++ b/src/graph/build.py @@ -28,6 +28,9 @@ # A preprint can be cited a year before the venue date OpenAlex gives it. PREPRINT_SLACK = 1 +# Citations `graph.semanticscholar` recovered, kept apart from the crawl's own. +RECOVERED_EDGES = "recovered_edges.parquet" + NODE_ATTRIBUTES = ( "title", "abstract", @@ -95,9 +98,11 @@ def is_missing(value: object) -> bool: def load_raw(raw_dir: Path) -> tuple[pd.DataFrame, pd.DataFrame]: - """Read the crawler's raw node and edge tables.""" + """Read the crawler's raw node and edge tables, with any recovered citations.""" nodes = pd.read_parquet(raw_dir / "nodes.parquet") edges = pd.read_parquet(raw_dir / "edges.parquet") + if (recovered := raw_dir / RECOVERED_EDGES).exists(): + edges = pd.concat([edges, pd.read_parquet(recovered)], ignore_index=True) return nodes, edges diff --git a/src/graph/semanticscholar.py b/src/graph/semanticscholar.py new file mode 100644 index 0000000..0aa6b86 --- /dev/null +++ b/src/graph/semanticscholar.py @@ -0,0 +1,203 @@ +"""Recover the references OpenAlex lost, from Semantic Scholar. + +About one reference in eight that OpenAlex lists points at a record it has since +deleted: the id answers 404, and nothing about the work it named survives. GCN +loses its links to Bruna's spectral networks, to Henaff and to Planetoid that +way. Semantic Scholar keeps its own reference lists, and knows most works by the +Microsoft Academic Graph ids that OpenAlex's older ids are made from. + +So every crawled work with such an id has its Semantic Scholar references +matched back to crawled works: by MAG id where OpenAlex kept it, or by title and +year where OpenAlex re-filed the work under a new id. Only citations between +works the crawl holds are added. A work it never reached stays out, as it would +had OpenAlex kept the reference. + +An API key raises the rate limit. It is read from `S2_API_KEY`, never from code. +""" + +from __future__ import annotations + +import argparse +import logging +import os +import re +import time +from collections import defaultdict +from collections.abc import Iterable, Iterator +from pathlib import Path +from typing import Any + +import httpx +import pandas as pd +from pydantic import BaseModel, Field + +from graph.build import PREPRINT_SLACK, RECOVERED_EDGES, is_missing +from graph.config import load_field +from graph.crawler import DATA_DIR + +logger = logging.getLogger(__name__) + +S2_API = "https://api.semanticscholar.org/graph/v1" +MAX_IDS_PER_BATCH = 500 +REFERENCE_FIELDS = "references.title,references.year,references.externalIds" + +# OpenAlex numbers the works it added itself from 4.2 billion up, past every MAG id. +FIRST_OPENALEX_ID = 4_000_000_000 + +_RETRY_STATUS = frozenset({429, 500, 502, 503, 504}) + + +class SemanticScholarError(RuntimeError): + """The API could not satisfy a request.""" + + +class Reference(BaseModel): + title: str | None = None + year: int | None = None + external_ids: dict[str, Any] | None = Field(default=None, alias="externalIds") + + +def mag_id(work_id: str) -> str | None: + """The MAG id an OpenAlex work id was made from, if it was made from one.""" + digits = work_id.removeprefix("W") + return digits if digits.isdigit() and int(digits) < FIRST_OPENALEX_ID else None + + +class SemanticScholarClient: + """Reference lists for batches of works, with backoff on the shared rate limit.""" + + def __init__( + self, + api_key: str | None = None, + *, + base_url: str = S2_API, + transport: httpx.BaseTransport | None = None, + max_retries: int = 8, + backoff: float = 2.0, + timeout: float = 120.0, + ) -> None: + self.max_retries = max_retries + self.backoff = backoff + key = api_key or os.environ.get("S2_API_KEY") + self._client = httpx.Client( + base_url=base_url, + transport=transport, + timeout=timeout, + headers={"x-api-key": key} if key else None, + ) + + def __enter__(self) -> SemanticScholarClient: + return self + + def __exit__(self, *exc: object) -> None: + self._client.close() + + def references(self, work_ids: list[str]) -> Iterator[tuple[str, list[Reference]]]: + """Each work Semantic Scholar knows by its MAG id, with the works it cites.""" + for start in range(0, len(work_ids), MAX_IDS_PER_BATCH): + chunk = work_ids[start : start + MAX_IDS_PER_BATCH] + papers = self._post({"ids": [f"MAG:{mag_id(work_id)}" for work_id in chunk]}) + for work_id, paper in zip(chunk, papers, strict=True): + if paper: + cited = paper.get("references") or [] + yield work_id, [Reference.model_validate(reference) for reference in cited] + + def _post(self, body: dict[str, list[str]]) -> list[dict[str, Any] | None]: + for attempt in range(self.max_retries): + response = self._client.post( + "/paper/batch", params={"fields": REFERENCE_FIELDS}, json=body + ) + if response.status_code in _RETRY_STATUS: + retry_after = response.headers.get("Retry-After", "") + time.sleep( + float(retry_after) if retry_after.isdigit() else self.backoff * 2**attempt + ) + continue + if response.is_error: + raise SemanticScholarError(f"{response.status_code} from {response.request.url}") + return response.json() + raise SemanticScholarError(f"gave up after {self.max_retries} retries") + + +def _normal(title: object) -> str: + return "" if is_missing(title) else re.sub(r"[^a-z0-9]", "", str(title).lower()) + + +def recover_edges( + nodes: pd.DataFrame, references: Iterable[tuple[str, list[Reference]]] +) -> pd.DataFrame: + """Citations Semantic Scholar records between crawled works, in OpenAlex ids. + + A title shared by works more than a year apart, or by two records of one + year, names no single work, so it matches nothing. + """ + known = set(nodes["id"]) + by_title: dict[str, list[tuple[str, object]]] = defaultdict(list) + for work_id, title, year in nodes[["id", "title", "publication_year"]].itertuples(index=False): + if key := _normal(title): + by_title[key].append((work_id, year)) + + def match(reference: Reference) -> str | None: + mag = (reference.external_ids or {}).get("MAG") + if mag and f"W{mag}" in known: + return f"W{mag}" + candidates = [ + work_id + for work_id, year in by_title.get(_normal(reference.title), []) + if reference.year is None + or is_missing(year) + or abs(year - reference.year) <= PREPRINT_SLACK + ] + return candidates[0] if len(candidates) == 1 else None + + edges = { + (source, target) + for source, cited in references + for reference in cited + if (target := match(reference)) and target != source + } + return pd.DataFrame(sorted(edges), columns=["source", "target"]) + + +def recover_field( + field_name: str, + *, + data_dir: Path = DATA_DIR, + client: SemanticScholarClient | None = None, +) -> tuple[Path, int, int]: + """Recover a field's lost citations beside its raw crawl. + + Returns where they went, how many were matched, and how many of those + OpenAlex had not listed. + """ + raw = data_dir / load_field(field_name).name / "raw" + nodes = pd.read_parquet(raw / "nodes.parquet") + listed = pd.read_parquet(raw / "edges.parquet") + queried = [work_id for work_id in nodes["id"] if mag_id(work_id)] + logger.info("asking Semantic Scholar for %d works' references", len(queried)) + + with client or SemanticScholarClient() as s2: + recovered = recover_edges(nodes, s2.references(queried)) + + already = set(zip(listed["source"], listed["target"], strict=True)) + new = sum(edge not in already for edge in recovered.itertuples(index=False, name=None)) + path = raw / RECOVERED_EDGES + recovered.to_parquet(path, index=False) + return path, len(recovered), new + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--field", default="gnn", help="field config name under fields/") + parser.add_argument("--data-dir", type=Path, default=DATA_DIR) + args = parser.parse_args(argv) + + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + logging.getLogger("httpx").setLevel(logging.WARNING) + path, matched, new = recover_field(args.field, data_dir=args.data_dir) + logger.info("%d citations matched, %d missing from OpenAlex", matched, new) + logger.info("wrote %s", path) + + +if __name__ == "__main__": + main() diff --git a/tests/test_semanticscholar.py b/tests/test_semanticscholar.py new file mode 100644 index 0000000..43fdd35 --- /dev/null +++ b/tests/test_semanticscholar.py @@ -0,0 +1,144 @@ +import json + +import httpx +import pandas as pd +import pytest + +from graph import semanticscholar +from graph.build import RECOVERED_EDGES, load_raw +from graph.semanticscholar import ( + Reference, + SemanticScholarClient, + SemanticScholarError, + mag_id, + recover_edges, +) + + +def client_for(handler, **kwargs): + return SemanticScholarClient(transport=httpx.MockTransport(handler), backoff=0, **kwargs) + + +def reference(title, year=2015, mag=None): + return Reference(title=title, year=year, externalIds={"MAG": mag} if mag else None) + + +def nodes(*rows): + return pd.DataFrame(rows, columns=["id", "title", "publication_year"]) + + +def test_only_ids_made_from_mag_ids_are_asked_about(): + assert mag_id("W2519887557") == "2519887557" + assert mag_id("W4297733535") is None + assert mag_id("not-an-id") is None + + +def test_references_are_read_by_mag_id_and_unknown_works_skipped(): + asked = {} + + def handler(request): + asked["ids"] = json.loads(request.content)["ids"] + asked["fields"] = request.url.params["fields"] + chebnet = {"references": [{"title": "ChebNet", "year": 2016}]} + return httpx.Response(200, json=[chebnet, None]) + + found = list(client_for(handler).references(["W1", "W2"])) + + assert asked["ids"] == ["MAG:1", "MAG:2"] + assert "references.externalIds" in asked["fields"] + assert found == [("W1", [Reference(title="ChebNet", year=2016)])] + + +def test_a_rate_limited_batch_is_retried(): + calls = [] + + def handler(request): + calls.append(request) + return httpx.Response(429) if len(calls) == 1 else httpx.Response(200, json=[None]) + + assert list(client_for(handler).references(["W1"])) == [] + assert len(calls) == 2 + + +def test_a_refused_batch_raises_rather_than_passing_for_empty(): + with pytest.raises(SemanticScholarError): + list(client_for(lambda request: httpx.Response(403)).references(["W1"])) + + +def test_the_api_key_is_sent_only_when_one_is_set(monkeypatch): + headers = [] + + def handler(request): + headers.append(request.headers.get("x-api-key")) + return httpx.Response(200, json=[None]) + + monkeypatch.delenv("S2_API_KEY", raising=False) + list(client_for(handler).references(["W1"])) + monkeypatch.setenv("S2_API_KEY", "from-the-environment") + list(client_for(handler).references(["W1"])) + + assert headers == [None, "from-the-environment"] + + +def test_a_reference_is_matched_by_mag_id_first_then_by_title_and_year(): + crawled = nodes( + ("W1", "Citing Work", 2016), + ("W2", "Kept By Id", 2014), + ("W3", "Spectral Networks and Locally Connected Networks on Graphs", 2013), + ) + cited = [ + reference("anything", mag="2"), + reference("Spectral networks and locally-connected networks on graphs.", 2014, mag="9"), + ] + + edges = recover_edges(crawled, [("W1", cited)]) + + assert list(edges.itertuples(index=False, name=None)) == [("W1", "W2"), ("W1", "W3")] + + +def test_a_title_that_names_no_single_work_matches_nothing(): + crawled = nodes( + ("W1", "Citing Work", 2016), + ("W2", "Graph Kernels", 2008), + ("W3", "Graph Kernels", 2008), + ("W4", "Introduction", 1990), + ) + cited = [reference("Graph Kernels", 2008), reference("Introduction", 2010)] + + assert recover_edges(crawled, [("W1", cited)]).empty + + +def test_a_work_is_never_recovered_as_citing_itself(): + crawled = nodes(("W1", "Self Aware", 2016)) + + assert recover_edges(crawled, [("W1", [reference("Self Aware", 2016)])]).empty + + +def test_recovered_citations_join_the_crawl_when_the_graph_is_built(tmp_path): + pd.DataFrame({"id": ["W1", "W2"]}).to_parquet(tmp_path / "nodes.parquet") + pd.DataFrame({"source": ["W1"], "target": ["W2"]}).to_parquet(tmp_path / "edges.parquet") + pd.DataFrame({"source": ["W2"], "target": ["W1"]}).to_parquet(tmp_path / RECOVERED_EDGES) + + _, edges = load_raw(tmp_path) + + assert list(edges.itertuples(index=False, name=None)) == [("W1", "W2"), ("W2", "W1")] + + +def test_a_field_recovers_only_what_openalex_did_not_list(tmp_path, monkeypatch): + raw = tmp_path / "gnn" / "raw" + raw.mkdir(parents=True) + nodes(("W1", "Citing", 2016), ("W2", "Listed", 2014), ("W3", "Lost", 2013)).to_parquet( + raw / "nodes.parquet" + ) + pd.DataFrame({"source": ["W1"], "target": ["W2"]}).to_parquet(raw / "edges.parquet") + + def handler(request): + cited = [{"title": "Listed", "year": 2014}, {"title": "Lost", "year": 2013}] + return httpx.Response(200, json=[{"references": cited}, None, None]) + + path, matched, new = semanticscholar.recover_field( + "gnn", data_dir=tmp_path, client=client_for(handler) + ) + + assert (matched, new) == (2, 1) + assert path == raw / RECOVERED_EDGES