src/graphtrace/corpus.py

"""Deterministic selection of a bounded, well-connected AI-research corpus from QASPER.

Procedure (documented in docs/PROTOCOL.md, reproduced exactly by ``graphtrace corpus``):

1. Load every QASPER paper in the configured splits (default: test + validation).
2. Scan each paper's text, excluding sections whose heading contains "related work",
   "background" or "reference", against ``DATASET_LEXICON``: a fixed list of public NLP
   benchmark/dataset names with regex variants. A paper "mentions" a dataset if the
   pattern matches at least ``min_mentions`` times.
3. Keep datasets mentioned by at least two candidate papers ("shared datasets").
4. Greedy selection: start with the paper mentioning the most shared datasets; then
   repeatedly add the paper that shares the most datasets with the papers already
   selected (ties: more shared datasets overall, then paper ID). When no remaining
   paper shares a dataset with the selection, continue by shared-dataset count.
   Stop at ``n_papers``.

The lexicon only shapes *which papers* enter the corpus. Graph edges come from LLM
extraction with verbatim evidence; the lexicon is never used as gold.
"""

from __future__ import annotations

import re
from collections import Counter, defaultdict
from dataclasses import dataclass

from .schemas import Document

# name -> regex (case-sensitive unless (?i) is given). Word boundaries are added.
DATASET_LEXICON: dict[str, str] = {
    "SQuAD": r"SQuAD(?:\s?(?:1\.1|2\.0|v1\.1|v2\.0))?",
    "GLUE": r"GLUE",
    "SuperGLUE": r"SuperGLUE",
    "SNLI": r"SNLI",
    "MultiNLI": r"(?:MultiNLI|MNLI)",
    "XNLI": r"XNLI",
    "SICK": r"SICK",
    "SST": r"SST(?:-2|-5)?|Stanford Sentiment Treebank",
    "IMDb": r"(?i:IMDb)",
    "Yelp": r"Yelp",
    "AG News": r"AG(?:'s)? ?News",
    "DBpedia": r"DBpedia",
    "TREC": r"TREC(?:-QA)?",
    "MRPC": r"MRPC",
    "QQP": r"QQP|Quora Question Pairs",
    "STS-B": r"STS-?B",
    "RTE": r"RTE",
    "CoLA": r"CoLA",
    "CoNLL-2003": r"CoNLL[- ]?(?:20)?03",
    "CoNLL-2012": r"CoNLL[- ]?(?:20)?12",
    "OntoNotes": r"OntoNotes(?:\s?5\.0)?",
    "Penn Treebank": r"Penn Treebank|PTB",
    "Universal Dependencies": r"Universal Dependencies",
    "WikiText-103": r"WikiText-?(?:103|2)",
    "One Billion Word": r"One Billion Word|1B Word",
    "WMT": r"WMT(?:\s?'?\d{2,4})?",
    "IWSLT": r"IWSLT(?:\s?'?\d{2,4})?",
    "Europarl": r"Europarl",
    "OpenSubtitles": r"OpenSubtitles",
    "MS MARCO": r"MS ?MARCO",
    "Natural Questions": r"Natural Questions",
    "TriviaQA": r"TriviaQA",
    "HotpotQA": r"HotpotQA",
    "WikiQA": r"WikiQA",
    "NewsQA": r"NewsQA",
    "SearchQA": r"SearchQA",
    "CoQA": r"CoQA",
    "QuAC": r"QuAC",
    "RACE": r"RACE",
    "DROP": r"DROP",
    "BoolQ": r"BoolQ",
    "bAbI": r"bAbI",
    "WikiSQL": r"WikiSQL",
    "CNN/Daily Mail": r"CNN ?/ ?(?:Daily ?Mail|DM)|CNN-DailyMail|CNN/DailyMail",
    "Gigaword": r"Gigaword",
    "XSum": r"XSum",
    "DUC": r"DUC(?:[- ]?20\d\d)?",
    "SemEval": r"SemEval(?:[- ]?20\d\d)?",
    "Switchboard": r"Switchboard",
    "LibriSpeech": r"LibriSpeech",
    "TIMIT": r"TIMIT",
    "MS COCO": r"(?:MS[- ]?)?COCO",
    "VQA": r"VQA(?:\s?v2)?",
    "Visual Genome": r"Visual Genome",
    "Ubuntu Dialogue Corpus": r"Ubuntu Dialogue Corpus",
    "DSTC": r"DSTC\d*",
    "MultiWOZ": r"MultiWOZ",
    "PersonaChat": r"Persona-?Chat",
    "FEVER": r"FEVER",
    "LIAR": r"LIAR",
    "TACRED": r"TACRED",
    "ACE 2005": r"ACE(?:[- ]?(?:20)?05)",
    "FB15k": r"FB15[kK](?:-237)?",
    "WN18": r"WN18(?:RR)?",
    "BioASQ": r"BioASQ",
    "MIMIC": r"MIMIC(?:-III)?",
    "BookCorpus": r"BookCorpus",
    "Amazon Reviews": r"Amazon (?:product )?reviews?",
    "Twitter SemEval": r"SemEval Twitter",
}

