| """Deterministic, group-safe and approximately stratified split assignment.""" |
|
|
| from __future__ import annotations |
|
|
| from collections import Counter, defaultdict |
| import hashlib |
| import heapq |
| import json |
| import math |
| from typing import Any, Iterable, Iterator, Mapping, MutableMapping, Sequence |
|
|
| from .grouping import attach_group_ids |
|
|
|
|
| DEFAULT_SPLIT_RATIOS: dict[str, float] = {"train": 0.9, "validation": 0.1} |
|
|
|
|
| class SplitValidationError(ValueError): |
| """Raised when a split manifest contains leakage or invalid assignments.""" |
|
|
|
|
| def normalize_split_ratios(ratios: Mapping[str, float]) -> dict[str, float]: |
| """Validate and normalize positive split weights to sum to one.""" |
|
|
| if not ratios: |
| raise ValueError("at least one split ratio is required") |
| normalized: dict[str, float] = {} |
| for name, raw_value in ratios.items(): |
| split_name = str(name).strip() |
| if not split_name: |
| raise ValueError("split names cannot be empty") |
| value = float(raw_value) |
| if not math.isfinite(value) or value <= 0: |
| raise ValueError(f"ratio for {split_name!r} must be finite and positive") |
| if split_name in normalized: |
| raise ValueError(f"duplicate split name: {split_name}") |
| normalized[split_name] = value |
| total = sum(normalized.values()) |
| return {name: value / total for name, value in normalized.items()} |
|
|
|
|
| def parse_split_ratios(values: Sequence[str]) -> dict[str, float]: |
| """Parse CLI values such as ``train=0.9`` and ``validation=0.1``.""" |
|
|
| parsed: dict[str, float] = {} |
| for value in values: |
| if "=" not in value: |
| raise ValueError(f"split ratio must use NAME=WEIGHT syntax: {value!r}") |
| name, raw_ratio = value.split("=", 1) |
| name = name.strip() |
| if name in parsed: |
| raise ValueError(f"duplicate split name: {name}") |
| try: |
| parsed[name] = float(raw_ratio) |
| except ValueError as exc: |
| raise ValueError(f"invalid ratio in {value!r}") from exc |
| return normalize_split_ratios(parsed) |
|
|
|
|
| def _stable_hash(*values: Any) -> int: |
| payload = "\0".join(str(value) for value in values).encode("utf-8") |
| return int.from_bytes(hashlib.sha256(payload).digest()[:8], "big") |
|
|
|
|
| def _category(value: Any) -> str: |
| if value is None: |
| return "<null>" |
| if isinstance(value, bool): |
| return "true" if value else "false" |
| return str(value).strip() or "<empty>" |
|
|
|
|
| def _field_value(row: Mapping[str, Any], field: str) -> Any: |
| current: Any = row |
| for part in field.split("."): |
| if not isinstance(current, Mapping) or part not in current: |
| return None |
| current = current[part] |
| return current |
|
|
|
|
| def _add_holdout_links(rows: Sequence[MutableMapping[str, Any]], fields: Sequence[str]) -> None: |
| if not fields: |
| return |
| for row in rows: |
| raw_keys = row.get("group_keys") or [] |
| keys = [str(raw_keys)] if isinstance(raw_keys, str) else [str(key) for key in raw_keys] |
| for field in fields: |
| value = _field_value(row, field) |
| if value is None or value == "": |
| continue |
| canonical = json.dumps(value, sort_keys=True, ensure_ascii=False, default=str) |
| digest = hashlib.sha256(f"{field}\0{canonical}".encode("utf-8")).hexdigest()[:32] |
| keys.append(f"holdout:{field}:{digest}") |
| row["group_keys"] = sorted(set(keys)) |
| attach_group_ids(rows) |
|
|
|
|
| def assign_splits( |
| rows: Sequence[Mapping[str, Any]], |
| *, |
| ratios: Mapping[str, float] | None = None, |
| seed: int = 42, |
| stratify_fields: Sequence[str] = ("endpoint", "language", "dataset"), |
| holdout_fields: Sequence[str] = (), |
| ) -> list[dict[str, Any]]: |
| """Assign whole leakage components using deterministic greedy balancing. |
| |
| The objective balances total rows and each requested marginal stratum. It |
| is deterministic for a fixed set of rows regardless of input row order. |
| ``holdout_fields`` can enforce source-, speaker-, or conversation-held-out |
| evaluation by linking all equal values before assignment. |
| """ |
|
|
| split_ratios = normalize_split_ratios(ratios or DEFAULT_SPLIT_RATIOS) |
| output = [dict(row) for row in rows] |
| if not output: |
| return output |
| if any(not row.get("group_id") for row in output): |
| attach_group_ids(output) |
| _add_holdout_links(output, holdout_fields) |
|
|
| group_indices: dict[str, list[int]] = defaultdict(list) |
| for index, row in enumerate(output): |
| group_indices[str(row["group_id"])].append(index) |
|
|
| total_rows = len(output) |
| all_strata: Counter[str] = Counter() |
| group_strata: dict[str, Counter[str]] = {} |
| for group_id, indices in group_indices.items(): |
| counts: Counter[str] = Counter() |
| for index in indices: |
| for field in stratify_fields: |
| key = f"{field}={_category(_field_value(output[index], field))}" |
| counts[key] += 1 |
| all_strata[key] += 1 |
| group_strata[group_id] = counts |
|
|
| |
| |
| |
| |
| remaining_groups = set(group_indices) |
| remaining_strata = Counter(all_strata) |
| desired_total = {name: total_rows * ratio for name, ratio in split_ratios.items()} |
| desired_strata = { |
| name: {key: count * ratio for key, count in all_strata.items()} |
| for name, ratio in split_ratios.items() |
| } |
| candidate_heaps: dict[str, list[tuple[int, int, int, str]]] = defaultdict(list) |
| fallback_heap: list[tuple[int, int, str]] = [] |
| for group_id, indices in group_indices.items(): |
| group_size = len(indices) |
| heapq.heappush( |
| fallback_heap, |
| (-group_size, _stable_hash(seed, "fallback-group", group_id), group_id), |
| ) |
| for stratum, contribution in group_strata[group_id].items(): |
| heapq.heappush( |
| candidate_heaps[stratum], |
| ( |
| -contribution, |
| -group_size, |
| _stable_hash(seed, "stratum-group", stratum, group_id), |
| group_id, |
| ), |
| ) |
| group_assignment: dict[str, str] = {} |
|
|
| while remaining_groups: |
| active_strata = [ |
| (count, _stable_hash(seed, "stratum-order", stratum), stratum) |
| for stratum, count in remaining_strata.items() |
| if count > 0 |
| ] |
| focus_stratum = min(active_strata)[2] if active_strata else None |
| if focus_stratum is not None: |
| candidates = candidate_heaps[focus_stratum] |
| while candidates and candidates[0][3] not in remaining_groups: |
| heapq.heappop(candidates) |
| if not candidates: |
| |
| focus_stratum = None |
| if focus_stratum is None: |
| while fallback_heap and fallback_heap[0][2] not in remaining_groups: |
| heapq.heappop(fallback_heap) |
| if not fallback_heap: |
| raise RuntimeError("internal split assignment error: no remaining group candidate") |
| group_id = fallback_heap[0][2] |
| else: |
| group_id = candidate_heaps[focus_stratum][0][3] |
|
|
| group_size = len(group_indices[group_id]) |
| strata = group_strata[group_id] |
| candidate_scores: list[tuple[float, float, float, int, str]] = [] |
| for split_name in split_ratios: |
| primary_need = ( |
| desired_strata[split_name].get(focus_stratum, 0.0) |
| if focus_stratum is not None |
| else desired_total[split_name] |
| ) |
| |
| |
| secondary_need = sum( |
| max(desired_strata[split_name].get(stratum, 0.0), 0.0) |
| * contribution |
| / max(all_strata[stratum], 1) |
| for stratum, contribution in strata.items() |
| ) |
| candidate_scores.append( |
| ( |
| -primary_need, |
| -secondary_need, |
| -desired_total[split_name], |
| _stable_hash(seed, "split-choice", group_id, split_name), |
| split_name, |
| ) |
| ) |
|
|
| selected = min(candidate_scores)[4] |
| group_assignment[group_id] = selected |
| remaining_groups.remove(group_id) |
| desired_total[selected] -= group_size |
| for stratum, contribution in strata.items(): |
| desired_strata[selected][stratum] -= contribution |
| remaining_strata[stratum] -= contribution |
|
|
| for row in output: |
| row["split"] = group_assignment[str(row["group_id"])] |
| return output |
|
|
|
|
| def assign_leave_one_out( |
| rows: Sequence[Mapping[str, Any]], |
| *, |
| field: str, |
| held_out_value: Any, |
| train_split: str = "train", |
| held_out_split: str = "validation", |
| ) -> list[dict[str, Any]]: |
| """Assign a named source/domain value to one held-out stress-test split. |
| |
| If a pre-existing leakage component contains both held-out and non-held-out |
| values, the entire component is held out. This preserves leakage safety at |
| the cost of a small amount of train-domain spillover, which is surfaced by |
| :func:`build_split_report`. |
| """ |
|
|
| if not field.strip(): |
| raise ValueError("leave-one-out field cannot be empty") |
| if not train_split or not held_out_split or train_split == held_out_split: |
| raise ValueError("train and held-out split names must be distinct and non-empty") |
| output = [dict(row) for row in rows] |
| if any(not row.get("group_id") for row in output): |
| attach_group_ids(output) |
| expected = _category(held_out_value) |
| held_out_groups = { |
| str(row["group_id"]) |
| for row in output |
| if _category(_field_value(row, field)) == expected |
| } |
| if not held_out_groups: |
| raise ValueError(f"held-out value {held_out_value!r} was not found in field {field!r}") |
| for row in output: |
| row["split"] = held_out_split if str(row["group_id"]) in held_out_groups else train_split |
| assert_no_split_leakage(output) |
| return output |
|
|
|
|
| def iter_leave_one_out_folds( |
| rows: Sequence[Mapping[str, Any]], |
| *, |
| field: str = "dataset", |
| values: Sequence[Any] | None = None, |
| train_split: str = "train", |
| held_out_split: str = "validation", |
| ) -> Iterator[tuple[str, list[dict[str, Any]]]]: |
| """Yield deterministic leave-one-source/domain-out manifests one at a time.""" |
|
|
| if values is None: |
| observed = { |
| _category(_field_value(row, field)) |
| for row in rows |
| if _field_value(row, field) is not None and _field_value(row, field) != "" |
| } |
| selected_values: Sequence[Any] = sorted(observed) |
| else: |
| selected_values = values |
| for value in selected_values: |
| canonical = _category(value) |
| yield canonical, assign_leave_one_out( |
| rows, |
| field=field, |
| held_out_value=value, |
| train_split=train_split, |
| held_out_split=held_out_split, |
| ) |
|
|
|
|
| def find_split_leakage(rows: Iterable[Mapping[str, Any]]) -> dict[str, Any]: |
| """Detect group, exact-audio and metadata-key crossings between splits.""" |
|
|
| split_by_group: dict[str, set[str]] = defaultdict(set) |
| split_by_audio: dict[str, set[str]] = defaultdict(set) |
| split_by_key: dict[str, set[str]] = defaultdict(set) |
| missing_split = 0 |
| for row in rows: |
| split = row.get("split") |
| if not split: |
| missing_split += 1 |
| continue |
| split_name = str(split) |
| if row.get("group_id"): |
| split_by_group[str(row["group_id"])].add(split_name) |
| if row.get("audio_sha256"): |
| split_by_audio[str(row["audio_sha256"])].add(split_name) |
| raw_keys = row.get("group_keys") or [] |
| if isinstance(raw_keys, str): |
| raw_keys = [raw_keys] |
| for key in raw_keys: |
| split_by_key[str(key)].add(split_name) |
|
|
| group_crossings = {key: sorted(value) for key, value in split_by_group.items() if len(value) > 1} |
| audio_crossings = {key: sorted(value) for key, value in split_by_audio.items() if len(value) > 1} |
| key_crossings = {key: sorted(value) for key, value in split_by_key.items() if len(value) > 1} |
| return { |
| "missing_split_rows": missing_split, |
| "group_crossings": group_crossings, |
| "audio_hash_crossings": audio_crossings, |
| "metadata_key_crossings": key_crossings, |
| "is_valid": not (missing_split or group_crossings or audio_crossings or key_crossings), |
| } |
|
|
|
|
| def assert_no_split_leakage(rows: Iterable[Mapping[str, Any]]) -> None: |
| """Raise :class:`SplitValidationError` if any linkage crosses splits.""" |
|
|
| report = find_split_leakage(rows) |
| if not report["is_valid"]: |
| raise SplitValidationError( |
| "split leakage detected: " |
| f"missing={report['missing_split_rows']}, groups={len(report['group_crossings'])}, " |
| f"audio={len(report['audio_hash_crossings'])}, keys={len(report['metadata_key_crossings'])}" |
| ) |
|
|
|
|
| def build_split_report( |
| rows: Sequence[Mapping[str, Any]], |
| *, |
| stratify_fields: Sequence[str] = ("endpoint", "language", "dataset"), |
| holdout_fields: Sequence[str] = (), |
| ) -> dict[str, Any]: |
| """Summarize split sizes, strata and leakage validation.""" |
|
|
| counts = Counter(str(row.get("split") or "<missing>") for row in rows) |
| by_field: dict[str, dict[str, dict[str, int]]] = {} |
| for field in stratify_fields: |
| split_values: dict[str, Counter[str]] = defaultdict(Counter) |
| for row in rows: |
| split_values[str(row.get("split") or "<missing>")][_category(_field_value(row, field))] += 1 |
| by_field[field] = { |
| split: dict(sorted(values.items())) for split, values in sorted(split_values.items()) |
| } |
| holdout_report: dict[str, Any] = {} |
| for field in holdout_fields: |
| splits_by_value: dict[str, set[str]] = defaultdict(set) |
| for row in rows: |
| value = _field_value(row, field) |
| if value is None or value == "": |
| continue |
| splits_by_value[_category(value)].add(str(row.get("split") or "<missing>")) |
| values_by_split: dict[str, list[str]] = defaultdict(list) |
| for value, splits in splits_by_value.items(): |
| for split in splits: |
| values_by_split[split].append(value) |
| holdout_report[field] = { |
| "by_split": { |
| split: sorted(values) for split, values in sorted(values_by_split.items()) |
| }, |
| "crossing_values": { |
| value: sorted(splits) for value, splits in splits_by_value.items() if len(splits) > 1 |
| }, |
| } |
| domain_shift: dict[str, Any] | None = None |
| if holdout_fields: |
| confounder_fields = ("endpoint", "language", "dataset", "synthetic", "midfiller", "endfiller") |
| confounder_counts: dict[str, dict[str, dict[str, int]]] = {} |
| for field in confounder_fields: |
| split_values: dict[str, Counter[str]] = defaultdict(Counter) |
| for row in rows: |
| split_values[str(row.get("split") or "<missing>")][ |
| _category(_field_value(row, field)) |
| ] += 1 |
| confounder_counts[field] = { |
| split: dict(sorted(values.items())) |
| for split, values in sorted(split_values.items()) |
| } |
| domain_shift = { |
| "kind": "domain_shift_stress_test", |
| "caution": ( |
| "Held-out source/domain values are confounded with language, synthetic status, " |
| "label and collection process; results measure joint domain shift and must not " |
| "be interpreted as a clean causal source effect." |
| ), |
| "slice_counts": confounder_counts, |
| } |
| return { |
| "records": len(rows), |
| "split_counts": dict(sorted(counts.items())), |
| "strata": by_field, |
| "holdouts": holdout_report, |
| "domain_shift": domain_shift, |
| "leakage": find_split_leakage(rows), |
| } |
|
|
