File size: 1,945 Bytes
1a5ba1e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | """Papers with Code enrichment — has-code / repo link signal.
Best-effort lookup by arXiv id. This is one HTTP call per paper, so callers should pass only
the shortlist (e.g. post-filter core/adjacent papers) and respect ``max_lookups``.
"""
from __future__ import annotations
import time
import requests
from wam.config import Config
from wam.logging import get_logger
from wam.models import PaperRecord
log = get_logger("source.pwc")
API = "https://paperswithcode.com/api/v1/papers/"
def _repo_for_arxiv(aid: str, timeout: int) -> str | None:
try:
r = requests.get(API, params={"arxiv_id": aid}, timeout=timeout)
r.raise_for_status()
results = r.json().get("results") or []
if not results:
return None
pid = results[0]["id"]
rr = requests.get(f"{API}{pid}/repositories/", timeout=timeout)
rr.raise_for_status()
repos = rr.json().get("results") or []
if not repos:
return None
# Prefer the official repo if flagged, else the most-starred.
repos.sort(key=lambda x: (x.get("is_official", False), x.get("stars", 0)), reverse=True)
return repos[0].get("url")
except Exception as e: # noqa: BLE001
log.debug("pwc lookup failed for %s: %s", aid, e)
return None
def enrich(cfg: Config, records: list[PaperRecord], *, max_lookups: int = 60) -> list[PaperRecord]:
if not (cfg.get("sources.papers_with_code", {}) or {}).get("enabled", True):
return records
timeout = cfg.get("constants.request_timeout", 90)
found = 0
for rec in records[:max_lookups]:
if not rec.arxiv_id or rec.links.code:
continue
url = _repo_for_arxiv(rec.arxiv_id, timeout)
if url:
rec.links.code = url
rec.has_code = True
found += 1
time.sleep(0.5)
log.info("pwc found code for %d papers", found)
return records
|