"""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"(? 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, )