src/wam/pipeline/links.py

"""Links enrichment — find each paper's GitHub / Hugging Face repo.

Repo URLs are almost always in the PDF body ("Code is available at https://github.com/…"),
not the abstract — so we scan the (cached) full text. A keyword-proximity heuristic picks the
authors' own repo over baselines/comparisons. Found URLs are written into ``links_json`` as
``code`` (GitHub) and ``hf`` (Hugging Face); the README and email then surface them.
"""

from __future__ import annotations

import json
import re
import sqlite3

from wam.config import Config
from wam.logging import get_logger
from wam.sources import papers_with_code, pdf

log = get_logger("pipeline.links")

_GH = re.compile(r"https?://github\.com/[A-Za-z0-9_.\-]+/[A-Za-z0-9_.\-]+")
_HF = re.compile(r"https?://huggingface\.co/[A-Za-z0-9_.\-/]+")
# repos that are almost always dependencies/baselines, not the paper's own.
_GH_DENY = ("github.com/huggingface", "github.com/pytorch", "github.com/openai",
            "github.com/google", "github.com/facebookresearch/detectron")
_KEYWORDS = ("code", "project", "available", "github", "hugging", "weights", "released",
             "release", "repo", "page", "checkpoint")


def _clean(url: str) -> str:
    return url.rstrip(").,;:'\"]}").rstrip("/")


def _prep(text: str) -> str:
    """Rejoin URLs broken across PDF lines: drop hyphenated breaks, unwrap newlines."""
    return text.replace("-\n", "").replace("\n", " ")


def _best(text: str, pattern: re.Pattern, deny: tuple = ()) -> str | None:
    cands = [(m.start(), _clean(m.group(0))) for m in pattern.finditer(text)]
    cands = [(p, u) for p, u in cands if not any(d in u.lower() for d in deny)]
    if not cands:
        return None
    low = text.lower()
    for pos, url in cands:  # prefer a URL sitting right after a "code/project/…" cue
        if any(k in low[max(0, pos - 60):pos] for k in _KEYWORDS):
            return url
    return cands[0][1]


def run(cfg: Config, conn: sqlite3.Connection, limit: int | None = None,
        download: bool = True, use_pwc: bool = True) -> int:
    q = ("SELECT id, abstract, links_json FROM papers WHERE track IN ('core','adjacent')")
    if limit:
        q += f" LIMIT {int(limit)}"
    rows = conn.execute(q).fetchall()
    log.info("links: scanning %d papers (download=%s, pwc=%s)", len(rows), download, use_pwc)
    found = 0
    for i, r in enumerate(rows):
        links = json.loads(r["links_json"] or "{}")
        if links.get("code") and links.get("hf"):
            continue
        pdf_url = links.get("pdf") or (f"https://arxiv.org/pdf/{r['id'].split(':',1)[1]}"
                                       if r["id"].startswith("arxiv:") else "")
        text = pdf.get_text(cfg, r["id"], pdf_url if download else "", max_chars=80000)
        blob = _prep((text or "") + "\n" + (r["abstract"] or ""))  # rejoin line-broken URLs
        changed = False
        if blob.strip():
            if not links.get("code"):
                gh = _best(blob, _GH, _GH_DENY)
                if gh:
                    links["code"], changed = gh, True
            if not links.get("hf"):
                hf = _best(blob, _HF)
                if hf:
                    links["hf"], changed = hf, True
        # Fallback: Papers-with-Code (canonical arXiv->repo) when still no code link.
        if use_pwc and not links.get("code") and r["id"].startswith("arxiv:"):
            url = papers_with_code._repo_for_arxiv(r["id"].split(":", 1)[1],
                                                   cfg.get("constants.request_timeout", 90))
            if url:
                links["code"], changed = _clean(url), True
        if changed:
            conn.execute("UPDATE papers SET links_json=? WHERE id=?",
                         (json.dumps(links), r["id"]))
            found += 1
        if i and i % 25 == 0:
            conn.commit()
            log.info("links progress: %d/%d (found %d)", i, len(rows), found)
    conn.commit()
    log.info("links: enriched %d papers with code/HF", found)
    return found