leakcheck

Audit train/test leakage within and across ML benchmark datasets.

Benchmarks are routinely assembled from one another, so a "held-out" test function can already sit in a different dataset's training split. leakcheck finds those overlaps, tells you which ids are affected, flags records that carry conflicting labels, and writes a clean copy of a test split with the leaked records removed.

It was written after a real finding: about a third of DiverseVul's test functions also occur in PrimeVul's training split (see Worked example).

• Inputs: JSONL / JSON / CSV / Parquet files (also globs, directories, .gz), or Hugging Face dataset ids (offline-capable), with a text field and optional label / split / id fields.
• Matching: exact, whitespace (default), code (comment-aware), plus optional --near near-duplicate detection (MinHash/LSH).
• Outputs: split x split and dataset x dataset overlap matrices, the overlapping ids, label-conflict cases, and a filtered test split.
• Pure standard library by default. Extras: [parquet] (pyarrow), [hf] (datasets), [fast] (numpy).

PyPI distribution name: dataset-leakcheck (the name leakcheck is taken on PyPI); the import name and the CLI are both leakcheck. Not yet published from this checkout.

pip install dataset-leakcheck                 # core
pip install "dataset-leakcheck[parquet,hf,fast]"
# from a clone:
pip install -e ".[test]"

Quick start

# Audit two datasets. NAME:SPLIT=SOURCE puts every row of a file into that split.
leakcheck audit \
  -d primevul:train=pv_train.jsonl -d primevul:test=pv_test.jsonl \
  -d diversevul:test=dv_test.jsonl \
  --text-field func --label-field target \
  -o diversevul.label=label -o primevul.id=idx -o diversevul.id=id \
  --out-dir leak_report

# Write a copy of the test split without anything that occurs in the training split.
leakcheck clean \
  --test diversevul:test=dv_test.jsonl --against primevul:train=pv_train.jsonl \
  --text-field func -o diversevul.id=id --out dv_test_clean.jsonl --report dropped.json

-d NAME[:SPLIT]=SOURCE is repeatable. If a file has a split column use --split-field; Hugging Face datasets contribute their own split names. Field flags set defaults for all datasets; -o NAME.KEY=VALUE overrides one dataset (KEY = text|label|split|id|config|hf_split). If no id field is given the row index is the id.

What it reports

For audit (text to the terminal, plus --json FILE|- and --out-dir DIR):

Records whose text is empty after normalisation never match anything and are counted as empty. --fail-on-overlap makes audit exit 1 if any cross-group overlap exists (useful in CI).

clean keeps the original rows of the test split (all columns, unchanged), drops every record that matches anything in --against, and prints how many were dropped. Options: --dedup (also drop repeats inside the test split, keep first), --drop-conflicts (also drop test records whose text appears with a different label). --out may be .jsonl, .csv or .parquet; --report records every dropped id and the reason.

Normalisers and matching

--max-chars N truncates each text to N characters before normalising, to compare exactly the view a model was trained on (e.g. --max-chars 4000).

Near-duplicates (--near)

Two texts are near-duplicates when the **Jaccard similarity of their sets of 5-token shingles is

= 0.8** (--shingle 5, --threshold 0.8). Tokens are words and punctuation after comment stripping (--lang). MinHash (128 permutations) with banded LSH is used only to find candidates cheaply; every candidate is verified with the exact Jaccard, so reported similarities are exact and nothing below the threshold is reported. The band layout follows the threshold (16x8 at 0.8: about 95% recall at J=0.8, ~100% at J>=0.9; pairs far above the threshold are essentially never missed). Pairs near the threshold can be missed, so treat --near as a lower bound on near-duplication. Exact-matching records are collapsed first so large duplicate groups stay cheap. Near matches count in the matrices, the id list and clean; label-conflict detection stays exact-key only. numpy, if installed, speeds up signatures (results are identical to the pure-Python path).

Library API

import leakcheck as L

src = [
    L.Source("primevul", "pv_train.jsonl", split="train", text_field="func", id_field="idx"),
    L.Source("diversevul", "dv_test.jsonl", split="test", text_field="func", id_field="id",
             label_field="label"),
]
res = L.audit(src, L.Options(normalizer="code", near=True))
res.split_matrix["diversevul/test"]["primevul/train"]   # {"count": ..., "fraction": ...}
res.overlaps          # [{"from", "id", "to", "match_ids", "match", "similarity", ...}]
res.label_conflicts

