| |
| """Score an RLCD checkpoint on a labelled split, the way the published model was judged. |
| |
| python evaluate.py --model rlcd-mine --data data/test.jsonl --latency 200 \ |
| --prior public=0.146,internal=0.527,confidential=0.297,restricted=0.024,sovereign-only=0.006 |
| |
| Reports, per head: accuracy, macro-F1, expected calibration error, and coverage / accuracy at confidence |
| thresholds (to choose your own). For the ordered sensitivity head, under the tighten-only rule (the highest |
| level L with P(class >= L) >= --tighten): how many sensitive requests are classified too low, overall and per |
| level, and how many public or internal ones are raised, with 95% Wilson intervals. If rows carry a "dlp" |
| level (your DLP's verdict), it also reports the rate with DLP alongside. `--prior` re-weights rows to your |
| production sensitivity mix: weighted figures are reported beside the as-sampled ones, never instead of them. |
| Macro-F1 averages over every class of a head, as the published figures do (a class absent from the split |
| scores 0); `--macro-f1-over present` averages over the classes present instead, for small or partial splits. |
| `--latency N` times N single requests on this machine, tokenizer included. |
| `--model` is a folder or a Hub repo; `--revision` pins a Hub revision. |
| """ |
|
|
| import argparse |
| import itertools |
| import json |
| import math |
| import time |
| from pathlib import Path |
|
|
| import torch |
| from huggingface_hub import snapshot_download |
|
|
| from rlcd import RlcdClassifier, render |
|
|
|
|
| def wilson(hits: float, n: float) -> tuple[float, float]: |
| if n == 0: |
| return (0.0, 1.0) |
| p, z = hits / n, 1.96 |
| centre, half = p + z * z / (2 * n), z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) |
| return ((centre - half) / (1 + z * z / n), (centre + half) / (1 + z * z / n)) |
|
|
|
|
| def tighten(distribution: dict[str, float], order: list[str], threshold: float) -> str: |
| mass = 0.0 |
| for level in reversed(order): |
| mass += distribution[level] |
| if mass >= threshold: |
| return level |
| return order[0] |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| parser.add_argument("--model", required=True) |
| parser.add_argument("--revision", default=None) |
| parser.add_argument("--data", type=Path, required=True) |
| parser.add_argument("--prior", default=None, help="level=share,... to weight rows to production") |
| parser.add_argument("--tighten", type=float, default=0.5) |
| parser.add_argument("--sensitive-from", default="confidential", help="the first level that counts as sensitive") |
| parser.add_argument("--thresholds", default="0.5,0.6,0.7,0.8,0.9,0.95") |
| parser.add_argument("--macro-f1-over", choices=["all", "present"], default="all") |
| parser.add_argument("--latency", type=int, default=0) |
| parser.add_argument("--out", type=Path, default=None, help="write the report as JSON") |
| args = parser.parse_args() |
|
|
| path = Path(args.model) if Path(args.model).is_dir() else Path(snapshot_download(args.model, revision=args.revision)) |
| model = RlcdClassifier(path) |
| heads, temperatures = model.config["heads"], model.config["temperatures"] |
| rows = [json.loads(line) for line in args.data.read_text().splitlines() if line.strip()] |
| texts = [r.get("text") or render(r["messages"]) for r in rows] |
|
|
| probs: dict[str, list[torch.Tensor]] = {h: [] for h in heads} |
| with torch.no_grad(): |
| for i in range(0, len(texts), 16): |
| logits = model.logits(texts[i : i + 16]) |
| for h in heads: |
| probs[h].append((logits[h] / temperatures[h]).softmax(-1)) |
| probs = {h: torch.cat(p) for h, p in probs.items()} |
|
|
| weights = torch.ones(len(rows)) |
| if args.prior: |
| prior = {k: float(v) for k, v in (item.split("=") for item in args.prior.split(","))} |
| levels = [r.get("sensitivity") for r in rows] |
| share = {lvl: levels.count(lvl) / len(levels) for lvl in set(levels)} |
| weights = torch.tensor([prior.get(lvl, 0.0) / share[lvl] for lvl in levels]) |
|
|
| report: dict = {"rows": len(rows), "provenance": "measured", "heads": {}} |
| thresholds = [float(t) for t in args.thresholds.split(",")] |
|
|
| def metrics(p: torch.Tensor, labels: torch.Tensor, w: torch.Tensor, classes: int) -> dict: |
| confidence, prediction = p.max(-1) |
| correct = prediction.eq(labels).float() |
| over = range(classes) if args.macro_f1_over == "all" else sorted(set(labels.tolist()) | set(prediction.tolist())) |
| f1s = [] |
| for c in over: |
| tp = (w * ((prediction == c) & (labels == c))).sum() |
| fp, fn = (w * ((prediction == c) & (labels != c))).sum(), (w * ((prediction != c) & (labels == c))).sum() |
| f1s.append(0.0 if tp == 0 else (2 * tp / (2 * tp + fp + fn)).item()) |
| ece = 0.0 |
| for low, high in itertools.pairwise(torch.linspace(0, 1, 16)): |
| in_bin = (confidence > low) & (confidence <= high) |
| if in_bin.any(): |
| bw = w[in_bin] |
| gap = ((confidence[in_bin] * bw).sum() - (correct[in_bin] * bw).sum()) / bw.sum() |
| ece += (bw.sum() / w.sum()).item() * abs(gap.item()) |
| curve = {} |
| for t in thresholds: |
| kept = confidence >= t |
| accuracy = ((correct[kept] * w[kept]).sum() / w[kept].sum()).item() if kept.any() else None |
| curve[str(t)] = {"coverage": round((w[kept].sum() / w.sum()).item(), 4), |
| "accuracy": None if accuracy is None else round(accuracy, 4)} |
| return {"accuracy": round(((correct * w).sum() / w.sum()).item(), 4), |
| "macro_f1": round(sum(f1s) / max(1, len(f1s)), 4), "ece": round(ece, 4), "coverage_accuracy": curve} |
|
|
| for h, classes in heads.items(): |
| keep = [i for i, r in enumerate(rows) if r.get(h) in classes] |
| if not keep: |
| continue |
| labels = torch.tensor([classes.index(rows[i][h]) for i in keep]) |
| report["heads"][h] = {"rows": len(keep), **metrics(probs[h][keep], labels, torch.ones(len(keep)), len(classes))} |
| if args.prior: |
| report["heads"][h]["under_prior"] = metrics(probs[h][keep], labels, weights[keep], len(classes)) |
|
|
| if "sensitivity" in heads: |
| order = heads["sensitivity"] |
| start = order.index(args.sensitive_from) |
| rank = order.index |
| safety: dict = {"tighten_at": args.tighten} |
| variants = {"rlcd": lambda i, lvl: lvl} |
| if any("dlp" in r for r in rows): |
| variants["rlcd_or_dlp"] = lambda i, lvl: max(lvl, rows[i].get("dlp") or order[0], key=rank) |
| for name, combine in variants.items(): |
| effective = [combine(i, tighten(dict(zip(order, probs["sensitivity"][i].tolist(), strict=True)), order, args.tighten)) |
| for i in range(len(rows))] |
| sensitive = [i for i, r in enumerate(rows) if r.get("sensitivity") in order[start:]] |
| under = [i for i in sensitive if rank(effective[i]) < rank(rows[i]["sensitivity"])] |
| total, hits = sum(weights[i].item() for i in sensitive), sum(weights[i].item() for i in under) |
| entry = {"under_classified": {"rate": round(len(under) / len(sensitive), 4) if sensitive else None, |
| "of": len(sensitive), "ci95": [round(x, 4) for x in wilson(len(under), len(sensitive))], |
| **({"rate_under_prior": round(hits / total, 4) if total else None} if args.prior else {})}} |
| for level in order[start:]: |
| at = [i for i in sensitive if rows[i]["sensitivity"] == level] |
| n_under = sum(1 for i in at if rank(effective[i]) < rank(level)) |
| entry[f"under_{level}"] = {"rate": round(n_under / len(at), 4) if at else None, "of": len(at), |
| "ci95": [round(x, 4) for x in wilson(n_under, len(at))]} |
| plain = [i for i, r in enumerate(rows) if r.get("sensitivity") in order[:start]] |
| for level in order[start:]: |
| raised = sum(1 for i in plain if rank(effective[i]) >= rank(level)) |
| entry[f"raised_to_{level}_or_above"] = {"rate": round(raised / len(plain), 4) if plain else None, |
| "of": len(plain)} |
| safety[name] = entry |
| report["sensitivity_safety"] = safety |
|
|
| if args.latency: |
| sample = [rows[i % len(rows)] for i in range(args.latency)] |
| timings = [] |
| for row in sample: |
| started = time.perf_counter() |
| model.predict(row.get("messages") or [{"role": "user", "content": row["text"]}]) |
| timings.append((time.perf_counter() - started) * 1000) |
| timings.sort() |
| report["latency_ms"] = {"p50": round(timings[len(timings) // 2], 1), |
| "p95": round(timings[min(len(timings) - 1, int(0.95 * len(timings)))], 1), |
| "requests": len(timings), "threads": torch.get_num_threads()} |
| print(json.dumps(report, indent=1)) |
| if args.out: |
| args.out.write_text(json.dumps(report, indent=1)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
