Add trainer/validate_dataset.py

trainer/validate_dataset.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+import json
+import re
+from collections import Counter, defaultdict
+from pathlib import Path
+from typing import Any
+
+from .common import DATASET_ROOT, atomic_json, ensure_dirs, log_event, update_status
+
+
+ALLOWED_ACTIONS = {
+    "CONTINUE_MONITORING",
+    "RECHECK_SENSOR",
+    "HOLD_ASCENT",
+    "ACTIVATE_EMERGENCY",
+    "REQUEST_HUMAN_REVIEW",
+}
+SENSITIVE_PATTERNS = [
+    re.compile(r"\b\d{3}-\d{2}-\d{4}\b"),
+    re.compile(r"\b(?:\d[ -]*?){13,16}\b"),
+    re.compile(r"\b[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}\b", re.I),
+]
+
+
+def rows(path: Path) -> list[dict[str, Any]]:
+    result = []
+    with path.open("r", encoding="utf-8") as handle:
+        for line_no, line in enumerate(handle, 1):
+            try:
+                result.append(json.loads(line))
+            except json.JSONDecodeError as exc:
+                raise ValueError(f"{path.name}:{line_no}: invalid JSON: {exc}") from exc
+    return result
+
+
+def validate() -> dict[str, Any]:
+    errors: list[str] = []
+    warnings: list[str] = []
+    ids: set[str] = set()
+    content_hashes: set[str] = set()
+    group_splits: defaultdict[str, set[str]] = defaultdict(set)
+    counts: Counter[str] = Counter()
+    categories: Counter[str] = Counter()
+    actions: Counter[str] = Counter()
+
+    for split in ("train", "validation", "test"):
+        path = DATASET_ROOT / f"{split}.jsonl"
+        if not path.exists():
+            errors.append(f"Missing {path.name}")
+            continue
+        for row in rows(path):
+            counts[split] += 1
+            if row.get("split") != split:
+                errors.append(f"{row.get('id')}: split field does not match {path.name}")
+            if row.get("id") in ids:
+                errors.append(f"Duplicate id {row.get('id')}")
+            ids.add(str(row.get("id")))
+            group_splits[str(row.get("group"))].add(split)
+            categories[str(row.get("category"))] += 1
+
+            messages = row.get("messages")
+            if not isinstance(messages, list) or len(messages) < 3:
+                errors.append(f"{row.get('id')}: messages must contain system, user, assistant")
+                continue
+            roles = [message.get("role") for message in messages]
+            if roles[0] != "system" or roles[-1] != "assistant" or roles[-2] != "user":
+                errors.append(f"{row.get('id')}: invalid role sequence {roles}")
+            if any(not isinstance(message.get("content"), str) or not message["content"].strip() for message in messages):
+                errors.append(f"{row.get('id')}: empty or non-string message")
+
+            body = "\n".join(message["content"] for message in messages)
+            digest = __import__("hashlib").sha256(body.encode("utf-8")).hexdigest()
+            if digest in content_hashes:
+                errors.append(f"{row.get('id')}: exact content duplicate")
+            content_hashes.add(digest)
+            if any(pattern.search(body) for pattern in SENSITIVE_PATTERNS):
+                errors.append(f"{row.get('id')}: possible sensitive identifier")
+
+            try:
+                output = json.loads(messages[-1]["content"])
+            except json.JSONDecodeError:
+                errors.append(f"{row.get('id')}: assistant output is not JSON")
+                continue
+            if set(output) != {"speech", "next_action", "question", "uncertainties"}:
+                errors.append(f"{row.get('id')}: assistant schema mismatch")
+            action = output.get("next_action")
+            actions[str(action)] += 1
+            if action not in ALLOWED_ACTIONS:
+                errors.append(f"{row.get('id')}: unsupported action {action}")
+            if not isinstance(output.get("uncertainties"), list):
+                errors.append(f"{row.get('id')}: uncertainties must be a list")
+            if row.get("source") == "med-safety-bench" and split == "train":
+                errors.append("MedSafetyBench must not enter the SFT train split")
+
+    leaks = {group: sorted(splits) for group, splits in group_splits.items() if len(splits) > 1}
+    if leaks:
+        errors.append(f"Group leakage across splits: {list(leaks.items())[:10]}")
+    if counts["train"] < 1000:
+        errors.append(f"Train split is too small: {counts['train']}")
+    if counts["validation"] < 100 or counts["test"] < 100:
+        errors.append(f"Validation/test split too small: {dict(counts)}")
+    if actions["ACTIVATE_EMERGENCY"] < 100:
+        warnings.append("Emergency examples are below 100")
+
+    report = {
+        "ok": not errors,
+        "errors": errors,
+        "warnings": warnings,
+        "counts": dict(counts),
+        "categories": dict(categories),
+        "actions": dict(actions),
+        "unique_groups": len(group_splits),
+    }
+    return report
+
+
+def main() -> None:
+    ensure_dirs()
+    update_status("validate_dataset", "Validating schema, splits, duplication, and safety contract")
+    report = validate()
+    atomic_json(DATASET_ROOT / "validation_report.json", report)
+    log_event("dataset_validated", **report)
+    if not report["ok"]:
+        raise RuntimeError("Dataset validation failed: " + "; ".join(report["errors"][:8]))
+    update_status("dataset_valid", "Dataset validation passed", validation=report)
+
+
+if __name__ == "__main__":
+    main()
+