Merge nucleic/sleek-ember-seal-uady into dev

prepare_data.py

@@ -13,10 +13,12 @@ from typing import Any, Sequence
 from purpose_data import (
     HARD_SLICES,
     LABELS,
+    CurationResult,
     DataError,
     SourceRecord,
     curate_records,
     distribution,
+    exclude_reviewed_duplicates,
     file_sha256,
     jsonl_bytes,
     load_classifiable_fixtures,
@@ -42,6 +44,7 @@ DEFAULT_FIXTURES = (
 DEFAULT_OUTPUT_DIR = SCRIPT_DIR / ".artifacts" / "dataset-v1"
 DEFAULT_FROZEN_TEST = DATA_DIR / "frozen-test-v1.jsonl"
 DEFAULT_SPLIT_MANIFEST = DATA_DIR / "dataset-v1-manifest.json"
+DEFAULT_CURATION_REVIEW = DATA_DIR / "curation-review-v1.json"
 DATASET_VERSION = "purpose-dataset-v1"
 DEFAULT_SEED = 0xC1A551F1
 
@@ -85,6 +88,21 @@ def _sha256_bytes(value: bytes) -> str:
     return hashlib.sha256(value).hexdigest()
 
 
+def load_reviewed_semantic_exclusions(path: Path) -> list[dict[str, Any]]:
+    try:
+        value = json.loads(path.read_text(encoding="utf-8"))
+        review = value["semanticDuplicateReview"]
+        decisions = review["excluded"]
+        status = review["status"]
+    except (OSError, UnicodeError, json.JSONDecodeError, KeyError, TypeError) as exc:
+        raise DataError(f"{path}: cannot read semantic exclusion review: {exc}") from exc
+    if status != "complete":
+        raise DataError(f"{path}: semantic duplicate review is not complete")
+    if not isinstance(decisions, list):
+        raise DataError(f"{path}: semantic excluded must be an array")
+    return decisions
+
+
 def _records(values: Sequence[SourceRecord]) -> list[dict[str, Any]]:
     return [record.value for record in values]
 
@@ -102,8 +120,9 @@ def _build_manifest(
     output_hashes: dict[str, str],
     seed: int,
     near_duplicate_threshold: float,
+    curation_review_path: Path | None,
 ) -> dict[str, Any]:
-    return {
+    manifest = {
         "schemaVersion": 1,
         "datasetVersion": DATASET_VERSION,
         "seed": seed,
@@ -159,6 +178,10 @@ def _build_manifest(
             },
         },
     }
+    if curation_review_path is not None:
+        manifest["curation"]["reviewPath"] = _relative(curation_review_path)
+        manifest["curation"]["reviewSha256"] = file_sha256(curation_review_path)
+    return manifest
 
 
 def prepare(
@@ -171,6 +194,7 @@ def prepare(
     refresh_frozen_test: bool,
     seed: int,
     near_duplicate_threshold: float,
+    curation_review_path: Path | None = None,
 ) -> dict[str, Any]:
     records = load_sources(sources)
     fixtures = load_classifiable_fixtures(fixtures_path)
@@ -178,11 +202,23 @@ def prepare(
     if classifiable_fixture_count != len(fixtures):
         raise DataError("fixture accounting mismatch")
 
-    curated = curate_records(
+    lexical_curation = curate_records(
         records,
         fixtures,
         near_duplicate_threshold=near_duplicate_threshold,
     )
+    reviewed_curation = (
+        exclude_reviewed_duplicates(
+            lexical_curation.records,
+            load_reviewed_semantic_exclusions(curation_review_path),
+        )
+        if curation_review_path is not None
+        else CurationResult(records=lexical_curation.records, duplicates=[])
+    )
+    curated = CurationResult(
+        records=reviewed_curation.records,
+        duplicates=lexical_curation.duplicates + reviewed_curation.duplicates,
+    )
     splits = split_records(
         curated.records,
         fixture_count=classifiable_fixture_count,
@@ -246,6 +282,7 @@ def prepare(
         output_hashes=output_hashes,
         seed=seed,
         near_duplicate_threshold=near_duplicate_threshold,
+        curation_review_path=curation_review_path,
     )
     # The frozen path can be overridden in tests or experiments.
     manifest["frozenEval"]["syntheticPath"] = _relative(frozen_test_path)
@@ -285,6 +322,12 @@ def build_parser() -> argparse.ArgumentParser:
     )
     parser.add_argument("--seed", type=int, default=DEFAULT_SEED)
     parser.add_argument("--near-duplicate-threshold", type=float, default=0.92)
+    parser.add_argument(
+        "--curation-review",
+        type=Path,
+        default=DEFAULT_CURATION_REVIEW,
+        help="completed semantic duplicate review applied after lexical deduplication",
+    )
     return parser
 
 
@@ -305,6 +348,7 @@ def main(argv: Sequence[str] | None = None) -> int:
             refresh_frozen_test=args.refresh_frozen_test,
             seed=args.seed,
             near_duplicate_threshold=args.near_duplicate_threshold,
+            curation_review_path=args.curation_review.resolve(),
         )
     except (DataError, OSError, UnicodeError, json.JSONDecodeError) as exc:
         print(f"error: {exc}", file=sys.stderr)