"""Example 9 — Cut a dataset into train/val/test with the split engine: rules, determinism, selectors. Splits are cut over a **dataset**, never a live workset: the engine needs a member set that cannot move underneath it (that is why 08 converts first). Given a frozen parent and a canonical spec — an ``assignment_seed``, the ``splits``, and the ``rules`` — the engine assigns every canonical sample to exactly one split and materializes the results as ordinary child datasets named ``{split_set}/{split}``. Assignment is deterministic: same parent + same spec → same partition, on any machine, forever. This script is the engine itself, in four moves: 1. DISCOVER which keys can drive rules (get_dataset_split_fields) 2. THE LEAK ratio-only scatters a speaker across splits ── add EXCLUSIVITY ──▶ 0 leak 3. DETERMINISM the seed drives the random planner: same seed → identical, different → different (exclusivity/stratification use a sorted-greedy planner the seed does not move) 4. SELECTORS a split can be defined by RATIO (a share of the rest) or by ATTRIBUTE value (e.g. test = the newest recordings) — a declarative, non-random carve-out The parent has 8 speakers × 6 utterances, three locales spread independently of speaker, and a recording ``year`` — so exclusivity (keep a speaker whole), stratification (balance locales), and an attribute carve-out (hold out the newest year) each have something real to do. DATAVO_SERVER=https://dev.datavo.io python examples/09_split_a_dataset.py """ from __future__ import annotations from _common import ( build_dataset, cleanup, cleanup_datasets, cleanup_split_sets, make_client, seed_source_store, step, unique, wait_for_split_set, ) from datavo_sdk import ( DatasetSplitDraftRequest, DatasetSplitSetCreateRequest, ExclusivityRule, KeyGroupSpec, SampleStream, ShardPlan, SourceSpec, SplitAttributeSelector, SplitNumericRangePredicate, SplitOutputSpec, SplitRatioSelector, StratificationRule, ) # A parent corpus where each rule has a job. Every speaker records SPEAKER_UTTERANCES utterances # (the leak vector: one speaker's samples must not straddle splits). Locale is assigned by sample # index, *independently* of speaker, so stratification has something to balance that exclusivity # does not already fix. Year is the attribute a temporal carve-out selects on. SPEAKERS = [f"spk_{i}" for i in range(8)] SPEAKER_UTTERANCES = 6 LOCALES = ["de-DE", "en-US", "fr-FR"] YEARS = [2022, 2023, 2024] SAMPLES = [ { "__key__": f"{s}-{u:02d}", "audio.wav": f"audio-{s}-{u}".encode(), "speaker.txt": speaker, # categorical scalar (.txt) → the leak vector "locale.txt": LOCALES[n % len(LOCALES)], # categorical scalar → stratify target "year.cls": YEARS[n % len(YEARS)], # numeric scalar (.cls) → attribute selector } for n, (s, u, speaker) in enumerate( (speaker, u, speaker) for speaker in SPEAKERS for u in range(SPEAKER_UTTERANCES) ) ] RATIOS = {"train": 0.6, "val": 0.2, "test": 0.2} def ratio_splits(ratios: dict[str, float]) -> list[SplitOutputSpec]: """The plain proportional shape: each split gets a ``ratio`` share, greedily balanced.""" return [ SplitOutputSpec(name=name, selector=SplitRatioSelector(ratio=ratio)) for name, ratio in ratios.items() ] def ratio_only_spec(name: str, *, seed: int) -> DatasetSplitSetCreateRequest: """A 60/20/20 split with no rules — the *random* planner, whose one input is the seed.""" return DatasetSplitSetCreateRequest( name=name, assignment_seed=seed, splits=ratio_splits(RATIOS), shard_plan=ShardPlan(shard_size=8), ) def exclusive_spec(name: str, *, seed: int) -> DatasetSplitSetCreateRequest: """A 60/20/20 split with the two soft/hard rules that make it publishable: * **exclusivity** on ``speaker.txt`` (hard) — a speaker's samples land wholly in one split, never across two. This is the anti-leak rule; without it, ratios alone leak. * **stratification** on ``locale.txt`` (best-effort) — keep each locale's share close to the parent's, so no split is accidentally all-``de-DE``. """ return DatasetSplitSetCreateRequest( name=name, assignment_seed=seed, splits=ratio_splits(RATIOS), rules=[ ExclusivityRule(keys=["speaker.txt"]), StratificationRule(keys=["locale.txt"]), ], shard_plan=ShardPlan(shard_size=8), ) def member_ids(client, dataset_name: str) -> set[str]: """The canonical sample ids in a (child) dataset — the membership itself. ``decode=False``: we compare *which* samples landed here, not their contents.""" return { sample["__key__"] for sample in SampleStream.from_dataset(dataset_name, keys=["speaker.txt"], client=client, decode=False) } def speakers_spanning_splits(client, split_set: str, split_names: list[str]) -> set[str]: """Speakers that appear in more than one split — i.e. the leak. Streams each child for its ``speaker.txt`` values and returns the speakers seen in ≥2 splits.""" seen_in: dict[str, set[str]] = {} for split_name in split_names: for sample in SampleStream.from_dataset( f"{split_set}/{split_name}", keys=["speaker.txt"], client=client, decode=True ): speaker = sample["speaker.txt"] seen_in.setdefault(speaker, set()).add(split_name) return {speaker for speaker, splits in seen_in.items() if len(splits) > 1} def print_preview(client, dataset: str, request: DatasetSplitSetCreateRequest) -> None: """Preview writes nothing — it runs the same planner as create and reports the diagnostics. Read them before you commit: per-split ratio drift, the exclusivity groups, and any warnings. Preview takes the naked *draft* (no name, no team), so we hand it the definition from the create request rather than the create request itself.""" draft = DatasetSplitDraftRequest( assignment_seed=request.assignment_seed, splits=request.splits, rules=request.rules, shard_plan=request.shard_plan, ) preview = client.preview_dataset_split_set(dataset, draft) print(f" preview: {preview.candidate_count} candidates, " f"{preview.exclusive_group_count} exclusivity groups " f"(largest: {', '.join(f'{g.label}={g.size}' for g in preview.largest_groups[:3])})") for split in preview.splits: err = f"{split.ratio_error:.3f}" if split.ratio_error is not None else "—" print(f" {split.name:<5} target={split.target_ratio:<4} " f"got={split.assigned_ratio:.3f} (n={split.assigned_count}, |err| {err})") for result in preview.rule_results: print(f" rule[{result.kind}] {result.status}: {result.summary}") for warning in preview.warnings: print(f" ⚠ {warning}") def main() -> None: client = make_client() store = unique("example_split_store") dataset = unique("example_split_ds") # Five split sets over one parent: the leak before/after, two determinism probes, and an # attribute carve-out. Names are unique so none collide on create (split-set names are unique). ratio_only = unique("ratio_only") ratio_seed7b = unique("ratio_seed7b") ratio_seed8 = unique("ratio_seed8") exclusive = unique("exclusive") by_year = unique("by_year") all_split_sets = [ratio_only, ratio_seed7b, ratio_seed8, exclusive, by_year] split_names = list(RATIOS) try: step(f"seed {store!r}: {len(SAMPLES)} samples ({len(SPEAKERS)} speakers × {SPEAKER_UTTERANCES})") seed_source_store(client, store, samples=SAMPLES) step(f"freeze it into dataset {dataset!r} (splits are cut over a dataset, not a workset)") build_dataset( client, dataset, sources=[SourceSpec(sample_store=store, keys=["audio.wav", "speaker.txt", "locale.txt", "year.cls"])], key_groups=[ KeyGroupSpec(name="audio", keys=["audio.wav"]), KeyGroupSpec(name="speaker", keys=["speaker.txt"]), KeyGroupSpec(name="locale", keys=["locale.txt"]), KeyGroupSpec(name="year", keys=["year.cls"]), ], shard_plan=ShardPlan(shard_size=8), ) # 1. DISCOVER — which keys can drive rules, and how each resolves. Only scalar keys # (.txt categorical, .cls numeric) are eligible; the payload audio.wav is not. step("discover split fields: which keys can drive rules") fields = client.get_dataset_split_fields(dataset) for field in fields.fields: kinds = ", ".join(field.supported_rule_kinds) or "—" print(f" {field.key:<12} {field.kind:<11} present={field.present_count:<3} → rules: {kinds}") # 2. THE LEAK, MEASURED — ratio-only vs ratio+exclusivity, over the same seed. step("cut a ratio-ONLY split (no rules): watch speakers scatter across splits") client.create_dataset_split_set(dataset, ratio_only_spec(ratio_only, seed=7)) wait_for_split_set(client, dataset, ratio_only) leaked = speakers_spanning_splits(client, ratio_only, split_names) print(f" → {len(leaked)}/{len(SPEAKERS)} speakers leak across splits (train↔val↔test): " f"{', '.join(sorted(leaked))}") step("now add EXCLUSIVITY on speaker.txt (+ stratify locale.txt) — preview, then create") excl_request = exclusive_spec(exclusive, seed=7) print_preview(client, dataset, excl_request) client.create_dataset_split_set(dataset, excl_request) detail = wait_for_split_set(client, dataset, exclusive) leaked = speakers_spanning_splits(client, exclusive, split_names) print(f" → {len(leaked)} speakers leak (exclusivity keeps each speaker whole)") for child in detail.datasets: print(f" {child.split_name:<5} n={child.sample_count} locales={child.stratify_summary.get('locale.txt', {})}") # 3. DETERMINISM — reproducibility is per (parent dataset, spec). The seed is the one # knob, and it moves the *random* planner (no exclusivity/stratification): same seed → # identical membership, different seed → a different partition. Record the seed. step("determinism: re-cut ratio-only with the SAME seed (7), and once with a DIFFERENT seed (8)") client.create_dataset_split_set(dataset, ratio_only_spec(ratio_seed7b, seed=7)) client.create_dataset_split_set(dataset, ratio_only_spec(ratio_seed8, seed=8)) wait_for_split_set(client, dataset, ratio_seed7b) wait_for_split_set(client, dataset, ratio_seed8) same = all(member_ids(client, f"{ratio_only}/{s}") == member_ids(client, f"{ratio_seed7b}/{s}") for s in split_names) diff = any(member_ids(client, f"{ratio_only}/{s}") != member_ids(client, f"{ratio_seed8}/{s}") for s in split_names) print(f" → same seed → identical membership: {same}; different seed → different membership: {diff}") print(" (exclusivity/stratification use a sorted-greedy planner instead — just as " "reproducible, but the seed does not move it)") # 4. SELECTORS — a split need not be a random share. An ATTRIBUTE selector carves out # samples by value: hold out the newest recordings as a fixed test set, then split the # remainder by ratio. shuffle is allowed on an attribute split (it has no seeded order # of its own); on a ratio split it is rejected — ratio order already comes from the seed. step("attribute selector: test = the newest year (year.cls ≥ 2024); train/val split the rest") try: SplitOutputSpec(name="train", selector=SplitRatioSelector(ratio=0.6), shuffle=True) except Exception: print(" (shuffle on a ratio split is rejected — its order comes from assignment_seed)") client.create_dataset_split_set( dataset, DatasetSplitSetCreateRequest( name=by_year, assignment_seed=7, splits=[ SplitOutputSpec( name="test", selector=SplitAttributeSelector( name="year.cls", predicate=SplitNumericRangePredicate(gte=2024), ), shuffle=True, # attribute splits may shuffle their output order ), SplitOutputSpec(name="train", selector=SplitRatioSelector(ratio=0.75)), SplitOutputSpec(name="val", selector=SplitRatioSelector(ratio=0.25)), ], shard_plan=ShardPlan(shard_size=8), ), ) detail = wait_for_split_set(client, dataset, by_year) sizes = {d.split_name: d.sample_count for d in detail.datasets} print(f" → sizes {sizes}: test is exactly the year-2024 samples; train/val share the rest") finally: # Teardown runs downward: split sets (archive → delete, cascading children) → parent # dataset → store. A split set blocks its parent dataset's deletion until it is gone. cleanup_split_sets(client, dataset, *all_split_sets) cleanup_datasets(client, dataset) cleanup(client, store) if __name__ == "__main__": main()