| """De-duplication (spec §3.3.5): exact dedup on normalized function text, then |
| near-duplicate detection via 5-gram Jaccard >= 0.9 using MinHash/LSH. |
| |
| Operates at the pair level (a pair is one dedup unit — dropping only one half |
| of a pair would violate pair integrity), keyed by each pair's func_vuln text. |
| """ |
| from __future__ import annotations |
|
|
| from collections import defaultdict |
|
|
| from datasketch import MinHash, MinHashLSH |
|
|
| from src.data.normalize_text import shingles, strip_comments_and_whitespace |
|
|
| MINHASH_PERM = 128 |
| JACCARD_THRESHOLD = 0.9 |
|
|
|
|
| def exact_dedup_pairs(rows: list[dict]) -> tuple[list[dict], dict]: |
| """Drops whole pairs whose func_vuln normalizes to a text seen before |
| (first pair with that text wins). Returns (kept_rows, stats).""" |
| by_pair: dict[str, list[dict]] = defaultdict(list) |
| for r in rows: |
| by_pair[r["pair_id"]].append(r) |
|
|
| seen_norm_text: set[str] = set() |
| kept_pair_ids: set[str] = set() |
| n_dropped_pairs = 0 |
|
|
| for pair_id, pair_rows in by_pair.items(): |
| vuln_row = next((r for r in pair_rows if r["target"] == 1), pair_rows[0]) |
| norm = strip_comments_and_whitespace(vuln_row["func_vuln"]) |
| if norm in seen_norm_text: |
| n_dropped_pairs += 1 |
| continue |
| seen_norm_text.add(norm) |
| kept_pair_ids.add(pair_id) |
|
|
| kept_rows = [r for r in rows if r["pair_id"] in kept_pair_ids] |
| stats = { |
| "n_pairs_before": len(by_pair), |
| "n_pairs_dropped_exact": n_dropped_pairs, |
| "n_pairs_after": len(by_pair) - n_dropped_pairs, |
| } |
| return kept_rows, stats |
|
|
|
|
| def near_dup_pair_groups(rows: list[dict], threshold: float = JACCARD_THRESHOLD) -> list[set[str]]: |
| """Returns groups (as sets of pair_id) of pairs whose func_vuln text is a |
| near-duplicate of another's (Jaccard over 5-gram shingles >= threshold), |
| via MinHash/LSH. Does not drop anything itself — spec §3.3.5 only asks to |
| "log how many samples were removed" and spec §3.3's near-dup check is |
| primarily used as a cross-partition leakage check in the split builder |
| (`tests/test_splits.py`: "no near-dup across partitions (spot-check)"), |
| not necessarily a blanket corpus-wide drop. |
| """ |
| by_pair: dict[str, list[dict]] = defaultdict(list) |
| for r in rows: |
| by_pair[r["pair_id"]].append(r) |
|
|
| lsh = MinHashLSH(threshold=threshold, num_perm=MINHASH_PERM) |
| minhashes: dict[str, MinHash] = {} |
|
|
| for pair_id, pair_rows in by_pair.items(): |
| vuln_row = next((r for r in pair_rows if r["target"] == 1), pair_rows[0]) |
| norm = strip_comments_and_whitespace(vuln_row["func_vuln"]) |
| sh = shingles(norm, n=5) |
| mh = MinHash(num_perm=MINHASH_PERM) |
| for s in sh: |
| mh.update(s.encode("utf-8")) |
| minhashes[pair_id] = mh |
| lsh.insert(pair_id, mh) |
|
|
| visited: set[str] = set() |
| groups: list[set[str]] = [] |
| for pair_id, mh in minhashes.items(): |
| if pair_id in visited: |
| continue |
| matches = set(lsh.query(mh)) |
| if len(matches) > 1: |
| groups.append(matches) |
| visited.update(matches) |
| return groups |
|
|