out = L.clean(src[1], [src[0]], L.Options(normalizer="code"), dedup=True)
out.kept              # list[Record]; r.raw is the original row
out.dropped           # [{"id", "reason", ...}]

Also: audit_records / clean_records for rows you already hold as Records, LeakIndex for custom queries, and the normalisers in leakcheck.normalize.

Hugging Face ids and offline use

-d prime=hf:colin/PrimeVul loads through the optional datasets package, one group per HF split (or just hf_split via -o prime.hf_split=train). leakcheck never downloads by itself beyond what datasets does; set HF_HUB_OFFLINE=1 to force the local cache. Locally cached parquet snapshots can also be passed as a glob, e.g. ~/.cache/huggingface/hub/datasets--bstee615--diversevul/snapshots/*/data/test-*.parquet.

Worked example: PrimeVul -> DiverseVul

Commands only; no dataset contents are included in this repository. Obtain both datasets under their own licences (for example into the Hugging Face cache) and set HF_HUB_OFFLINE=1 if offline.

pip install "dataset-leakcheck[hf,parquet,fast]"

# 1. Cross-dataset audit: how much of DiverseVul's test split sits in PrimeVul's splits?
#    PrimeVul columns: func, target.  DiverseVul columns: func, target.
#    --max-chars 4000 reproduces the 4000-character view the original run trained on.
leakcheck audit \
  -d primevul=hf:colin/PrimeVul \
  -d diversevul=hf:bstee615/diversevul -o diversevul.hf_split=test \
  --text-field func --label-field target \
  --normalizer whitespace --max-chars 4000 \
  --out-dir leak_primevul_diversevul

# Read:  dataset x dataset  [diversevul][primevul]   -> fraction of DiverseVul test records
#        that also occur in PrimeVul;  split x split [diversevul/test][primevul/train] -> the
#        train-overlap that invalidates a "train on PrimeVul, test on DiverseVul" claim.
#        label_conflicts.jsonl lists texts the two datasets label differently.

# 2. A clean DiverseVul test split: remove everything present in PrimeVul train.
leakcheck clean \
  --test diversevul=hf:bstee615/diversevul -o diversevul.hf_split=test \
  --against primevul=hf:colin/PrimeVul -o primevul.hf_split=train \
  --text-field func --label-field target \
  --normalizer whitespace --max-chars 4000 \
  --out diversevul_test_clean.parquet --report dropped.json

# 3. Stricter variants: ignore comment edits, and also catch near-duplicates.
leakcheck audit -d primevul=hf:colin/PrimeVul -d diversevul=hf:bstee615/diversevul \
  --text-field func --normalizer code --near --threshold 0.8 --json -

The motivating finding (whitespace-normalised exact match, 4000-char truncation, DiverseVul test sampled against all of PrimeVul train) was roughly one third of the test functions; re-scoring a model on the strict-clean subset is the point of clean. Your numbers depend on the dataset revisions you use, so run the audit rather than quoting this one.

Development

pip install -e ".[test]" && pytest          # synthetic fixtures only, no network
python -m build                              # sdist + wheel

License

MIT. Copyright (c) 2026 Pallav Pokhriyal.

Input validation (0.2.0)

Bad input is an error, never a silent success. Every failure is {code, message, hint} (--json) with exit 2: LC_E_EMPTY_INPUT (empty, whitespace-only, comment-only or header-only file, []), LC_E_BAD_VALUE (text that is not a string or null), LC_E_MISSING_FIELD (--text-field/--id-field absent from a row, or --label-field/--split-field present in no row; -o NAME.label= with an empty value switches it off for one dataset), LC_E_BAD_SPEC (-o naming a dataset that was not given), LC_E_BAD_FORMAT (rows that are not objects, corrupt/fake .gz, bad parquet, invalid JSON/UTF-8), LC_E_IO (unreadable file), LC_E_USAGE (--shingle/--max-chars <= 0, --threshold outside (0, 1], negative --max-rows, unknown or abbreviated options). A UTF-8 BOM is accepted. A null text is kept and counted as empty.