"""Build a deterministic, leakage-filtered OraRL JSONL mixture. Run from the OraRL directory with:: cp configs/data_sources.example.yaml ./data_sources.local.yaml python -m orarl.data.build --config ./data_sources.local.yaml --output ./prepared/train.jsonl """ from __future__ import annotations import argparse import hashlib import json import os import sys import tempfile from collections import Counter, defaultdict from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from pathlib import Path from typing import Any, Optional, Union from .identity import media_anchors, media_leakage_tokens, prompt_identity, stable_rank from .schema import ( SchemaError, canonicalize_record, is_remote_reference, read_records, ) PathLike = Union[str, os.PathLike[str]] class ConfigError(ValueError): """Raised for an invalid source manifest.""" class ShortfallError(RuntimeError): """Raised when a requested train or canary count cannot be reached.""" @dataclass(frozen=True) class SourceSpec: name: str input_path: Path task: str family: str quota: int media_root: Path preserve_problem_type: bool license: Optional[str] url: Optional[str] @dataclass(frozen=True) class Candidate: record: dict[str, Any] spec: SourceSpec source_index: int identity: str anchors: tuple[str, ...] leakage_tokens: frozenset[str] row_digest: str @dataclass(frozen=True) class Exclusions: identities: frozenset[str] media_tokens: frozenset[str] rows: int paths: tuple[Path, ...] checksums: Mapping[str, str] def _sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def _json_bytes(value: Any) -> bytes: return json.dumps( value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), allow_nan=False, ).encode("utf-8") def _record_digest(record: Mapping[str, Any]) -> str: return hashlib.sha256(_json_bytes(record)).hexdigest() def _atomic_write_jsonl(path: Path, records: Iterable[Mapping[str, Any]]) -> None: path.parent.mkdir(parents=True, exist_ok=True) descriptor, temporary_name = tempfile.mkstemp( prefix=path.name + ".", suffix=".tmp", dir=str(path.parent), ) try: with os.fdopen(descriptor, "wb") as handle: for record in records: handle.write(_json_bytes(record)) handle.write(b"\n") os.replace(temporary_name, path) except Exception: try: os.unlink(temporary_name) except FileNotFoundError: pass raise def _atomic_write_json(path: Path, value: Any) -> None: path.parent.mkdir(parents=True, exist_ok=True) descriptor, temporary_name = tempfile.mkstemp( prefix=path.name + ".", suffix=".tmp", dir=str(path.parent), ) try: with os.fdopen(descriptor, "w", encoding="utf-8") as handle: json.dump( value, handle, ensure_ascii=False, sort_keys=True, indent=2, allow_nan=False, ) handle.write("\n") os.replace(temporary_name, path) except Exception: try: os.unlink(temporary_name) except FileNotFoundError: pass raise def _load_yaml(path: Path) -> Any: try: import yaml except ImportError as error: raise ConfigError( f"{path}: YAML input requires PyYAML; use a JSON manifest if it is unavailable" ) from error try: with path.open("r", encoding="utf-8") as handle: return yaml.safe_load(handle) except yaml.YAMLError as error: raise ConfigError(f"{path}: invalid YAML: {error}") from error def _load_config(path: Path) -> dict[str, Any]: suffix = path.suffix.casefold() if suffix in {".yaml", ".yml"}: payload = _load_yaml(path) else: try: with path.open("r", encoding="utf-8") as handle: payload = json.load(handle) except json.JSONDecodeError: payload = _load_yaml(path) if not isinstance(payload, dict): raise ConfigError(f"{path}: source manifest must be a mapping") return payload def _local_path(base: Path, value: Any, label: str) -> Path: if not isinstance(value, str) or not value.strip(): raise ConfigError(f"{label} must be a nonempty local path") if is_remote_reference(value): raise ConfigError(f"{label} must be local; URL fields are metadata only") path = Path(os.path.expanduser(value.strip())) if not path.is_absolute(): path = base / path return Path(os.path.abspath(os.path.normpath(str(path)))) def _optional_text(value: Any, label: str) -> Optional[str]: if value is None: return None if not isinstance(value, str) or not value.strip(): raise ConfigError(f"{label} must be a nonempty string when supplied") return value.strip() def _integer(value: Any, label: str, minimum: int = 0) -> int: if isinstance(value, bool): raise ConfigError(f"{label} must be an integer") try: parsed = int(value) except (TypeError, ValueError) as error: raise ConfigError(f"{label} must be an integer") from error if parsed != value and not (isinstance(value, str) and str(parsed) == value.strip()): raise ConfigError(f"{label} must be an integer") if parsed < minimum: raise ConfigError(f"{label} must be at least {minimum}") return parsed def _boolean(value: Any, label: str, default: bool) -> bool: if value is None: return default if not isinstance(value, bool): raise ConfigError(f"{label} must be true or false") return value def _source_specs(config_path: Path, config: Mapping[str, Any]) -> list[SourceSpec]: raw_sources = config.get("sources") if not isinstance(raw_sources, list) or not raw_sources: raise ConfigError(f"{config_path}: sources must be a nonempty list") specs: list[SourceSpec] = [] names: set[str] = set() for index, raw in enumerate(raw_sources): label = f"{config_path}: sources[{index}]" if not isinstance(raw, Mapping): raise ConfigError(f"{label} must be a mapping") raw_input = raw.get("input", raw.get("path")) input_path = _local_path(config_path.parent, raw_input, f"{label}.input") raw_name = raw.get("name", raw.get("source")) name = _optional_text(raw_name, f"{label}.name") if name is None: name = input_path.stem if name in names: raise ConfigError(f"{label}.name is duplicated: {name}") names.add(name) task = _optional_text(raw.get("task"), f"{label}.task") family = _optional_text(raw.get("family"), f"{label}.family") if task is None or family is None: raise ConfigError(f"{label} must define task and family") quota = _integer(raw.get("quota"), f"{label}.quota") media_root = _local_path( config_path.parent, raw.get("media_root"), f"{label}.media_root", ) if not input_path.is_file(): raise ConfigError(f"{label}.input does not exist: {input_path}") specs.append( SourceSpec( name=name, input_path=input_path, task=task, family=family, quota=quota, media_root=media_root, preserve_problem_type=_boolean( raw.get("preserve_problem_type"), f"{label}.preserve_problem_type", False, ), license=_optional_text(raw.get("license"), f"{label}.license"), url=_optional_text(raw.get("url"), f"{label}.url"), ) ) return specs def _benchmark_paths( config_path: Path, config: Mapping[str, Any], command_line_paths: Sequence[PathLike], ) -> list[Path]: raw_paths = config.get( "benchmark_excludes", config.get("benchmark_exclude", []), ) if raw_paths is None: raw_paths = [] if isinstance(raw_paths, (str, Mapping)): raw_paths = [raw_paths] if not isinstance(raw_paths, list): raise ConfigError(f"{config_path}: benchmark_excludes must be a list") paths: list[Path] = [] for index, raw in enumerate(raw_paths): value = raw.get("path") if isinstance(raw, Mapping) else raw paths.append( _local_path( config_path.parent, value, f"{config_path}: benchmark_excludes[{index}]", ) ) for index, raw in enumerate(command_line_paths): paths.append(_local_path(Path.cwd(), str(raw), f"--benchmark-exclude[{index}]")) unique: list[Path] = [] seen: set[str] = set() for path in paths: key = os.path.normcase(str(path)) if key not in seen: seen.add(key) if not path.is_file(): raise ConfigError(f"benchmark exclusion file does not exist: {path}") unique.append(path) return unique def _loose_benchmark_record(raw: Mapping[str, Any], root: Path) -> tuple[dict[str, Any], bool]: has_problem = any( isinstance(raw.get(key), str) and bool(str(raw[key]).strip()) for key in ("problem", "question", "prompt") ) working = dict(raw) if not has_problem: working["problem"] = "media exclusion" working["answer"] = "excluded" if not working.get("problem_type"): working["problem_type"] = "benchmark" if not any( working.get(key) for key in ("images", "videos", "image", "video", "image_path", "video_path") ): generic_media = raw.get("media", raw.get("media_path", raw.get("path"))) if generic_media: working["videos"] = generic_media return ( canonicalize_record( working, source="benchmark", media_root=root, require_media=False, context="benchmark exclusion", ), has_problem and bool(raw.get("problem_type")), ) def _load_exclusions(paths: Sequence[Path]) -> Exclusions: identities: set[str] = set() tokens: set[str] = set() rows = 0 checksums: dict[str, str] = {} for path in paths: checksums[str(path)] = _sha256_file(path) for raw in read_records(path): rows += 1 explicit_identity = raw.get("prompt_identity", raw.get("identity")) if isinstance(explicit_identity, str) and explicit_identity.strip(): identities.add(explicit_identity.strip()) try: record, has_full_identity = _loose_benchmark_record(raw, path.parent) except SchemaError: if explicit_identity: continue raise tokens.update(media_leakage_tokens(record)) if has_full_identity: identities.add(prompt_identity(record)) return Exclusions( identities=frozenset(identities), media_tokens=frozenset(tokens), rows=rows, paths=tuple(paths), checksums=checksums, ) def _load_candidates( specs: Sequence[SourceSpec], require_media: bool, ) -> tuple[list[Candidate], dict[str, int], dict[str, str]]: candidates: list[Candidate] = [] input_counts: dict[str, int] = {} checksums: dict[str, str] = {} for spec in specs: raw_records = read_records(spec.input_path) input_counts[spec.name] = len(raw_records) checksums[spec.name] = _sha256_file(spec.input_path) for index, raw in enumerate(raw_records): prepared = dict(raw) if ( not spec.preserve_problem_type or not str(prepared.get("problem_type") or "").strip() ): prepared["problem_type"] = spec.task context = f"{spec.input_path}: record {index + 1}" record = canonicalize_record( prepared, source=spec.name, family=spec.family, media_root=spec.media_root, require_media=require_media, context=context, ) try: digest = _record_digest(record) except (TypeError, ValueError) as error: raise SchemaError(f"{context}: record is not strict JSON: {error}") from error candidates.append( Candidate( record=record, spec=spec, source_index=index, identity=prompt_identity(record), anchors=media_anchors(record), leakage_tokens=media_leakage_tokens(record), row_digest=digest, ) ) return candidates, input_counts, checksums def _candidate_order(candidate: Candidate, seed: int, namespace: str) -> tuple[int, str, int]: key = f"{candidate.identity}|{candidate.row_digest}" return ( stable_rank(seed, namespace, key), candidate.row_digest, candidate.source_index, ) def _select_train( specs: Sequence[SourceSpec], candidates: Sequence[Candidate], seed: int, max_prompts_per_media: int, ) -> tuple[list[Candidate], dict[str, int], dict[str, int]]: grouped: dict[str, list[Candidate]] = defaultdict(list) for candidate in candidates: grouped[candidate.spec.name].append(candidate) unique_capacity = { spec.name: len({candidate.identity for candidate in grouped.get(spec.name, [])}) for spec in specs } spec_order = sorted( specs, key=lambda spec: ( unique_capacity[spec.name] - spec.quota, spec.name.casefold(), ), ) selected: list[Candidate] = [] selected_identities: set[str] = set() anchor_counts: Counter[str] = Counter() cap_rejections: Counter[str] = Counter() selected_counts: Counter[str] = Counter() for spec in spec_order: ordered = sorted( grouped.get(spec.name, []), key=lambda candidate: _candidate_order( candidate, seed, f"train:{spec.name}", ), ) for candidate in ordered: if selected_counts[spec.name] >= spec.quota: break if candidate.identity in selected_identities: continue if any(anchor_counts[anchor] >= max_prompts_per_media for anchor in candidate.anchors): cap_rejections[spec.name] += 1 continue selected.append(candidate) selected_counts[spec.name] += 1 selected_identities.add(candidate.identity) anchor_counts.update(candidate.anchors) shortfalls = { spec.name: spec.quota - selected_counts[spec.name] for spec in specs if selected_counts[spec.name] < spec.quota } return selected, shortfalls, dict(cap_rejections) def _select_canary( candidates: Sequence[Candidate], train: Sequence[Candidate], canary_size: int, seed: int, max_prompts_per_media: int, ) -> tuple[list[Candidate], dict[str, int]]: train_identities = {candidate.identity for candidate in train} train_media_tokens: set[str] = set() for candidate in train: train_media_tokens.update(candidate.leakage_tokens) ordered = sorted( candidates, key=lambda candidate: _candidate_order(candidate, seed, "canary"), ) canary: list[Candidate] = [] canary_identities: set[str] = set() canary_anchor_counts: Counter[str] = Counter() stats: Counter[str] = Counter() for candidate in ordered: if len(canary) >= canary_size: break if candidate.identity in train_identities: stats["train_identity"] += 1 continue if candidate.identity in canary_identities: stats["duplicate_identity"] += 1 continue if candidate.leakage_tokens & train_media_tokens: stats["train_media"] += 1 continue if any( canary_anchor_counts[anchor] >= max_prompts_per_media for anchor in candidate.anchors ): stats["media_cap"] += 1 continue canary.append(candidate) canary_identities.add(candidate.identity) canary_anchor_counts.update(candidate.anchors) return canary, dict(stats) def _sorted_output(candidates: Sequence[Candidate], seed: int, namespace: str) -> list[Candidate]: return sorted( candidates, key=lambda candidate: _candidate_order(candidate, seed, namespace), ) def _counts(candidates: Sequence[Candidate]) -> dict[str, Any]: by_task = Counter(str(candidate.record["problem_type"]) for candidate in candidates) by_source = Counter(str(candidate.record["source"]) for candidate in candidates) by_family = Counter(str(candidate.record.get("family") or "") for candidate in candidates) return { "rows": len(candidates), "by_task": dict(sorted(by_task.items())), "by_source": dict(sorted(by_source.items())), "by_family": dict(sorted(by_family.items())), } def _default_canary_path(output: Path) -> Path: suffix = output.suffix or ".jsonl" stem = output.name[: -len(output.suffix)] if output.suffix else output.name return output.with_name(stem + ".canary" + suffix) def _default_manifest_path(output: Path) -> Path: stem = output.name[: -len(output.suffix)] if output.suffix else output.name return output.with_name(stem + ".manifest.json") def _output_path(value: PathLike) -> Path: path = Path(os.path.expanduser(str(value))) return Path(os.path.abspath(os.path.normpath(str(path)))) def _check_output_paths( paths: Sequence[Path], overwrite: bool, protected_paths: Sequence[Path], ) -> None: normalized = [os.path.normcase(str(path)) for path in paths] if len(normalized) != len(set(normalized)): raise ConfigError("train, canary, and manifest outputs must be different paths") protected = {os.path.normcase(str(path)) for path in protected_paths} collisions = [ str(path) for path, normalized_path in zip(paths, normalized) if normalized_path in protected ] if collisions: raise ConfigError( "output paths cannot replace configs, sources, or exclusions: " + ", ".join(collisions) ) if not overwrite: existing = [str(path) for path in paths if path.exists()] if existing: raise FileExistsError( "refusing to overwrite existing output(s): " + ", ".join(existing) + "; pass --overwrite" ) def _output_leakage( train: Sequence[Candidate], canary: Sequence[Candidate], exclusions: Exclusions, ) -> dict[str, int]: train_ids = {candidate.identity for candidate in train} canary_ids = {candidate.identity for candidate in canary} train_tokens: set[str] = set() canary_tokens: set[str] = set() for candidate in train: train_tokens.update(candidate.leakage_tokens) for candidate in canary: canary_tokens.update(candidate.leakage_tokens) return { "benchmark_identity": len((train_ids | canary_ids) & exclusions.identities), "benchmark_media": len((train_tokens | canary_tokens) & exclusions.media_tokens), "train_canary_identity": len(train_ids & canary_ids), "train_canary_media": len(train_tokens & canary_tokens), } def build_dataset( config_path: PathLike, output_path: PathLike, *, canary_output: Optional[PathLike] = None, manifest_output: Optional[PathLike] = None, seed: Optional[int] = None, target: Optional[int] = None, canary_size: Optional[int] = None, max_prompts_per_media: Optional[int] = None, require_media: Optional[bool] = None, benchmark_excludes: Sequence[PathLike] = (), allow_shortfall: Optional[bool] = None, overwrite: bool = False, ) -> dict[str, Any]: """Build train, canary, and audit manifest files.""" config_file = _output_path(config_path) if not config_file.is_file(): raise ConfigError(f"source manifest does not exist: {config_file}") config = _load_config(config_file) specs = _source_specs(config_file, config) effective_seed = _integer( seed if seed is not None else config.get("seed", 42), "seed", ) effective_canary_size = _integer( canary_size if canary_size is not None else config.get("canary_size", 0), "canary_size", ) effective_cap = _integer( ( max_prompts_per_media if max_prompts_per_media is not None else config.get( "max_prompts_per_media", config.get("max_prompts_per_anchor", 1), ) ), "max_prompts_per_media", minimum=1, ) effective_require_media = ( require_media if require_media is not None else _boolean(config.get("require_media"), "require_media", False) ) effective_allow_shortfall = ( allow_shortfall if allow_shortfall is not None else _boolean(config.get("allow_shortfall"), "allow_shortfall", False) ) quota_total = sum(spec.quota for spec in specs) effective_target = _integer( target if target is not None else config.get("target", quota_total), "target", ) if effective_target != quota_total: raise ConfigError( f"target ({effective_target}) must equal the sum of source quotas ({quota_total})" ) output = _output_path(output_path) canary_path = ( _output_path(canary_output) if canary_output is not None else _default_canary_path(output) ) manifest_path = ( _output_path(manifest_output) if manifest_output is not None else _default_manifest_path(output) ) benchmark_paths = _benchmark_paths(config_file, config, benchmark_excludes) protected_paths = [ config_file, *(spec.input_path for spec in specs), *benchmark_paths, ] _check_output_paths( (output, canary_path, manifest_path), overwrite, protected_paths, ) exclusions = _load_exclusions(benchmark_paths) all_candidates, input_counts, source_checksums = _load_candidates( specs, require_media=effective_require_media, ) identity_counts = Counter(candidate.identity for candidate in all_candidates) duplicate_rows = sum(count - 1 for count in identity_counts.values()) duplicate_groups = sum(count > 1 for count in identity_counts.values()) eligible: list[Candidate] = [] identity_leak_rows = 0 media_leak_rows = 0 excluded_rows = 0 for candidate in all_candidates: identity_hit = candidate.identity in exclusions.identities media_hit = bool(candidate.leakage_tokens & exclusions.media_tokens) identity_leak_rows += int(identity_hit) media_leak_rows += int(media_hit) if identity_hit or media_hit: excluded_rows += 1 else: eligible.append(candidate) train, source_shortfalls, cap_rejections = _select_train( specs, eligible, effective_seed, effective_cap, ) if source_shortfalls and not effective_allow_shortfall: detail = ", ".join( f"{name} short by {count}" for name, count in sorted(source_shortfalls.items()) ) raise ShortfallError(f"source quota shortfall: {detail}") canary, canary_rejections = _select_canary( eligible, train, effective_canary_size, effective_seed, effective_cap, ) canary_shortfall = effective_canary_size - len(canary) if canary_shortfall and not effective_allow_shortfall: raise ShortfallError( f"canary shortfall: requested {effective_canary_size}, got {len(canary)}" ) train = _sorted_output(train, effective_seed, "train-output") canary = _sorted_output(canary, effective_seed, "canary-output") output_leakage = _output_leakage(train, canary, exclusions) if any(output_leakage.values()): raise RuntimeError(f"internal leakage invariant failed: {output_leakage}") _atomic_write_jsonl(output, (candidate.record for candidate in train)) _atomic_write_jsonl(canary_path, (candidate.record for candidate in canary)) source_details = [] selected_by_source = Counter(candidate.spec.name for candidate in train) for spec in sorted(specs, key=lambda value: value.name.casefold()): source_details.append( { "name": spec.name, "input": str(spec.input_path), "media_root": str(spec.media_root), "task": spec.task, "family": spec.family, "quota": spec.quota, "selected": selected_by_source[spec.name], "shortfall": spec.quota - selected_by_source[spec.name], "preserve_problem_type": spec.preserve_problem_type, "license": spec.license, "url": spec.url, "input_sha256": source_checksums[spec.name], } ) manifest: dict[str, Any] = { "schema_version": 1, "seed": effective_seed, "target": effective_target, "actual_train_rows": len(train), "canary_target": effective_canary_size, "actual_canary_rows": len(canary), "allow_shortfall": effective_allow_shortfall, "media": { "require_existing": effective_require_media, "max_prompts_per_anchor": effective_cap, }, "files": { "train": str(output), "canary": str(canary_path), "manifest": str(manifest_path), }, "checksums": { "algorithm": "sha256", "config": _sha256_file(config_file), "train": _sha256_file(output), "canary": _sha256_file(canary_path), "sources": dict(sorted(source_checksums.items())), "benchmark_excludes": dict(sorted(exclusions.checksums.items())), }, "counts": { "input_rows": len(all_candidates), "input_by_source": dict(sorted(input_counts.items())), "eligible_rows": len(eligible), "train": _counts(train), "canary": _counts(canary), }, "duplicates": { "candidate_prompt_identity_groups": duplicate_groups, "candidate_prompt_rows": duplicate_rows, "train_prompt_rows": len(train) - len({candidate.identity for candidate in train}), "canary_prompt_rows": len(canary) - len({candidate.identity for candidate in canary}), }, "leakage": { "benchmark_reference_rows": exclusions.rows, "candidate_identity_rows_excluded": identity_leak_rows, "candidate_media_rows_excluded": media_leak_rows, "candidate_rows_excluded": excluded_rows, "output": output_leakage, }, "selection": { "source_shortfalls": dict(sorted(source_shortfalls.items())), "canary_shortfall": canary_shortfall, "media_cap_rejections_by_source": dict(sorted(cap_rejections.items())), "canary_rejections": dict(sorted(canary_rejections.items())), }, "sources": source_details, "benchmark_excludes": [str(path) for path in exclusions.paths], } _atomic_write_json(manifest_path, manifest) return manifest def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--config", required=True, help="YAML or JSON source manifest") parser.add_argument("--output", required=True, help="Training JSONL output") parser.add_argument("--canary-output", default=None) parser.add_argument("--manifest-output", default=None) parser.add_argument("--seed", type=int, default=None) parser.add_argument("--target", type=int, default=None) parser.add_argument("--canary-size", type=int, default=None) parser.add_argument( "--max-prompts-per-media", "--max-prompts-per-anchor", dest="max_prompts_per_media", type=int, default=None, ) parser.add_argument( "--require-media", action=argparse.BooleanOptionalAction, default=None, help="Require every normalized media path to exist", ) parser.add_argument( "--require-existing-media", dest="require_media", action="store_true", help=argparse.SUPPRESS, ) parser.add_argument( "--benchmark-exclude", action="append", default=[], help="JSON or JSONL benchmark identities/media to exclude; repeatable", ) parser.add_argument( "--allow-shortfall", action="store_true", default=None, help="Write available rows instead of failing a source or canary shortfall", ) parser.add_argument("--overwrite", action="store_true") return parser def main(argv: Optional[Sequence[str]] = None) -> int: args = build_parser().parse_args(argv) try: manifest = build_dataset( args.config, args.output, canary_output=args.canary_output, manifest_output=args.manifest_output, seed=args.seed, target=args.target, canary_size=args.canary_size, max_prompts_per_media=args.max_prompts_per_media, require_media=args.require_media, benchmark_excludes=args.benchmark_exclude, allow_shortfall=args.allow_shortfall, overwrite=args.overwrite, ) except (ConfigError, FileExistsError, OSError, SchemaError, ShortfallError) as error: print(f"[FATAL] {error}", file=sys.stderr) return 2 print( "OraRL data prepared: " f"train={manifest['actual_train_rows']} " f"canary={manifest['actual_canary_rows']} " f"manifest={manifest['files']['manifest']}" ) return 0 if __name__ == "__main__": raise SystemExit(main())