Blog · Kaggle

Finding Data in the Footnotes: An LLM + Constrained-Decoding Pipeline for Make Data Count

Make Data Count — Finding Data References · 10 min read

The problem: citations hiding in plain text

The Make Data Count — Finding Data References competition asks a deceptively simple question: given the full text of a scientific paper, which datasets does it cite, and how? Concretely, for each article I had to emit tuples of (article_id, dataset_id, type), where dataset_id is a data citation (a DOI, or an accession ID like a GEO GSE, an SRA run, a PDB entry, and so on) and type is either Primary (data the authors generated for this study) or Secondary (data they reused from elsewhere). Scoring is F1 over those exact tuples, so a wrong type or a malformed id is just as bad as a miss.

Two things make this hard. First, dataset references live everywhere — methods sections, data-availability statements, figure captions, footnotes — and look almost identical to ordinary paper citations. Second, Primary vs Secondary is a genuinely semantic judgement that hinges on phrasing like "we collected" versus "obtained from." This is the story of how my solution evolved from pure feature engineering to an LLM pipeline with constrained decoding.

Version 1–4: regex candidates and a voting ensemble

My early versions (v2 through v4) were entirely classical ML — no LLM at all. The design was a DataCitationDetector class that did three things.

Text extraction. I read each document from both the XML and the PDF versions, using PyMuPDF (fitz) for the PDFs, preferring XML when available.

Candidate mining via regex. Rather than let a model hallucinate ids, I harvested candidates with a battery of patterns — one for DOIs, and a set for common repository/accession formats:

"doi":    r"(?:https?://)?(?:dx\.)?doi\.org/(10\.\d{4,9}/[\w./-]+)",
"zenodo": r"(?:https?://)?zenodo\.org/record/(\d+)",
"github": r"(?:https?://)?github\.com/[\w\-_]+/[\w\-_]+",
"gse":    r"GSE\d+",  "sra": r"SRA\d+",  "prjna": r"PRJNA\d+",
"chembl": r"CHEMBL\d+", "pdb": r"PDB:\w+", "uniprot": r"UniProt:\w+"

Classification via a feature-engineered ensemble. For each candidate I built a feature dictionary from a ±500-character window around it: counts of "primary" cue phrases (we generated, this study, our data) and "secondary" cue phrases (obtained from, previously published, derived from), which section the mention appeared in (methods, results, data availability, references), proximity to figures/tables/supplements, the citation's own type flags, and a spaCy named-entity count. These went through a DictVectorizer into a soft-voting VotingClassifier over a RandomForest, a GradientBoosting model, and Logistic Regression, evaluated with 5-fold StratifiedKFold weighted F1.

The important lesson came in v4: I added negative sampling. Instead of training only on true citations, I ran the regex miner over the training set and labelled every extracted-but-unlabelled candidate as a third class, Not_A_Citation. That taught the classifier to reject the flood of false positives the regex produced — which is the real bottleneck when precision counts as much as recall in an F1 metric.

The pivot: a keyword-search LLM baseline

The feature engineering hit a ceiling because Primary vs Secondary is fundamentally about reading comprehension. So I built a bridge notebook: keyword search to find candidates, then an LLM to read the context and decide. The flow was:

  1. Extract PDF text with PyMuPDF, truncating at the first references heading to drop the bibliography (where citations are noise, not data).
  2. Find DOI candidates with a loose 10\.\d{4} match, skip any DOI that is the article's own, and grab a ±200-character context chunk.
  3. Feed each chunk to Qwen-3 8B (AWQ) served through vLLM, prompting it to return a JSON object with the normalized DOI and a classification of Primary / Secondary / Not Relevant.
  4. Parse the JSON, keep only Primary/Secondary, then drop any dataset_id appearing in ≥3 articles — a frequency filter to kill boilerplate DOIs (journal, license, and template links) that are never real data.

This notebook wired in the official mdc_fine_grained_eval.compute_metrics so I could track the entity-level F1 (ents_f1) on the 524 training PDFs locally before submitting. Free-form JSON generation, though, is fragile: a stray token breaks the parse and the whole citation is silently dropped in the try/except.

