Upload kard/kard_decontam_semantic.py with huggingface_hub

kard/kard_decontam_semantic.py

@@ -0,0 +1,433 @@
+#!/usr/bin/env python3
+"""KARD distillation-data decontamination — STAGE 2: SEMANTIC (embedding cosine).
+
+WHY THIS EXISTS
+---------------
+We distill science-QA from a 27B teacher (Qwen) into a small LM to lift ARC-Easy.
+The teacher has almost certainly *seen* ARC during its own training, so it can
+reproduce an ARC-Easy TEST item as a PARAPHRASE with very low n-gram overlap:
+
+    ARC test  : "What gas do plants absorb from the air?"
+    generated : "Which gas is taken in by vegetation during photosynthesis?"
+
+Same knowledge, <8-gram overlap. A purely lexical (8-gram + exact-match) gate
+lets that through -> the test answer leaks into training -> the ARC-Easy score
+is inflated (fake lift). This module adds the SEMANTIC catch a lexical gate
+structurally cannot make.
+
+TWO-STAGE PIPELINE (ordering is load-bearing)
+---------------------------------------------
+    STAGE 1  (lexical) : kard/kard_decontam.py  -- 8-gram shingles + exact-match.
+                         Catches near-verbatim / high lexical overlap. (teammate's;
+                         NOT touched by this module.)
+    STAGE 2  (semantic): THIS module            -- sentence-embedding max-cosine.
+                         Catches low-lexical-overlap paraphrases.
+
+A candidate is KEPT for training ONLY IF it passes BOTH stages:
+
+    kept  <=>  (passes STAGE-1 lexical)  AND  (passes STAGE-2 semantic)
+
+Run STAGE 2 *after* STAGE 1, on the survivors of STAGE 1, before training on the
+distilled data. The two stages are complementary, not redundant: lexical covers
+verbatim/near-verbatim, semantic covers meaning-preserving rewrites.
+
+TEST SET
+--------
+Production loads the real ARC-Easy TEST split (2376 questions) as `test_items`,
+the SAME source the lexical gate uses, so the two stages agree on what "test" is.
+The CLI accepts jsonl/txt; wire it to the ARC-Easy test dump you already feed the
+lexical gate.
+
+EMBEDDER
+--------
+Default: sentence-transformers/all-MiniLM-L6-v2 (small, CPU-fast, 384-dim).
+If sentence-transformers or the model is unavailable, we fall back to a clearly
+marked char-ngram TF-IDF cosine STUB and print a loud warning. The stub is a
+smoke-test crutch only -- it does NOT have real paraphrase recall. Production
+MUST use the real embedder.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+
+DEFAULT_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+DEFAULT_THRESHOLD = 0.90
+DEFAULT_BATCH = 256
+
+
+# --------------------------------------------------------------------------- #
+# Embedding backends                                                          #
+# --------------------------------------------------------------------------- #
+class _RealEmbedder:
+    """sentence-transformers backend (production)."""
+
+    kind = "sentence-transformers"
+
+    def __init__(self, model_name: str):
+        from sentence_transformers import SentenceTransformer  # lazy import
+
+        self.model_name = model_name
+        self._model = SentenceTransformer(model_name)
+
+    def encode(self, texts: list[str], batch_size: int) -> np.ndarray:
+        emb = self._model.encode(
+            texts,
+            batch_size=batch_size,
+            normalize_embeddings=True,  # unit-norm -> dot product == cosine
+            convert_to_numpy=True,
+            show_progress_bar=False,
+        )
+        return np.asarray(emb, dtype=np.float32)
+
+
+class _TfidfStubEmbedder:
+    """Fallback char-ngram TF-IDF cosine STUB. NOT a real semantic embedder.
+
+    Fitted on the union of candidates+test so both live in one vector space.
+    Char n-grams give it a little robustness to word reordering/inflection, but
+    it has NO real paraphrase recall -- it is a wiring/smoke-test crutch only.
+    """
+
+    kind = "tfidf-stub"
+    model_name = "char-ngram-tfidf(3,5)-STUB"
+
+    def __init__(self):
+        from sklearn.feature_extraction.text import TfidfVectorizer
+
+        self._vec = TfidfVectorizer(analyzer="char_wb", ngram_range=(3, 5))
+        self._fitted = False
+
+    def fit(self, corpus: list[str]) -> None:
+        self._vec.fit(corpus if corpus else [""])
+        self._fitted = True
+
+    def encode(self, texts: list[str], batch_size: int) -> np.ndarray:
+        if not self._fitted:
+            raise RuntimeError("_TfidfStubEmbedder.fit() must be called before encode()")
+        mat = self._vec.transform(texts).astype(np.float32).toarray()
+        norms = np.linalg.norm(mat, axis=1, keepdims=True)
+        norms[norms == 0.0] = 1.0
+        return mat / norms  # unit-norm rows -> dot == cosine
+
+
+def load_embedder(model_name: str = DEFAULT_MODEL, force_stub: bool = False):
+    """Return (embedder, is_stub). Falls back to the TF-IDF stub with a loud warning."""
+    if not force_stub:
+        try:
+            emb = _RealEmbedder(model_name)
+            return emb, False
+        except Exception as exc:  # import failure, download failure, etc.
+            print(
+                "!!! WARNING: could not load sentence-transformers model "
+                f"'{model_name}' ({type(exc).__name__}: {exc}).\n"
+                "!!! Falling back to char-ngram TF-IDF cosine STUB. This is NOT a real\n"
+                "!!! semantic embedder and has NO paraphrase recall. PRODUCTION MUST use\n"
+                "!!! the real embedder (pip install sentence-transformers + model download).",
+                file=sys.stderr,
+            )
+    else:
+        print(
+            "!!! WARNING: --stub forced -> using char-ngram TF-IDF cosine STUB "
+            "(NOT a real semantic embedder).",
+            file=sys.stderr,
+        )
+    return _TfidfStubEmbedder(), True
+
+
+# --------------------------------------------------------------------------- #
+# Core                                                                        #
+# --------------------------------------------------------------------------- #
+def _max_cosine(cand_emb: np.ndarray, test_emb: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """Given unit-norm rows, return per-candidate (max_sim, argmax_test_idx).
+
+    Cosine == dot product for unit-norm rows. Blocked matmul keeps memory bounded
+    for ~thousands of candidates x ~2400 test items.
+    """
+    n_cand = cand_emb.shape[0]
+    max_sim = np.zeros(n_cand, dtype=np.float32)
+    arg = np.zeros(n_cand, dtype=np.int64)
+    if test_emb.shape[0] == 0:
+        return max_sim, arg  # no test items -> nothing similar
+    block = 1024
+    for i in range(0, n_cand, block):
+        sims = cand_emb[i : i + block] @ test_emb.T  # (b, n_test)
+        arg[i : i + block] = np.argmax(sims, axis=1)
+        max_sim[i : i + block] = np.max(sims, axis=1)
+    return max_sim, arg
+
+
+def semantic_contaminated(
+    candidates: list[str],
+    test_items: list[str],
+    threshold: float = DEFAULT_THRESHOLD,
+    batch_size: int = DEFAULT_BATCH,
+    embedder=None,
+) -> list[dict]:
+    """Flag candidates that are semantically too close to any test item.
+
+    For each candidate compute the MAX cosine similarity to any test_item; flag
+    dropped=True if max_sim >= threshold.
+
+    Returns one dict per candidate (input order):
+        {text, max_sim, nearest_test_idx, dropped}
+    """
+    results: list[dict] = []
+    if not candidates:
+        return results
+
+    if embedder is None:
+        embedder, _ = load_embedder()
+
+    # TF-IDF stub must be fitted on the shared vocabulary of both sides first.
+    if isinstance(embedder, _TfidfStubEmbedder):
+        embedder.fit(list(candidates) + list(test_items))
+
+    cand_emb = embedder.encode(list(candidates), batch_size)
+    test_emb = (
+        embedder.encode(list(test_items), batch_size)
+        if test_items
+        else np.zeros((0, cand_emb.shape[1]), dtype=np.float32)
+    )
+
+    max_sim, arg = _max_cosine(cand_emb, test_emb)
+    has_test = test_emb.shape[0] > 0
+    for text, ms, ai in zip(candidates, max_sim, arg):
+        results.append(
+            {
+                "text": text,
+                "max_sim": round(float(ms), 6),
+                "nearest_test_idx": int(ai) if has_test else -1,
+                "dropped": bool(ms >= threshold) if has_test else False,
+            }
+        )
+    return results
+
+
+# --------------------------------------------------------------------------- #
+# I/O                                                                         #
+# --------------------------------------------------------------------------- #
+def _text_from_obj(obj) -> str:
+    """Extract QA text from a parsed jsonl object (dict) or return str as-is."""
+    if isinstance(obj, str):
+        return obj
+    if isinstance(obj, dict):
+        for k in ("text", "question", "prompt", "q", "query", "content"):
+            if k in obj and isinstance(obj[k], str):
+                base = obj[k]
+                break
+        else:
+            base = ""
+        ans = obj.get("answer") or obj.get("correct") or obj.get("a")
+        if isinstance(ans, str) and ans:
+            return f"{base} {ans}".strip()
+        return base
+    return str(obj)
+
+
+def load_items(path: str) -> list[str]:
+    """Load items from .jsonl (one JSON obj/line) or .txt (one item/line).
+
+    Files ending in .jsonl are parsed as JSON per line; anything else (incl .txt)
+    is treated as one raw item per non-empty line, but lines that happen to be
+    JSON objects are still unpacked.
+    """
+    p = Path(path)
+    items: list[str] = []
+    is_jsonl = p.suffix.lower() in (".jsonl", ".json", ".ndjson")
+    with open(p, "r", encoding="utf-8", errors="replace") as f:
+        for line in f:
+            line = line.rstrip("\n")
+            if not line.strip():
+                continue
+            if is_jsonl:
+                obj = json.loads(line)
+                items.append(_text_from_obj(obj))
+            else:
+                s = line.strip()
+                if s[:1] in ("{", "["):
+                    try:
+                        items.append(_text_from_obj(json.loads(s)))
+                        continue
+                    except json.JSONDecodeError:
+                        pass
+                items.append(s)
+    return items
+
+
+# --------------------------------------------------------------------------- #
+# Self-test                                                                   #
+# --------------------------------------------------------------------------- #
+def run_self_test(threshold: float | None = None, force_stub: bool = False) -> int:
+    """Hand-crafted fixture proving the SEMANTIC catch an 8-gram gate misses.
+
+    If `threshold` is None the test AUTO-CALIBRATES a drop threshold to the loaded
+    embedder (the point midway between the highest clean similarity and the lowest
+    leak similarity) and reports it. This is deliberate: with all-MiniLM a genuine
+    paraphrase scores well BELOW the module's 0.90 default, so the self-test proves
+    *separability* (a threshold exists that catches the paraphrase-leak while keeping
+    clean QA) and surfaces the embedder-specific calibration number instead of
+    faking a pass at 0.90.
+    """
+    test_items = ["What gas do plants absorb from the air?"]
+    fixtures = [
+        # (label, text, expected_dropped)
+        (
+            "A paraphrase (LOW lexical overlap -- would PASS an 8-gram gate)",
+            "Which gas is taken in by vegetation during photosynthesis?",
+            True,
+        ),
+        (
+            "B unrelated clean science-QA",
+            "What is the boiling point of water at sea level?",
+            False,
+        ),
+        (
+            "C near-verbatim",
+            "What gas do plants absorb from the air?",
+            True,
+        ),
+    ]
+    embedder, is_stub = load_embedder(force_stub=force_stub)
+    if is_stub:
+        print("[self-test] *** RUNNING ON TF-IDF STUB -- results are NOT a valid semantic "
+              "PASS; install sentence-transformers for a real run. ***")
+
+    cands = [f[1] for f in fixtures]
+    # threshold=-1 -> nothing auto-dropped; we only want the measured max_sim values.
+    res = semantic_contaminated(cands, test_items, threshold=-1.0, embedder=embedder)
+    sims = [r["max_sim"] for r in res]
+
+    leak_sims = [s for (_, _, exp), s in zip(fixtures, sims) if exp]      # A, C
+    clean_sims = [s for (_, _, exp), s in zip(fixtures, sims) if not exp]  # B
+
+    # Auto-calibrate: midpoint of the clean/leak gap, rounded to 2 dp toward drop side.
+    calibrated = None
+    if threshold is None:
+        gap_lo, gap_hi = max(clean_sims), min(leak_sims)
+        if gap_hi > gap_lo:
+            calibrated = round((gap_lo + gap_hi) / 2.0, 2)
+            thr = calibrated
+        else:
+            thr = DEFAULT_THRESHOLD  # no gap -> honest failure at module default
+    else:
+        thr = threshold
+
+    print(f"[self-test] embedder = {embedder.kind} ({embedder.model_name})")
+    print(f"[self-test] measured max-cosine: leak(A,C)={leak_sims} clean(B)={clean_sims}")
+    if calibrated is not None:
+        print(f"[self-test] auto-calibrated drop threshold for THIS embedder = {thr} "
+              f"(module default is {DEFAULT_THRESHOLD}; see note on embedder-specific tuning)")
+    else:
+        print(f"[self-test] using threshold = {thr}")
+
+    def word_8grams(s: str) -> set:
+        w = s.lower().replace("?", " ").split()
+        return {tuple(w[i : i + 8]) for i in range(len(w) - 7)} if len(w) >= 8 else set()
+
+    test_8g = word_8grams(test_items[0])
+
+    all_ok = True
+    for (label, text, expected), s in zip(fixtures, sims):
+        got = s >= thr
+        ok = got == expected
+        all_ok &= ok
+        shared8 = bool(word_8grams(text) & test_8g)
+        verdict = "PASS" if ok else "FAIL"
+        act = "DROPPED" if got else "KEPT"
+        exp = "DROPPED" if expected else "KEPT"
+        print(f"  [{verdict}] {label}")
+        print(f"         text             : {text!r}")
+        print(f"         max_sim          : {s:.4f}")
+        print(f"         expected/actual  : {exp} / {act}")
+        print(f"         shares 8-gram?   : {shared8}  "
+              f"(an 8-gram gate would {'DROP' if shared8 else 'KEEP'} this)")
+
+    print()
+    if all_ok:
+        print("SELF-TEST: PASS (all 3 expectations met)")
+        print("KEY DEMONSTRATION: candidate A is a paraphrase with NO shared 8-gram "
+              "(a lexical gate KEEPS it) yet semantic decontam DROPS it -> paraphrase-leak closed.")
+    else:
+        print("SELF-TEST: FAIL")
+        if calibrated is None and threshold is None:
+            print("  (no clean/leak gap found with this embedder -- paraphrase not separable "
+                  "from clean; use a stronger embedder such as bge/e5.)")
+    if is_stub and all_ok:
+        print("NOTE: PASS obtained on TF-IDF STUB, not the real embedder -- treat as "
+              "wiring-only, re-run with sentence-transformers installed for a valid PASS.")
+    return 0 if all_ok else 1
+
+
+# --------------------------------------------------------------------------- #
+# CLI                                                                         #
+# --------------------------------------------------------------------------- #
+def main(argv=None) -> int:
+    ap = argparse.ArgumentParser(
+        description="KARD STAGE-2 semantic decontamination (embedding max-cosine). "
+        "Run AFTER the STAGE-1 lexical 8-gram/exact gate; keep a candidate only if "
+        "it passes BOTH stages.",
+    )
+    ap.add_argument("--candidates", help="candidate QA file (.jsonl or .txt)")
+    ap.add_argument("--test", help="test items file, e.g. ARC-Easy test split "
+                    "(.jsonl or .txt); production: real ARC-Easy 2376-question test set")
+    ap.add_argument("--out", help="output drop-report path (.jsonl)")
+    ap.add_argument("--threshold", type=float, default=None,
+                    help=f"cosine drop threshold (module default {DEFAULT_THRESHOLD}; "
+                    "range 0.85-0.92, higher = stricter drop; for --self-test, omit to "
+                    "auto-calibrate to the loaded embedder)")
+    ap.add_argument("--batch-size", type=int, default=DEFAULT_BATCH)
+    ap.add_argument("--model", default=DEFAULT_MODEL, help="sentence-transformers model name")
+    ap.add_argument("--stub", action="store_true",
+                    help="force the TF-IDF stub embedder (debugging only, NOT production)")
+    ap.add_argument("--self-test", action="store_true",
+                    help="run the built-in paraphrase-catch fixture and exit")
+    args = ap.parse_args(argv)
+
+    if args.self_test:
+        return run_self_test(threshold=args.threshold, force_stub=args.stub)
+
+    threshold = DEFAULT_THRESHOLD if args.threshold is None else args.threshold
+
+    if not (args.candidates and args.test and args.out):
+        ap.error("--candidates, --test and --out are required (or use --self-test)")
+
+    candidates = load_items(args.candidates)
+    test_items = load_items(args.test)
+    print(f"loaded {len(candidates)} candidates, {len(test_items)} test items")
+    if len(test_items) < 2000:
+        print(f"note: only {len(test_items)} test items loaded; production ARC-Easy "
+              "test set is 2376 questions -- verify you passed the full split.")
+
+    embedder, is_stub = load_embedder(model_name=args.model, force_stub=args.stub)
+    print(f"embedder: {embedder.kind} ({embedder.model_name}); threshold={threshold}")
+
+    results = semantic_contaminated(
+        candidates, test_items,
+        threshold=threshold, batch_size=args.batch_size, embedder=embedder,
+    )
+
+    dropped = sum(1 for r in results if r["dropped"])
+    kept = len(results) - dropped
+    out = Path(args.out)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    with open(out, "w", encoding="utf-8") as f:
+        for r in results:
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")
+
+    print(f"STAGE-2 semantic decontam: kept={kept} dropped={dropped} "
+          f"(threshold={threshold}); report -> {out}")
+    if is_stub:
+        print("!!! report produced on TF-IDF STUB, NOT the real embedder -- rerun with "
+              "sentence-transformers for a production-valid report.")
+    print("REMINDER: this is STAGE 2. Feed it the survivors of STAGE-1 lexical decontam; "
+          "final training set = candidates that pass BOTH stages.")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())