Spaces:
Running on Zero
Running on Zero
| """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) | |
| 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, | |
| ) | |