The final pipeline: search + LLM + logits-processor-zoo

My strongest solution keeps the "regex finds candidates, LLM decides" spine but fixes the fragility with constrained decoding, and scales the model up to Qwen2.5-32B-Instruct (AWQ) on vLLM. Here is the shape of it:

flowchart TD
  A[Full-text PDF] --> B[PyMuPDF extract, cut at references/bibliography]
  B --> C[Regex DOI candidate search]
  C --> D[Build 150-char context chunk + data-keyword flag]
  D --> E[LLM pass 1: extract and normalize DOI]
  E --> F[LLM pass 2: classify A/B/C]
  F --> G[MultipleChoiceLogitsProcessor forces single token]
  G --> H[Confidence = top1 - top2 logprob margin]
  H --> I[Filter low-confidence None, dedup by data-context]
  I --> J[submission.csv]

Parsing. Each PDF is read page by page; as soon as a page contains references, bibliography, or works cited, I truncate at that marker and stop. A richer DOI regex — (?:doi\.org/|doi:|DOI:?\s*)?10\.\d{4,}/[^\s\]\)\,\;]+ — catches more real-world formatting than the earlier pattern, and I skip any DOI that overlaps the article's own id.

Candidate context. For every DOI hit I take a ±150-character window and set a boolean has_data_context flag if that window mentions data cues (dataset, data repository, deposited at, available at, and similar). That flag becomes a tie-breaker later.

LLM pass 1 — extraction. A first vLLM call asks the model to extract and normalize just the DOI from the chunk (lowercased, trailing punctuation stripped), with a regex fallback if the model's output doesn't start with 10.. Each id is rebuilt into a canonical https://doi.org/... URL.

LLM pass 2 — constrained classification. This is the heart of it. Instead of asking for free-form text, I frame the Primary/Secondary/None decision as multiple choice — A, B, or C — and constrain generation with MultipleChoiceLogitsProcessor from logits-processor-zoo:

mclp = MultipleChoiceLogitsProcessor(tokenizer, choices=["A", "B", "C"])
outputs = llm.generate(prompts, vllm.SamplingParams(
    max_tokens=1, logits_processors=[mclp],
    logprobs=len(mclp.choices), temperature=0))

The model emits exactly one token, guaranteed to be one of the three choices — no JSON to parse, no formatting failures, nothing to drop. One practical gotcha: vLLM's V1 engine doesn't accept logits processors yet, so the notebook sets VLLM_USE_V1="0" to fall back to the engine that does.

Confidence and dedup. Because I request the logprobs of all three choices, I get a free confidence signal: the margin between the top and second-best logprob. I map A/B/C to Primary / Secondary / None, then use a threshold of 0.5 on that margin to filter out low-confidence None predictions (a shaky "not data" is riskier than a confident one). Finally I dedup on (article_id, dataset_id), sorting so that rows with data context and higher confidence win the tie, drop remaining Nones, and write submission.csv.

What I'd tell my past self

Constrained decoding turned an unreliable text generator into a deterministic classifier. That single change — max_tokens=1 plus a multiple-choice logits processor — removed the entire class of "the JSON didn't parse" bugs that quietly cost recall in the earlier notebook.

Honest caveats. The regex-first, LLM-second design keeps the model on a leash so it can only classify ids that actually exist in the text — but it also means the final pipeline only chases DOIs. The rich accession-ID patterns (GEO, SRA, PDB, UniProt, PRJNA) that I built in the classical versions never made it into the final notebook, which is a clear avenue for more recall. The confidence threshold of 0.5 and the ±150-character window were tuned by feel rather than swept. And while the code wires up the official ents_f1 evaluator on the 524 training PDFs, my final notebook doesn't persist a score, so I'm reluctant to quote a leaderboard number I can't reproduce from the committed cells.

The broader arc is the interesting part: I started by trying to engineer the reading comprehension into features, and ended by letting an LLM do the comprehension while I engineered the guardrails — deterministic candidate mining on the way in, constrained decoding on the way out. For extraction-and-classification tasks judged on exact tuples, that division of labour is the pattern I'd reach for again.

← Back to all posts