324 lines
11 KiB
Python
324 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Build a training-only dataset augmented with labeled first prompts."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import shutil
|
|
import sys
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
from typing import Any, Sequence
|
|
|
|
from purpose_data import (
|
|
DataError,
|
|
SourceRecord,
|
|
curate_records,
|
|
distribution,
|
|
file_sha256,
|
|
jsonl_bytes,
|
|
load_classifiable_fixtures,
|
|
load_jsonl,
|
|
normalized_key,
|
|
prompt_hash,
|
|
validate_source_record,
|
|
write_json,
|
|
write_jsonl,
|
|
)
|
|
|
|
|
|
SCRIPT_DIR = Path(__file__).resolve().parent
|
|
REPOSITORY_ROOT = SCRIPT_DIR.parent.parent
|
|
DEFAULT_BASE_DATASET = SCRIPT_DIR / ".artifacts" / "dataset-v1"
|
|
DEFAULT_HISTORY = (
|
|
SCRIPT_DIR / ".artifacts" / "nucleic-history-first-prompts.labeled.jsonl"
|
|
)
|
|
DEFAULT_FIXTURES = (
|
|
REPOSITORY_ROOT
|
|
/ "Tests"
|
|
/ "NucleicCoreTests"
|
|
/ "Fixtures"
|
|
/ "purpose-prompts.json"
|
|
)
|
|
DEFAULT_OUTPUT = SCRIPT_DIR / ".artifacts" / "dataset-v1-history-first-prompts"
|
|
EXPERIMENT_VERSION = "history-first-prompts-training-augmentation-v1"
|
|
|
|
|
|
def _relative(path: Path) -> str:
|
|
try:
|
|
return str(path.resolve().relative_to(REPOSITORY_ROOT))
|
|
except ValueError:
|
|
return str(path.resolve())
|
|
|
|
|
|
def _validate_records(records: Sequence[dict[str, Any]], path: Path) -> None:
|
|
if not records:
|
|
raise DataError(f"{path}: split is empty")
|
|
for line, record in enumerate(records, 1):
|
|
validate_source_record(record, f"{path}:{line}")
|
|
|
|
|
|
def _indexed_records(
|
|
records: Sequence[dict[str, Any]], path: Path
|
|
) -> list[SourceRecord]:
|
|
return [
|
|
SourceRecord(value=record, source=path, line=line)
|
|
for line, record in enumerate(records, 1)
|
|
]
|
|
|
|
|
|
def _unique_split_keys(
|
|
splits: Sequence[tuple[str, Sequence[dict[str, Any]]]]
|
|
) -> dict[str, tuple[str, str]]:
|
|
seen: dict[str, tuple[str, str]] = {}
|
|
for split_name, records in splits:
|
|
for line, record in enumerate(records, 1):
|
|
key = normalized_key(record["prompt"])
|
|
previous = seen.get(key)
|
|
if previous is not None:
|
|
previous_split, previous_label = previous
|
|
detail = (
|
|
"conflicting labels"
|
|
if previous_label != record["purpose"]
|
|
else "duplicate prompt"
|
|
)
|
|
raise DataError(
|
|
f"{split_name}:{line}: {detail} also present in {previous_split}"
|
|
)
|
|
seen[key] = (split_name, record["purpose"])
|
|
return seen
|
|
|
|
|
|
def _sha256_bytes(value: bytes) -> str:
|
|
return hashlib.sha256(value).hexdigest()
|
|
|
|
|
|
def _exclusion(
|
|
record: SourceRecord,
|
|
*,
|
|
reason: str,
|
|
matched_prompt_hash: str | None = None,
|
|
similarity: float | None = None,
|
|
) -> dict[str, Any]:
|
|
value: dict[str, Any] = {
|
|
"promptHash": prompt_hash(record.value["prompt"]),
|
|
"sourceLine": record.line,
|
|
"reason": reason,
|
|
}
|
|
if matched_prompt_hash is not None:
|
|
value["matchedPromptHash"] = matched_prompt_hash
|
|
if similarity is not None:
|
|
value["similarity"] = round(similarity, 6)
|
|
return value
|
|
|
|
|
|
def prepare(
|
|
*,
|
|
base_dataset: Path,
|
|
history_path: Path,
|
|
fixtures_path: Path,
|
|
output_dir: Path,
|
|
near_duplicate_threshold: float,
|
|
overwrite_output: bool,
|
|
) -> dict[str, Any]:
|
|
base_paths = {
|
|
split: base_dataset / f"{split}.jsonl"
|
|
for split in ("train", "validation", "test")
|
|
}
|
|
base = {split: load_jsonl(path) for split, path in base_paths.items()}
|
|
for split, path in base_paths.items():
|
|
_validate_records(base[split], path)
|
|
|
|
base_keys = _unique_split_keys(
|
|
[(split, base[split]) for split in ("train", "validation", "test")]
|
|
)
|
|
history = load_jsonl(history_path)
|
|
_validate_records(history, history_path)
|
|
fixtures = load_classifiable_fixtures(fixtures_path)
|
|
|
|
eval_fixtures = [
|
|
{"prompt": record["prompt"], "purpose": record["purpose"]}
|
|
for split in ("validation", "test")
|
|
for record in base[split]
|
|
] + fixtures
|
|
eval_hashes = {prompt_hash(record["prompt"]) for record in eval_fixtures}
|
|
|
|
eligible: list[SourceRecord] = []
|
|
exclusions: list[dict[str, Any]] = []
|
|
for record in _indexed_records(history, history_path):
|
|
if record.value["slice"] == "vague-eval":
|
|
exclusions.append(_exclusion(record, reason="vague-eval"))
|
|
continue
|
|
|
|
key = normalized_key(record.value["prompt"])
|
|
previous = base_keys.get(key)
|
|
if previous is not None:
|
|
split_name, previous_label = previous
|
|
if previous_label != record.value["purpose"]:
|
|
raise DataError(
|
|
f"{history_path}:{record.line}: label conflicts with exact "
|
|
f"{split_name} prompt"
|
|
)
|
|
exclusions.append(
|
|
_exclusion(
|
|
record,
|
|
reason=(
|
|
"exact-base-train-overlap"
|
|
if split_name == "train"
|
|
else "exact-eval-overlap"
|
|
),
|
|
matched_prompt_hash=prompt_hash(record.value["prompt"]),
|
|
similarity=1.0,
|
|
)
|
|
)
|
|
continue
|
|
eligible.append(record)
|
|
|
|
curated = curate_records(
|
|
eligible,
|
|
eval_fixtures,
|
|
near_duplicate_threshold=near_duplicate_threshold,
|
|
)
|
|
for duplicate in curated.duplicates:
|
|
exclusions.append(
|
|
_exclusion(
|
|
duplicate.dropped,
|
|
reason=(
|
|
f"{duplicate.kind}-eval-overlap"
|
|
if duplicate.matched_prompt_hash in eval_hashes
|
|
else f"{duplicate.kind}-history-duplicate"
|
|
),
|
|
matched_prompt_hash=duplicate.matched_prompt_hash,
|
|
similarity=duplicate.similarity,
|
|
)
|
|
)
|
|
|
|
accepted_history = [record.value for record in curated.records]
|
|
output_train = base["train"] + accepted_history
|
|
train_bytes = jsonl_bytes(output_train)
|
|
validation_bytes = base_paths["validation"].read_bytes()
|
|
test_bytes = base_paths["test"].read_bytes()
|
|
|
|
if output_dir.exists() and any(output_dir.iterdir()):
|
|
if not overwrite_output:
|
|
raise DataError(
|
|
f"{output_dir}: output is not empty; pass --overwrite-output "
|
|
"intentionally"
|
|
)
|
|
shutil.rmtree(output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
(output_dir / "train.jsonl").write_bytes(train_bytes)
|
|
(output_dir / "validation.jsonl").write_bytes(validation_bytes)
|
|
(output_dir / "test.jsonl").write_bytes(test_bytes)
|
|
write_jsonl(
|
|
output_dir / "history-exclusions.jsonl",
|
|
sorted(exclusions, key=lambda value: value["sourceLine"]),
|
|
)
|
|
|
|
exclusion_counts = dict(
|
|
sorted(Counter(value["reason"] for value in exclusions).items())
|
|
)
|
|
accepted_sources = _indexed_records(accepted_history, history_path)
|
|
manifest = {
|
|
"schemaVersion": 1,
|
|
"experimentVersion": EXPERIMENT_VERSION,
|
|
"policy": {
|
|
"historyUsage": "training-only",
|
|
"vagueEval": "excluded from optimization",
|
|
"exactBaseTrainOverlap": "excluded",
|
|
"exactEvaluationOverlap": "excluded",
|
|
"nearEvaluationOverlap": "excluded",
|
|
"nearHistoryDuplicates": "excluded",
|
|
"nearDuplicateThreshold": near_duplicate_threshold,
|
|
"validationAndTestBytes": "identical to base dataset",
|
|
},
|
|
"sources": {
|
|
"baseDataset": {
|
|
"path": _relative(base_dataset),
|
|
"splits": {
|
|
split: {
|
|
"path": _relative(path),
|
|
"records": len(base[split]),
|
|
"sha256": file_sha256(path),
|
|
}
|
|
for split, path in base_paths.items()
|
|
},
|
|
},
|
|
"history": {
|
|
"path": _relative(history_path),
|
|
"records": len(history),
|
|
"sha256": file_sha256(history_path),
|
|
},
|
|
"fixtures": {
|
|
"path": _relative(fixtures_path),
|
|
"classifiableRecords": len(fixtures),
|
|
"sha256": file_sha256(fixtures_path),
|
|
},
|
|
},
|
|
"augmentation": {
|
|
"inputHistoryRecords": len(history),
|
|
"acceptedHistoryRecords": len(accepted_history),
|
|
"excludedHistoryRecords": len(exclusions),
|
|
"exclusions": exclusion_counts,
|
|
"acceptedDistribution": distribution(accepted_sources),
|
|
},
|
|
"outputs": {
|
|
"train": {
|
|
"records": len(output_train),
|
|
"sha256": _sha256_bytes(train_bytes),
|
|
},
|
|
"validation": {
|
|
"records": len(base["validation"]),
|
|
"sha256": _sha256_bytes(validation_bytes),
|
|
},
|
|
"test": {
|
|
"records": len(base["test"]),
|
|
"sha256": _sha256_bytes(test_bytes),
|
|
},
|
|
},
|
|
}
|
|
write_json(output_dir / "manifest.json", manifest)
|
|
return manifest
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--base-dataset", type=Path, default=DEFAULT_BASE_DATASET)
|
|
parser.add_argument("--history", type=Path, default=DEFAULT_HISTORY)
|
|
parser.add_argument("--fixtures", type=Path, default=DEFAULT_FIXTURES)
|
|
parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT)
|
|
parser.add_argument("--near-duplicate-threshold", type=float, default=0.92)
|
|
parser.add_argument("--overwrite-output", action="store_true")
|
|
return parser
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
|
args = build_parser().parse_args(argv)
|
|
try:
|
|
manifest = prepare(
|
|
base_dataset=args.base_dataset.resolve(),
|
|
history_path=args.history.resolve(),
|
|
fixtures_path=args.fixtures.resolve(),
|
|
output_dir=args.output_dir.resolve(),
|
|
near_duplicate_threshold=args.near_duplicate_threshold,
|
|
overwrite_output=args.overwrite_output,
|
|
)
|
|
except (DataError, OSError, UnicodeError, json.JSONDecodeError) as exc:
|
|
print(f"error: {exc}", file=sys.stderr)
|
|
return 1
|
|
|
|
augmentation = manifest["augmentation"]
|
|
outputs = manifest["outputs"]
|
|
print(
|
|
f"Prepared {outputs['train']['records']} training records "
|
|
f"({augmentation['acceptedHistoryRecords']} from history); validation and "
|
|
f"test remain byte-identical to the base dataset."
|
|
)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|