tinymind-native-8b-remote-handoff/bundle/evaluation/contamination_audit.py

"""Dataset-vs-probe contamination audit for leaderboard-safe training."""

from __future__ import annotations

from collections import Counter
from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
import re
from typing import Any

from evaluation.native_baseline_probe import PROBES as DEEP_CORE_PROBES
from evaluation.native_broad_probe import TASKS as BROAD_TASKS


WORD_RE = re.compile(r"[\w\u0E00-\u0E7F]+", re.UNICODE)


def _text_from_row(row: dict[str, Any]) -> str:
    if isinstance(row.get("messages"), list):
        return "\n".join(str(msg.get("content", "")) for msg in row["messages"] if isinstance(msg, dict))
    return "\n".join(str(row.get(key, "")) for key in ("prompt", "question", "answer", "response", "text"))


def _tokens(text: str) -> list[str]:
    return [tok.lower() for tok in WORD_RE.findall(text) if len(tok) >= 2]


def _ngrams(tokens: list[str], n: int) -> set[str]:
    if len(tokens) < n:
        return set()
    return {" ".join(tokens[i : i + n]) for i in range(len(tokens) - n + 1)}


def _sha(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def _probe_rows() -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for row in DEEP_CORE_PROBES:
        rows.append({"suite": "native_deep_core", "id": row["id"], "prompt": row["prompt"], "must": row.get("must", [])})
    for idx, row in enumerate(BROAD_TASKS):
        rows.append({"suite": "native_broad", "id": row["axis"], "prompt": row["prompt"], "must": row.get("must", []), "idx": idx})
    return rows


def audit_dataset_contamination(
    dataset_path: str | Path,
    out_dir: str | Path,
    *,
    ngram_n: int = 5,
    max_rows: int | None = None,
    high_overlap_threshold: float = 0.55,
) -> dict[str, Any]:
    dataset = Path(dataset_path)
    probes = _probe_rows()
    probe_features = []
    for probe in probes:
        ptoks = _tokens(str(probe["prompt"]) + " " + " ".join(map(str, probe.get("must", []))))
        probe_features.append(
            {
                **probe,
                "prompt_sha256": _sha(str(probe["prompt"])),
                "tokens": set(ptoks),
                "ngrams": _ngrams(ptoks, ngram_n),
            }
        )

    rows_seen = 0
    exact_hits = []
    high_overlap_hits = []
    max_overlap = 0.0
    suite_counts: Counter[str] = Counter()
    domain_counts: Counter[str] = Counter()
    with dataset.open("r", encoding="utf-8") as f:
        for line_no, line in enumerate(f, start=1):
            if max_rows is not None and rows_seen >= max_rows:
                break
            if not line.strip():
                continue
            rows_seen += 1
            row = json.loads(line)
            text = _text_from_row(row)
            low = text.lower()
            toks = set(_tokens(text))
            grams = _ngrams(list(_tokens(text)), ngram_n)
            domain = str(row.get("metadata", {}).get("domain") or row.get("quality_governor", {}).get("domain") or "unknown")
            domain_counts[domain] += 1
            for probe in probe_features:
                prompt = str(probe["prompt"])
                if prompt and prompt.lower() in low:
                    exact_hits.append(
                        {
                            "line_no": line_no,
                            "domain": domain,
                            "suite": probe["suite"],
                            "probe_id": probe["id"],
                            "reason": "exact_prompt_substring",
                            "prompt_sha256": probe["prompt_sha256"],
                        }
                    )
                    suite_counts[str(probe["suite"])] += 1
                denom = max(1, len(probe["tokens"]))
                token_overlap = len(toks & probe["tokens"]) / denom
                ngram_overlap = len(grams & probe["ngrams"]) / max(1, len(probe["ngrams"]))
                score = max(token_overlap, ngram_overlap)
                max_overlap = max(max_overlap, score)
                if score >= high_overlap_threshold:
                    high_overlap_hits.append(
                        {
                            "line_no": line_no,
                            "domain": domain,
                            "suite": probe["suite"],
                            "probe_id": probe["id"],
                            "token_overlap": round(token_overlap, 6),
                            "ngram_overlap": round(ngram_overlap, 6),
                            "score": round(score, 6),
                        }
                    )
                    suite_counts[str(probe["suite"])] += 1

    exact_count = len(exact_hits)
    high_count = len(high_overlap_hits)
    risk = min(100.0, exact_count * 15.0 + high_count * 2.0 + max(0.0, max_overlap - 0.35) * 50.0)
    cleared = exact_count == 0 and high_count == 0 and max_overlap < high_overlap_threshold
    report = {
        "schema_version": "tinymind-contamination-audit-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "dataset_path": str(dataset),
        "rows_seen": rows_seen,
        "probe_suites": sorted({probe["suite"] for probe in probes}),
        "settings": {
            "ngram_n": ngram_n,
            "high_overlap_threshold": high_overlap_threshold,
            "max_rows": max_rows,
        },
        "metrics": {
            "exact_prompt_hits": exact_count,
            "high_overlap_hits": high_count,
            "max_overlap": round(max_overlap, 6),
            "contamination_risk": round(risk, 6),
        },
        "domain_counts": dict(sorted(domain_counts.items())),
        "suite_hit_counts": dict(sorted(suite_counts.items())),
        "exact_hits": exact_hits[:200],
        "high_overlap_hits": high_overlap_hits[:200],
        "claim_gate": {
            "contamination_cleared": cleared,
            "leaderboard_safe_training_allowed": cleared,
            "promotion_blocked_by_contamination": not cleared,
            "reason": "Exact or high-overlap training rows against probe prompts block leaderboard-safe promotion.",
        },
    }
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    json_path = out / "contamination_audit_report.json"
    md_path = out / "contamination_audit_report.md"
    report["json_path"] = str(json_path)
    report["markdown_path"] = str(md_path)
    json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    md_path.write_text(_markdown(report), encoding="utf-8")
    return report


def _markdown(report: dict[str, Any]) -> str:
    lines = [
        "# TinyMind Contamination Audit",
        "",
        f"- Dataset: `{report['dataset_path']}`",
        f"- Rows seen: {report['rows_seen']}",
        f"- Exact prompt hits: {report['metrics']['exact_prompt_hits']}",
        f"- High overlap hits: {report['metrics']['high_overlap_hits']}",
        f"- Max overlap: {report['metrics']['max_overlap']}",
        f"- Contamination risk: {report['metrics']['contamination_risk']}",
        f"- Cleared: {report['claim_gate']['contamination_cleared']}",
    ]
    return "\n".join(lines) + "\n"