Files
nucleic/ml/purpose-classifier/prepare_data.py
T

403 lines
14 KiB
Python

#!/usr/bin/env python3
"""Curate source prompts and build deterministic purpose-classifier splits."""
from __future__ import annotations
import argparse
import hashlib
import json
import sys
from pathlib import Path
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,
load_sources,
prompt_hash,
split_records,
write_json,
write_jsonl,
)
from review_contract import apply_completed_human_review
SCRIPT_DIR = Path(__file__).resolve().parent
REPOSITORY_ROOT = SCRIPT_DIR.parent.parent
DATA_DIR = SCRIPT_DIR / "data"
GENERATION_MANIFEST = DATA_DIR / "generation-manifest.json"
DEFAULT_FIXTURES = (
REPOSITORY_ROOT
/ "Tests"
/ "NucleicCoreTests"
/ "Fixtures"
/ "purpose-prompts.json"
)
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
def _relative(path: Path) -> str:
try:
return str(path.resolve().relative_to(REPOSITORY_ROOT))
except ValueError:
return str(path.resolve())
def default_source_paths() -> list[Path]:
try:
value = json.loads(GENERATION_MANIFEST.read_text(encoding="utf-8"))
names = value["canonicalFiles"]
except (OSError, UnicodeError, json.JSONDecodeError, KeyError, TypeError) as exc:
raise DataError(f"{GENERATION_MANIFEST}: cannot read canonicalFiles: {exc}") from exc
if not isinstance(names, list) or not names or not all(
isinstance(name, str) and name for name in names
):
raise DataError(
f"{GENERATION_MANIFEST}: canonicalFiles must be a non-empty string array"
)
return [(DATA_DIR / name).resolve() for name in names]
def _fixture_counts(path: Path) -> tuple[int, int]:
try:
value = json.loads(path.read_text(encoding="utf-8"))
except (OSError, UnicodeError, json.JSONDecodeError) as exc:
raise DataError(f"{path}: cannot read fixtures: {exc}") from exc
if not isinstance(value, list):
raise DataError(f"{path}: fixture root must be an array")
classifiable = sum(
isinstance(item, dict) and item.get("purpose") in LABELS for item in value
)
return len(value), classifiable
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]
def _build_manifest(
*,
sources: Sequence[Path],
source_record_count: int,
curated_record_count: int,
duplicate_counts: dict[str, int],
splits: Any,
fixtures_path: Path,
total_fixture_count: int,
classifiable_fixture_count: int,
output_hashes: dict[str, str],
seed: int,
near_duplicate_threshold: float,
curation_review_path: Path | None,
human_review_path: Path | None,
human_review_summary: dict[str, Any] | None,
) -> dict[str, Any]:
manifest = {
"schemaVersion": 1,
"datasetVersion": DATASET_VERSION,
"seed": seed,
"ratios": {"train": 0.8, "validation": 0.1, "test": 0.1},
"sources": [
{
"path": _relative(path),
"sha256": file_sha256(path),
"records": len(load_sources([path])),
}
for path in sources
],
"curation": {
"inputRecords": source_record_count,
"retainedRecords": curated_record_count,
"excludedDuplicates": duplicate_counts,
"nearDuplicateMethod": "word-trigram Jaccard after SimHash LSH candidate search",
"nearDuplicateThreshold": near_duplicate_threshold,
"vagueEvalPolicy": "validation/test only",
},
"frozenEval": {
"syntheticPath": _relative(DEFAULT_FROZEN_TEST),
"syntheticSha256": output_hashes["test"],
"shippedFixturesPath": _relative(fixtures_path),
"shippedFixturesSha256": file_sha256(fixtures_path),
"shippedFixtureRecords": total_fixture_count,
"classifiableShippedFixtureRecords": classifiable_fixture_count,
"excludedGeneralFixtureRecords": (
total_fixture_count - classifiable_fixture_count
),
"hardSliceDefinition": sorted(HARD_SLICES),
},
"splits": {
"train": {
"records": len(splits.train),
"sha256": output_hashes["train"],
"distribution": distribution(splits.train),
},
"validation": {
"records": len(splits.validation),
"sha256": output_hashes["validation"],
"distribution": distribution(splits.validation),
},
"test": {
"syntheticRecords": len(splits.test),
"classifiableFixtureRecords": classifiable_fixture_count,
"logicalRecords": splits.logical_test_count,
"hardSyntheticRecords": sum(
record.value["slice"] in HARD_SLICES for record in splits.test
),
"sha256": output_hashes["test"],
"distribution": distribution(splits.test),
},
},
}
if curation_review_path is not None:
manifest["curation"]["reviewPath"] = _relative(curation_review_path)
manifest["curation"]["reviewSha256"] = file_sha256(curation_review_path)
if human_review_path is not None:
manifest["curation"]["humanReview"] = {
"path": _relative(human_review_path),
"sha256": file_sha256(human_review_path),
"summary": human_review_summary,
}
return manifest
def prepare(
*,
sources: Sequence[Path],
fixtures_path: Path,
output_dir: Path,
frozen_test_path: Path,
manifest_path: Path,
refresh_frozen_test: bool,
seed: int,
near_duplicate_threshold: float,
curation_review_path: Path | None = None,
human_review_path: Path | None = None,
) -> dict[str, Any]:
records = load_sources(sources)
fixtures = load_classifiable_fixtures(fixtures_path)
total_fixture_count, classifiable_fixture_count = _fixture_counts(fixtures_path)
if classifiable_fixture_count != len(fixtures):
raise DataError("fixture accounting mismatch")
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=[])
)
human_review_summary = None
reviewed_records = reviewed_curation.records
if human_review_path is not None:
human_review = apply_completed_human_review(
human_review_path,
reviewed_records,
dataset_version=DATASET_VERSION,
)
reviewed_records = human_review.records
human_review_summary = human_review.summary
curated = CurationResult(
records=reviewed_records,
duplicates=lexical_curation.duplicates + reviewed_curation.duplicates,
)
splits = split_records(
curated.records,
fixture_count=classifiable_fixture_count,
seed=seed,
)
train_values = _records(splits.train)
validation_values = _records(splits.validation)
test_values = _records(splits.test)
candidate_frozen_test = jsonl_bytes(test_values)
if refresh_frozen_test:
frozen_test_path.parent.mkdir(parents=True, exist_ok=True)
frozen_test_path.write_bytes(candidate_frozen_test)
elif not frozen_test_path.exists():
raise DataError(
f"{frozen_test_path}: frozen test is missing; review the candidate then run "
"--refresh-frozen-test"
)
elif frozen_test_path.read_bytes() != candidate_frozen_test:
raise DataError(
f"{frozen_test_path}: deterministic test split changed; inspect source/seed "
"changes and use --refresh-frozen-test only when intentionally versioning it"
)
output_dir.mkdir(parents=True, exist_ok=True)
write_jsonl(output_dir / "train.jsonl", train_values)
write_jsonl(output_dir / "validation.jsonl", validation_values)
write_jsonl(output_dir / "test.jsonl", test_values)
write_jsonl(
output_dir / "exclusions.jsonl",
(
{
"promptHash": prompt_hash(duplicate.dropped.value["prompt"]),
"matchedPromptHash": duplicate.matched_prompt_hash,
"source": _relative(duplicate.dropped.source),
"line": duplicate.dropped.line,
"kind": duplicate.kind,
"similarity": round(duplicate.similarity, 6),
}
for duplicate in curated.duplicates
),
)
duplicate_counts: dict[str, int] = {}
for duplicate in curated.duplicates:
duplicate_counts[duplicate.kind] = duplicate_counts.get(duplicate.kind, 0) + 1
output_hashes = {
"train": _sha256_bytes(jsonl_bytes(train_values)),
"validation": _sha256_bytes(jsonl_bytes(validation_values)),
"test": _sha256_bytes(candidate_frozen_test),
}
manifest = _build_manifest(
sources=sources,
source_record_count=len(records),
curated_record_count=len(curated.records),
duplicate_counts=dict(sorted(duplicate_counts.items())),
splits=splits,
fixtures_path=fixtures_path,
total_fixture_count=total_fixture_count,
classifiable_fixture_count=classifiable_fixture_count,
output_hashes=output_hashes,
seed=seed,
near_duplicate_threshold=near_duplicate_threshold,
curation_review_path=curation_review_path,
human_review_path=human_review_path,
human_review_summary=human_review_summary,
)
# The frozen path can be overridden in tests or experiments.
manifest["frozenEval"]["syntheticPath"] = _relative(frozen_test_path)
write_json(output_dir / "manifest.json", manifest)
if refresh_frozen_test:
write_json(manifest_path, manifest)
elif manifest_path.exists():
existing = json.loads(manifest_path.read_text(encoding="utf-8"))
if existing != manifest:
raise DataError(
f"{manifest_path}: split manifest changed; inspect and refresh the frozen "
"test intentionally"
)
else:
raise DataError(
f"{manifest_path}: frozen split manifest is missing; use --refresh-frozen-test"
)
return manifest
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--source",
action="append",
type=Path,
help="canonical source JSONL; repeat for multiple files (default: generation manifest)",
)
parser.add_argument("--fixtures", type=Path, default=DEFAULT_FIXTURES)
parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT_DIR)
parser.add_argument("--frozen-test", type=Path, default=DEFAULT_FROZEN_TEST)
parser.add_argument("--manifest", type=Path, default=DEFAULT_SPLIT_MANIFEST)
parser.add_argument(
"--refresh-frozen-test",
action="store_true",
help="replace the versioned test split and manifest after intentional review",
)
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",
)
parser.add_argument(
"--human-review",
type=Path,
help=(
"completed label-and-difficulty review ledger from review_data.py; "
"omitted until human review is complete"
),
)
return parser
def main(argv: Sequence[str] | None = None) -> int:
args = build_parser().parse_args(argv)
try:
sources = (
[path.resolve() for path in args.source]
if args.source
else default_source_paths()
)
manifest = prepare(
sources=sources,
fixtures_path=args.fixtures.resolve(),
output_dir=args.output_dir.resolve(),
frozen_test_path=args.frozen_test.resolve(),
manifest_path=args.manifest.resolve(),
refresh_frozen_test=args.refresh_frozen_test,
seed=args.seed,
near_duplicate_threshold=args.near_duplicate_threshold,
curation_review_path=args.curation_review.resolve(),
human_review_path=(
args.human_review.resolve() if args.human_review is not None else None
),
)
except (DataError, OSError, UnicodeError, json.JSONDecodeError) as exc:
print(f"error: {exc}", file=sys.stderr)
return 1
splits = manifest["splits"]
print(
f"Prepared {manifest['curation']['retainedRecords']} curated records: "
f"{splits['train']['records']} train, "
f"{splits['validation']['records']} validation, "
f"{splits['test']['logicalRecords']} frozen test "
f"({splits['test']['classifiableFixtureRecords']} shipped fixtures)."
)
return 0
if __name__ == "__main__":
raise SystemExit(main())