438 lines
15 KiB
Python
438 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
"""Audit semantic duplicates and emit the frozen human-review sample."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from collections import Counter
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Sequence
|
|
|
|
import numpy as np
|
|
|
|
from prepare_data import (
|
|
DEFAULT_CURATION_REVIEW,
|
|
DEFAULT_FIXTURES,
|
|
DEFAULT_FROZEN_TEST,
|
|
default_source_paths,
|
|
load_reviewed_semantic_exclusions,
|
|
)
|
|
from purpose_data import (
|
|
LABELS,
|
|
DataError,
|
|
SourceRecord,
|
|
curate_records,
|
|
exclude_reviewed_duplicates,
|
|
load_classifiable_fixtures,
|
|
load_sources,
|
|
prompt_hash,
|
|
write_json,
|
|
)
|
|
from review_contract import (
|
|
DEFAULT_SAMPLE_FRACTION,
|
|
DEFAULT_SAMPLE_SEED,
|
|
stratified_review_sample,
|
|
write_review_csv,
|
|
)
|
|
from train import (
|
|
DEFAULT_MODEL,
|
|
DEFAULT_MODEL_REVISION,
|
|
HEAD_TOKENS,
|
|
MAX_LENGTH,
|
|
TAIL_TOKENS,
|
|
encode_fixed_shape,
|
|
)
|
|
|
|
|
|
SCRIPT_DIR = Path(__file__).resolve().parent
|
|
REPOSITORY_ROOT = SCRIPT_DIR.parent.parent
|
|
DEFAULT_REPORT = SCRIPT_DIR / "data" / "semantic-audit-v1.json"
|
|
DEFAULT_REVIEW_CSV = SCRIPT_DIR / ".artifacts" / "human-review-v1.csv"
|
|
@dataclass(frozen=True)
|
|
class AuditRecord:
|
|
prompt: str
|
|
purpose: str
|
|
prompt_hash: str
|
|
origin: str
|
|
source: str
|
|
line: int | None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SemanticCandidate:
|
|
left: int
|
|
right: int
|
|
similarity: float
|
|
|
|
|
|
def _relative(path: Path) -> str:
|
|
try:
|
|
return str(path.resolve().relative_to(REPOSITORY_ROOT))
|
|
except ValueError:
|
|
return str(path.resolve())
|
|
|
|
|
|
def semantic_candidates(
|
|
embeddings: np.ndarray,
|
|
purposes: Sequence[str],
|
|
*,
|
|
same_label_threshold: float,
|
|
cross_label_threshold: float,
|
|
neighbors: int,
|
|
block_size: int,
|
|
) -> list[SemanticCandidate]:
|
|
"""Return de-duplicated high-cosine nearest-neighbor pairs."""
|
|
|
|
if embeddings.ndim != 2 or embeddings.shape[0] != len(purposes):
|
|
raise ValueError("embeddings and purposes must have matching record counts")
|
|
if len(purposes) < 2:
|
|
return []
|
|
if not (
|
|
0.0 < same_label_threshold <= 1.0
|
|
and 0.0 < cross_label_threshold <= 1.0
|
|
):
|
|
raise ValueError("semantic thresholds must be in (0, 1]")
|
|
if neighbors <= 0 or block_size <= 0:
|
|
raise ValueError("neighbors and block size must be positive")
|
|
|
|
norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
|
|
if np.any(norms == 0):
|
|
raise ValueError("semantic embeddings must be non-zero")
|
|
normalized = embeddings.astype(np.float32, copy=False) / norms
|
|
top_count = min(neighbors, len(purposes) - 1)
|
|
found: dict[tuple[int, int], float] = {}
|
|
for start in range(0, len(purposes), block_size):
|
|
stop = min(start + block_size, len(purposes))
|
|
similarities = normalized[start:stop] @ normalized.T
|
|
local_rows = np.arange(stop - start)
|
|
similarities[local_rows, np.arange(start, stop)] = -1.0
|
|
candidate_columns = np.argpartition(
|
|
similarities, -top_count, axis=1
|
|
)[:, -top_count:]
|
|
for local_index, columns in enumerate(candidate_columns):
|
|
left = start + local_index
|
|
for right in columns:
|
|
similarity = float(similarities[local_index, right])
|
|
threshold = (
|
|
same_label_threshold
|
|
if purposes[left] == purposes[right]
|
|
else cross_label_threshold
|
|
)
|
|
if similarity < threshold:
|
|
continue
|
|
pair = (min(left, int(right)), max(left, int(right)))
|
|
found[pair] = max(found.get(pair, -1.0), similarity)
|
|
return [
|
|
SemanticCandidate(left, right, similarity)
|
|
for (left, right), similarity in sorted(
|
|
found.items(), key=lambda item: (-item[1], item[0])
|
|
)
|
|
]
|
|
|
|
|
|
def _select_device(torch: Any, requested: str) -> Any:
|
|
if requested != "auto":
|
|
return torch.device(requested)
|
|
if torch.cuda.is_available():
|
|
return torch.device("cuda")
|
|
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
|
return torch.device("mps")
|
|
return torch.device("cpu")
|
|
|
|
|
|
def embed_prompts(
|
|
records: Sequence[AuditRecord],
|
|
*,
|
|
model_name: str,
|
|
model_revision: str,
|
|
device_name: str,
|
|
batch_size: int,
|
|
) -> tuple[np.ndarray, str]:
|
|
try:
|
|
import torch
|
|
from transformers import AutoModel, AutoTokenizer
|
|
except ImportError as exc:
|
|
raise DataError(
|
|
"semantic-audit dependencies are missing; install requirements.txt"
|
|
) from exc
|
|
|
|
device = _select_device(torch, device_name)
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
model_name, revision=model_revision, use_fast=True
|
|
)
|
|
model = AutoModel.from_pretrained(model_name, revision=model_revision).to(device)
|
|
model.eval()
|
|
chunks: list[np.ndarray] = []
|
|
with torch.inference_mode():
|
|
for start in range(0, len(records), batch_size):
|
|
batch = records[start : start + batch_size]
|
|
encoded = encode_fixed_shape(
|
|
tokenizer,
|
|
[record.prompt for record in batch],
|
|
torch,
|
|
)
|
|
encoded = {key: value.to(device) for key, value in encoded.items()}
|
|
hidden = model(**encoded).last_hidden_state
|
|
mask = encoded["attention_mask"].unsqueeze(-1).expand(hidden.size()).float()
|
|
pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
|
|
pooled = torch.nn.functional.normalize(pooled, p=2, dim=1)
|
|
chunks.append(pooled.cpu().numpy())
|
|
return np.concatenate(chunks), str(device)
|
|
|
|
|
|
def _audit_records(
|
|
curated: Sequence[SourceRecord],
|
|
fixtures: Sequence[dict[str, str]],
|
|
fixtures_path: Path,
|
|
) -> list[AuditRecord]:
|
|
records = [
|
|
AuditRecord(
|
|
prompt=record.value["prompt"],
|
|
purpose=record.value["purpose"],
|
|
prompt_hash=prompt_hash(record.value["prompt"]),
|
|
origin="synthetic",
|
|
source=_relative(record.source),
|
|
line=record.line,
|
|
)
|
|
for record in curated
|
|
]
|
|
records.extend(
|
|
AuditRecord(
|
|
prompt=fixture["prompt"],
|
|
purpose=fixture["purpose"],
|
|
prompt_hash=prompt_hash(fixture["prompt"]),
|
|
origin="shipped-fixture",
|
|
source=_relative(fixtures_path),
|
|
line=None,
|
|
)
|
|
for fixture in fixtures
|
|
)
|
|
return records
|
|
|
|
|
|
def _split_membership() -> dict[str, str]:
|
|
membership: dict[str, str] = {}
|
|
artifact_dir = SCRIPT_DIR / ".artifacts" / "dataset-v1"
|
|
for split in ("train", "validation"):
|
|
path = artifact_dir / f"{split}.jsonl"
|
|
if not path.exists():
|
|
continue
|
|
for line in path.read_text(encoding="utf-8").splitlines():
|
|
value = json.loads(line)
|
|
membership[prompt_hash(value["prompt"])] = split
|
|
for line in DEFAULT_FROZEN_TEST.read_text(encoding="utf-8").splitlines():
|
|
value = json.loads(line)
|
|
membership[prompt_hash(value["prompt"])] = "test"
|
|
return membership
|
|
|
|
|
|
def audit(args: argparse.Namespace) -> dict[str, Any]:
|
|
sources = (
|
|
[path.resolve() for path in args.source]
|
|
if args.source
|
|
else default_source_paths()
|
|
)
|
|
fixtures = load_classifiable_fixtures(args.fixtures.resolve())
|
|
source_records = load_sources(sources)
|
|
lexical_curation = curate_records(
|
|
source_records,
|
|
fixtures,
|
|
near_duplicate_threshold=args.word_duplicate_threshold,
|
|
)
|
|
curated = exclude_reviewed_duplicates(
|
|
lexical_curation.records,
|
|
load_reviewed_semantic_exclusions(args.curation_review.resolve()),
|
|
)
|
|
sample = stratified_review_sample(
|
|
curated.records,
|
|
fraction=args.review_fraction,
|
|
seed=args.review_seed,
|
|
)
|
|
write_review_csv(
|
|
args.review_csv.resolve(),
|
|
sample,
|
|
source_formatter=_relative,
|
|
)
|
|
|
|
audit_records = _audit_records(
|
|
curated.records, fixtures, args.fixtures.resolve()
|
|
)
|
|
embeddings, device = embed_prompts(
|
|
audit_records,
|
|
model_name=args.model,
|
|
model_revision=args.model_revision,
|
|
device_name=args.device,
|
|
batch_size=args.batch_size,
|
|
)
|
|
candidates = semantic_candidates(
|
|
embeddings,
|
|
[record.purpose for record in audit_records],
|
|
same_label_threshold=args.same_label_threshold,
|
|
cross_label_threshold=args.cross_label_threshold,
|
|
neighbors=args.neighbors,
|
|
block_size=args.block_size,
|
|
)
|
|
membership = _split_membership()
|
|
|
|
candidate_values = []
|
|
leakage_candidates = 0
|
|
for candidate in candidates:
|
|
left = audit_records[candidate.left]
|
|
right = audit_records[candidate.right]
|
|
left_split = (
|
|
"shipped-fixture"
|
|
if left.origin == "shipped-fixture"
|
|
else membership.get(left.prompt_hash, "unknown")
|
|
)
|
|
right_split = (
|
|
"shipped-fixture"
|
|
if right.origin == "shipped-fixture"
|
|
else membership.get(right.prompt_hash, "unknown")
|
|
)
|
|
crosses_frozen_boundary = (
|
|
left_split == "train"
|
|
and right_split in {"validation", "test", "shipped-fixture"}
|
|
) or (
|
|
right_split == "train"
|
|
and left_split in {"validation", "test", "shipped-fixture"}
|
|
)
|
|
leakage_candidates += int(crosses_frozen_boundary)
|
|
candidate_values.append(
|
|
{
|
|
"similarity": round(candidate.similarity, 6),
|
|
"sameLabel": left.purpose == right.purpose,
|
|
"crossesFrozenBoundary": crosses_frozen_boundary,
|
|
"left": {
|
|
"promptHash": left.prompt_hash,
|
|
"purpose": left.purpose,
|
|
"origin": left.origin,
|
|
"split": left_split,
|
|
"source": left.source,
|
|
"line": left.line,
|
|
},
|
|
"right": {
|
|
"promptHash": right.prompt_hash,
|
|
"purpose": right.purpose,
|
|
"origin": right.origin,
|
|
"split": right_split,
|
|
"source": right.source,
|
|
"line": right.line,
|
|
},
|
|
}
|
|
)
|
|
|
|
report = {
|
|
"schemaVersion": 1,
|
|
"auditVersion": "purpose-semantic-audit-v1",
|
|
"embedding": {
|
|
"model": args.model,
|
|
"revision": args.model_revision,
|
|
"fixedInputShape": [1, MAX_LENGTH],
|
|
"truncation": {
|
|
"strategy": "head-tail-pair",
|
|
"headTokens": HEAD_TOKENS,
|
|
"tailTokens": TAIL_TOKENS,
|
|
},
|
|
"pooling": "attention-mask mean pooling followed by L2 normalization",
|
|
"device": device,
|
|
},
|
|
"population": {
|
|
"canonicalSourceRecords": len(source_records),
|
|
"curatedSyntheticRecords": len(curated.records),
|
|
"classifiableShippedFixtures": len(fixtures),
|
|
"auditedRecords": len(audit_records),
|
|
},
|
|
"candidatePolicy": {
|
|
"sameLabelCosineThreshold": args.same_label_threshold,
|
|
"crossLabelCosineThreshold": args.cross_label_threshold,
|
|
"nearestNeighborsPerRecord": args.neighbors,
|
|
"automaticExclusion": False,
|
|
"note": (
|
|
"Embedding similarity proposes review candidates only. It cannot safely "
|
|
"distinguish deliberate boundary pairs from duplicated labels."
|
|
),
|
|
},
|
|
"summary": {
|
|
"candidates": len(candidate_values),
|
|
"sameLabelCandidates": sum(item["sameLabel"] for item in candidate_values),
|
|
"crossLabelCandidates": sum(
|
|
not item["sameLabel"] for item in candidate_values
|
|
),
|
|
"frozenBoundaryCandidates": leakage_candidates,
|
|
},
|
|
"humanReviewSample": {
|
|
"status": "planned",
|
|
"fraction": args.review_fraction,
|
|
"seed": args.review_seed,
|
|
"populationRecords": len(curated.records),
|
|
"sampleRecords": len(sample),
|
|
"strata": ["purpose", "slice", "primary language"],
|
|
"generatedCSV": _relative(args.review_csv),
|
|
"samplePurposeCounts": dict(
|
|
sorted(Counter(record.value["purpose"] for record in sample).items())
|
|
),
|
|
"sampleSliceCounts": dict(
|
|
sorted(Counter(record.value["slice"] for record in sample).items())
|
|
),
|
|
"requiredFields": [
|
|
"reviewedPurpose",
|
|
"reviewedSecondary",
|
|
"reviewedDifficulty",
|
|
"reviewedSlice",
|
|
"reviewStatus",
|
|
],
|
|
},
|
|
"candidates": candidate_values,
|
|
}
|
|
write_json(args.report.resolve(), report)
|
|
return report
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--source", action="append", type=Path)
|
|
parser.add_argument("--fixtures", type=Path, default=DEFAULT_FIXTURES)
|
|
parser.add_argument("--report", type=Path, default=DEFAULT_REPORT)
|
|
parser.add_argument("--review-csv", type=Path, default=DEFAULT_REVIEW_CSV)
|
|
parser.add_argument(
|
|
"--curation-review", type=Path, default=DEFAULT_CURATION_REVIEW
|
|
)
|
|
parser.add_argument("--review-fraction", type=float, default=DEFAULT_SAMPLE_FRACTION)
|
|
parser.add_argument("--review-seed", type=int, default=DEFAULT_SAMPLE_SEED)
|
|
parser.add_argument("--model", default=DEFAULT_MODEL)
|
|
parser.add_argument("--model-revision", default=DEFAULT_MODEL_REVISION)
|
|
parser.add_argument("--device", default="auto")
|
|
parser.add_argument("--batch-size", type=int, default=64)
|
|
parser.add_argument("--block-size", type=int, default=256)
|
|
parser.add_argument("--neighbors", type=int, default=8)
|
|
parser.add_argument("--same-label-threshold", type=float, default=0.97)
|
|
parser.add_argument("--cross-label-threshold", type=float, default=0.985)
|
|
parser.add_argument("--word-duplicate-threshold", type=float, default=0.92)
|
|
return parser
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
|
parser = build_parser()
|
|
args = parser.parse_args(argv)
|
|
if args.batch_size <= 0 or args.block_size <= 0 or args.neighbors <= 0:
|
|
parser.error("batch size, block size, and neighbors must be positive")
|
|
try:
|
|
report = audit(args)
|
|
except (DataError, OSError, UnicodeError, ValueError) as exc:
|
|
print(f"error: {exc}", file=sys.stderr)
|
|
return 1
|
|
print(
|
|
f"Semantic audit: {report['population']['auditedRecords']} records, "
|
|
f"{report['summary']['candidates']} candidates "
|
|
f"({report['summary']['frozenBoundaryCandidates']} cross a frozen boundary); "
|
|
f"{report['humanReviewSample']['sampleRecords']} prompts sampled for review."
|
|
)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|