METRIC_LEXICON: dict[str, str] = {
    "BLEU": r"BLEU",
    "ROUGE": r"ROUGE(?:-[12L])?",
    "METEOR": r"METEOR",
    "F1": r"F1|F-?score|F-measure",
    "Accuracy": r"(?i:accuracy)",
    "Exact Match": r"Exact Match|\bEM\b",
    "Perplexity": r"(?i:perplexity)",
    "WER": r"WER|word error rate",
    "MRR": r"MRR",
    "AUC": r"AUC|ROC",
}

EXCLUDED_SECTION = re.compile(r"(?i)related work|background|reference|acknowledg")


def _compile(lex: dict[str, str]) -> dict[str, re.Pattern]:
    return {name: re.compile(r"(?<![A-Za-z0-9])(?:" + pat + r")(?![A-Za-z0-9])") for name, pat in lex.items()}


_DATASETS = _compile(DATASET_LEXICON)
_METRICS = _compile(METRIC_LEXICON)


def paper_text(doc: Document) -> str:
    return "\n".join(p for s in doc.sections if not EXCLUDED_SECTION.search(s.heading) for p in s.paragraphs)


def lexicon_hits(doc: Document, patterns: dict[str, re.Pattern] | None = None) -> Counter:
    text = paper_text(doc)
    pats = patterns or _DATASETS
    return Counter({name: n for name, p in pats.items() if (n := len(p.findall(text)))})


def metric_hits(doc: Document) -> Counter:
    return lexicon_hits(doc, _METRICS)


@dataclass
class Selection:
    papers: list[str]
    mentions: dict[str, list[str]]  # paper -> datasets it mentions (min_mentions applied)
    dataset_df: dict[str, int]  # dataset -> number of candidate papers mentioning it
    stats: dict


def select_corpus(docs: dict[str, Document], n_papers: int, min_mentions: int = 2) -> Selection:
    mentions: dict[str, set[str]] = {}
    for doc_id in sorted(docs):
        hits = lexicon_hits(docs[doc_id])
        mentions[doc_id] = {d for d, n in hits.items() if n >= min_mentions}
    df = Counter(d for ds in mentions.values() for d in ds)
    shared = {d for d, n in df.items() if n >= 2}
    msh = {p: ds & shared for p, ds in mentions.items()}

    remaining = {p for p, ds in msh.items() if ds}
    selected: list[str] = []
    covered: Counter = Counter()
    if remaining:
        first = min(remaining, key=lambda p: (-len(msh[p]), p))
        selected.append(first)
        remaining.discard(first)
        covered.update(msh[first])
    while remaining and len(selected) < n_papers:
        best = min(remaining, key=lambda p: (-sum(1 for d in msh[p] if covered[d]), -len(msh[p]), p))
        selected.append(best)
        remaining.discard(best)
        covered.update(msh[best])

    sel = set(selected)
    sel_df = Counter(d for p in selected for d in msh[p])
    pairs = 0
    by_dataset: dict[str, list[str]] = defaultdict(list)
    for p in selected:
        for d in msh[p]:
            by_dataset[d].append(p)
    for d, ps in by_dataset.items():
        pairs += len(ps) * (len(ps) - 1) // 2
    # connected components of the paper-dataset bipartite graph restricted to the selection
    parent = {p: p for p in selected}

    def find(x: str) -> str:
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    for ps in by_dataset.values():
        for p in ps[1:]:
            parent[find(p)] = find(ps[0])
    components = Counter(find(p) for p in selected)
    stats = {
        "candidate_papers": len(docs),
        "papers_with_shared_dataset": sum(1 for ds in msh.values() if ds),
        "lexicon_datasets_seen": len(df),
        "shared_datasets": len(shared),
        "selected_papers": len(selected),
        "datasets_in_selection": len(sel_df),
        "datasets_shared_within_selection": sum(1 for n in sel_df.values() if n >= 2),
        "paper_pairs_sharing_a_dataset": pairs,
        "connected_components": len(components),
        "largest_component": max(components.values()) if components else 0,
        "top_datasets": sel_df.most_common(15),
        "min_mentions": min_mentions,
    }
    return Selection(
        papers=selected,
        mentions={p: sorted(msh[p]) for p in selected if p in sel},
        dataset_df=dict(df),
        stats=stats,
    )