#!/usr/bin/env python3 """Decontaminate the training mix against frozen eval sets. Removes any training row that overlaps an evaluation item, so CyberGym / vuln-detection / knowledge-MCQ numbers measure generalization, not memorization. Methods (in order of cost): 1. 13-gram collision -- word-level n-gram inverted index over the eval text; a train row sharing >= --min-shared-ngrams 13-grams with any eval item is contaminated (GPT-3/Llama style). 2. fuzzy whole-text -- difflib ratio for short rows that have no 13-grams, flagged when ratio >= --fuzzy-threshold. 3. embedding (opt-in) -- cosine match via sentence-transformers if installed and --use-embedding is set; otherwise skipped with a note. Eval sources can be: --eval-jsonl PATH JSONL whose user/assistant text (or a `text` field) is the eval item --eval-text PATH plain text / id-per-line file (e.g. CyberGym frozen task ids, or raw code snippets) -- each non-empty line is one eval item Outputs: --output-clean PATH train JSONL with contaminated rows removed (required for training) --report PATH markdown decontamination report (the gate artifact) Stdlib only, so it runs anywhere. Example: python training/scripts/decontaminate.py \ --train data/processed/stage1.ready.normalized.jsonl \ --eval-jsonl data/eval/vuln_detection_test.jsonl \ --eval-jsonl data/eval/knowledge_mcq.jsonl \ --eval-text reports/cybergym/frozen_level1_baseline_tasks.txt \ --output-clean data/processed/stage1.decontam.jsonl \ --report data/decontam/stage1_decontam_report.md """ from __future__ import annotations import argparse import json import re from difflib import SequenceMatcher from pathlib import Path from typing import Any, Iterable WORD_RE = re.compile(r"[A-Za-z0-9_]+") def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("--train", required=True, help="Training JSONL (rows with `messages` or `text`).") parser.add_argument("--eval-jsonl", action="append", default=[], help="Eval JSONL source (repeatable).") parser.add_argument("--eval-text", action="append", default=[], help="Eval text/id-per-line source (repeatable).") parser.add_argument("--output-clean", required=True) parser.add_argument("--report", required=True) parser.add_argument("--n", type=int, default=13, help="N-gram size (default 13).") parser.add_argument("--min-shared-ngrams", type=int, default=1, help="Min shared n-grams to flag a row (default 1 = any collision).") parser.add_argument("--fuzzy-threshold", type=float, default=0.90, help="difflib ratio for short rows without n-grams (default 0.90).") parser.add_argument("--fuzzy-max-eval", type=int, default=2000, help="Cap eval items scanned per short row (perf guard).") parser.add_argument("--use-embedding", action="store_true", help="Additionally use sentence-transformers cosine match if available.") parser.add_argument("--embedding-threshold", type=float, default=0.95) return parser.parse_args() def read_jsonl(path: Path) -> Iterable[dict[str, Any]]: with path.open("r", encoding="utf-8") as fh: for line_no, line in enumerate(fh, start=1): line = line.strip() if not line: continue try: row = json.loads(line) except json.JSONDecodeError as exc: raise ValueError(f"Invalid JSON in {path}:{line_no}: {exc}") from exc if isinstance(row, dict): yield row def row_text(row: dict[str, Any]) -> str: if isinstance(row.get("messages"), list): return " ".join( str(m.get("content", "")) for m in row["messages"] if m.get("role") in {"user", "assistant"} ) if isinstance(row.get("text"), str): return row["text"] # eval JSONL may use other field names parts = [str(row.get(k, "")) for k in ("question", "prompt", "input", "func", "code", "user", "assistant", "output")] return " ".join(p for p in parts if p) def tokens(text: str) -> list[str]: return WORD_RE.findall(text.lower()) def ngrams(toks: list[str], n: int) -> set[str]: if len(toks) < n: return set() return {" ".join(toks[i : i + n]) for i in range(len(toks) - n + 1)} def load_eval_items(eval_jsonl: list[str], eval_text: list[str]) -> list[dict[str, Any]]: items: list[dict[str, Any]] = [] for path in eval_jsonl: p = Path(path) if not p.is_file(): items.append({"_missing": str(p)}) continue for i, row in enumerate(read_jsonl(p)): items.append({"source": path, "id": row.get("id", f"{path}:{i}"), "text": row_text(row)}) for path in eval_text: p = Path(path) if not p.is_file(): items.append({"_missing": str(p)}) continue for i, line in enumerate(p.read_text(encoding="utf-8").splitlines()): line = line.strip() if line and not line.startswith("#"): items.append({"source": path, "id": f"{path}:{i}", "text": line}) return items def build_index(eval_items: list[dict[str, Any]], n: int) -> tuple[dict[str, set[int]], list[set[str]], list[str]]: """Return (ngram -> eval-idx set, per-eval ngram sets, per-eval short text).""" index: dict[str, set[int]] = {} eval_ngrams: list[set[str]] = [] eval_short: list[str] = [] for idx, item in enumerate(eval_items): text = item.get("text", "") toks = tokens(text) grams = ngrams(toks, n) eval_ngrams.append(grams) eval_short.append(text if len(toks) < n else "") for g in grams: index.setdefault(g, set()).add(idx) return index, eval_ngrams, eval_short def try_load_embedder(name: str = "all-MiniLM-L6-v2"): try: from sentence_transformers import SentenceTransformer # type: ignore return SentenceTransformer(name) except Exception: return None def main() -> int: args = parse_args() train_path = Path(args.train) clean_path = Path(args.output_clean) report_path = Path(args.report) clean_path.parent.mkdir(parents=True, exist_ok=True) report_path.parent.mkdir(parents=True, exist_ok=True) eval_items = load_eval_items(args.eval_jsonl, args.eval_text) missing = [i["_missing"] for i in eval_items if "_missing" in i] eval_items = [i for i in eval_items if "_missing" not in i] index, eval_ngrams, eval_short = build_index(eval_items, args.n) short_eval_idx = [i for i, s in enumerate(eval_short) if s] embedder = None eval_embeddings = None embed_note = "disabled" if args.use_embedding: embedder = try_load_embedder() if embedder is None: embed_note = "requested but sentence-transformers unavailable -> skipped" else: eval_embeddings = embedder.encode([i["text"] for i in eval_items], normalize_embeddings=True) embed_note = "enabled (all-MiniLM-L6-v2)" total = 0 kept = 0 flagged: list[dict[str, Any]] = [] reasons = {"ngram": 0, "fuzzy": 0, "embedding": 0} with clean_path.open("w", encoding="utf-8") as out: for row in read_jsonl(train_path): total += 1 text = row_text(row) toks = tokens(text) grams = ngrams(toks, args.n) hit_eval = None reason = None if grams: counts: dict[int, int] = {} for g in grams: for eidx in index.get(g, ()): # eval items sharing this n-gram counts[eidx] = counts.get(eidx, 0) + 1 if counts: best = max(counts, key=counts.get) if counts[best] >= args.min_shared_ngrams: hit_eval, reason = best, "ngram" reasons["ngram"] += 1 else: # short row: fuzzy compare against short eval items for eidx in short_eval_idx[: args.fuzzy_max_eval]: ratio = SequenceMatcher(None, text, eval_short[eidx]).ratio() if ratio >= args.fuzzy_threshold: hit_eval, reason = eidx, "fuzzy" reasons["fuzzy"] += 1 break if hit_eval is None and embedder is not None and eval_embeddings is not None and text.strip(): import numpy as np # type: ignore vec = embedder.encode([text], normalize_embeddings=True)[0] sims = np.asarray(eval_embeddings) @ np.asarray(vec) top = int(sims.argmax()) if float(sims[top]) >= args.embedding_threshold: hit_eval, reason = top, "embedding" reasons["embedding"] += 1 if hit_eval is not None: flagged.append( { "train_id": row.get("id", "?"), "reason": reason, "eval_id": eval_items[hit_eval].get("id"), "eval_source": eval_items[hit_eval].get("source"), } ) continue out.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n") kept += 1 _write_report(report_path, args, total, kept, flagged, reasons, eval_items, missing, embed_note) print(json.dumps( { "train_rows": total, "kept": kept, "removed": len(flagged), "reasons": reasons, "eval_items": len(eval_items), "missing_eval_sources": missing, "clean": str(clean_path), "report": str(report_path), }, indent=2, )) # Non-zero exit if eval sources were declared but missing (gate should not pass silently). return 3 if missing else 0 def _write_report(path, args, total, kept, flagged, reasons, eval_items, missing, embed_note) -> None: lines = [ "# Decontamination Report", "", f"- Train file: `{args.train}`", f"- Eval items: {len(eval_items)} (from {len(args.eval_jsonl)} jsonl + {len(args.eval_text)} text sources)", f"- N-gram size: {args.n}; min shared to flag: {args.min_shared_ngrams}", f"- Fuzzy threshold (short rows): {args.fuzzy_threshold}", f"- Embedding match: {embed_note}", "", f"- Train rows in: **{total}**", f"- Kept (clean): **{kept}**", f"- Removed (contaminated): **{len(flagged)}**", f" - by n-gram: {reasons['ngram']}", f" - by fuzzy: {reasons['fuzzy']}", f" - by embedding: {reasons['embedding']}", f"- Clean output: `{args.output_clean}`", "", ] if missing: lines.append("## ⚠️ Missing eval sources (gate must not pass until resolved)") for m in missing: lines.append(f"- {m}") lines.append("") if flagged: lines.append("## Sample of removed rows (first 50)") lines.append("") lines.append("| train_id | reason | eval_id | eval_source |") lines.append("|---|---|---|---|") for f in flagged[:50]: lines.append(f"| {f['train_id']} | {f['reason']} | {f['eval_id']} | {f['eval_source']} |") if len(flagged) > 50: lines.append("") lines.append(f"... and {len(flagged) - 50} more.") else: lines.append("No contamination detected against the provided eval sources.") lines.append("") Path(path).write_text("\n".join(lines) + "\n", encoding="utf-8") if __name__ == "__main__": raise SystemExit(main())