| |
| """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"] |
| |
| 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 |
|
|
| 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, ()): |
| 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: |
| |
| 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 |
|
|
| 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, |
| )) |
| |
| 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()) |
|
|
