Add dataset sharding + resume on top of qwen-upstream-eval

src/openbench/cli/commands/evaluate.py

@@ -35,6 +35,20 @@
 PipelineConfigOptions = dict[str, dict[str, Any] | dict[str, dict[str, Any]]]
 
 
+def parse_dataset_shard(spec: str) -> tuple[int, int]:
+    """Parse a ``INDEX/TOTAL`` shard spec into 0-based (index, total)."""
+    index_str, sep, total_str = spec.partition("/")
+    if not sep:
+        raise typer.BadParameter(f"Expected INDEX/TOTAL, got {spec!r}", param_hint="--dataset-shard")
+    try:
+        index, total = int(index_str), int(total_str)
+    except ValueError:
+        raise typer.BadParameter(f"INDEX and TOTAL must be integers, got {spec!r}", param_hint="--dataset-shard")
+    if total < 1 or not 0 <= index < total:
+        raise typer.BadParameter(f"Need 0 <= INDEX < TOTAL and TOTAL >= 1, got {spec!r}", param_hint="--dataset-shard")
+    return index, total
+
+
 class EvaluationConfig(BaseModel):
     benchmark_config: BenchmarkConfig = Field(..., description="The benchmark config to use for evaluation")
     pipeline_config: dict[str, dict[str, Any]] = Field(
@@ -181,6 +195,9 @@ def run_alias_mode(
     verbose: bool,
     hf_results_repo: str | None = None,
     hf_results_flush_every: int = 100,
+    hf_results_chunk_tag: str | None = None,
+    dataset_shard: str | None = None,
+    skip_completed_in: str | None = None,
     metric_config: list[str] | None = None,
 ) -> BenchmarkResult:
     """Run evaluation using pipeline and dataset aliases."""
@@ -227,6 +244,24 @@ def run_alias_mode(
         typer.echo(f"📊 Loading dataset: {dataset_name}")
         dataset_config = DatasetRegistry.get_alias_config(dataset_name)
 
+        # Row selection for split-up sweeps: keep one shard, and/or drop samples
+        # that already have results elsewhere.
+        dataset_overrides: dict[str, Any] = {}
+        if dataset_shard:
+            shard_index, num_shards = parse_dataset_shard(dataset_shard)
+            dataset_overrides.update(num_shards=num_shards, shard_index=shard_index)
+            typer.echo(f"🔀 Dataset shard {shard_index} of {num_shards} (interleaved)")
+        if skip_completed_in:
+            from openbench.runner.speech_generation_sink import completed_sample_ids
+
+            repos = [repo for repo in (r.strip() for r in skip_completed_in.split(",")) if repo]
+            completed = completed_sample_ids(repos)
+            typer.echo(f"⏭️  Skipping {len(completed)} samples already scored in {', '.join(repos)}")
+            if completed:
+                dataset_overrides["exclude_sample_ids"] = frozenset(completed)
+        if dataset_overrides:
+            dataset_config = dataset_config.model_copy(update=dataset_overrides)
+
         wandb_config = WandbConfig(
             project_name=wandb_project,
             run_name=wandb_run_name,
@@ -262,6 +297,7 @@ def run_alias_mode(
             metrics=metric_kwargs,
             hf_results_repo=hf_results_repo,
             hf_results_flush_every=hf_results_flush_every,
+            hf_results_chunk_tag=hf_results_chunk_tag,
         )
 
         # Create runner
@@ -420,6 +456,35 @@ def evaluate(
         "--hf-results-flush-every",
         help="Flush buffered per-sample results to the HF repo every N samples (used with --hf-results-repo).",
     ),
+    hf_results_chunk_tag: str | None = typer.Option(
+        None,
+        "--hf-results-chunk-tag",
+        help=(
+            "Suffix for the uploaded parquet shard filenames, e.g. `--hf-results-chunk-tag s3`. "
+            "Required when several runs push to the same --hf-results-repo concurrently: each "
+            "picks its starting chunk index independently, so without distinct tags they would "
+            "overwrite each other's shards."
+        ),
+    ),
+    dataset_shard: str | None = typer.Option(
+        None,
+        "--dataset-shard",
+        help=(
+            "Evaluate only one interleaved shard of the dataset, given as INDEX/TOTAL with a "
+            "0-based index, e.g. `--dataset-shard 3/10`. Lets one sweep run as several independent "
+            "jobs, each short enough to finish inside a CI job's time limit. Shard membership "
+            "depends only on INDEX/TOTAL, so a single shard can be retried on its own. Alias mode only."
+        ),
+    ),
+    skip_completed_in: str | None = typer.Option(
+        None,
+        "--skip-completed-in",
+        help=(
+            "Comma-separated HF results repos to resume past: every sample already scored there is "
+            "dropped before evaluation. Pass the current --hf-results-repo plus any older repo "
+            "holding partial results. Alias mode only."
+        ),
+    ),
     verbose: bool = typer.Option(False, "--verbose", "-v", help="Enable verbose output"),
 ) -> None:
     """Run evaluation benchmarks.
@@ -467,6 +532,17 @@ def evaluate(
 
         # Validate mutually exclusive modes
         if evaluation_config_path is not None:
+            alias_only = {
+                "--dataset-shard": dataset_shard,
+                "--skip-completed-in": skip_completed_in,
+                "--hf-results-chunk-tag": hf_results_chunk_tag,
+            }
+            unsupported = [flag for flag, value in alias_only.items() if value]
+            if unsupported:
+                raise typer.BadParameter(
+                    f"{', '.join(unsupported)} only apply in alias mode; set the equivalent fields in "
+                    "the evaluation config instead."
+                )
             typer.echo("🔧 Running with config file mode")
             result = run_config_file_mode(evaluation_config_path, evaluation_config_overrides, verbose)
         else:
@@ -484,6 +560,9 @@ def evaluate(
                 pipeline_config=pipeline_config,
                 hf_results_repo=hf_results_repo,
                 hf_results_flush_every=hf_results_flush_every,
+                hf_results_chunk_tag=hf_results_chunk_tag,
+                dataset_shard=dataset_shard,
+                skip_completed_in=skip_completed_in,
                 metric_config=metric_config,
                 verbose=verbose,
             )

src/openbench/dataset/dataset_base.py

@@ -10,7 +10,7 @@
 from argmaxtools.utils import get_logger
 from datasets import Dataset as HfDataset
 from datasets import load_dataset
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, field_serializer
 
 from ..types import PredictionProtocol
 from .dataset_utils import validate_hf_dataset_schema
@@ -32,6 +32,28 @@ class DatasetConfig(BaseModel):
     num_samples: int | None = Field(
         None, description="Number of samples to take from the dataset. If None, take all samples."
     )
+    num_shards: int | None = Field(
+        None,
+        description=(
+            "Split the dataset into this many shards and keep only `shard_index`, so one sweep "
+            "can be spread over several machines or CI jobs. Shards are interleaved (row i goes "
+            "to shard i % num_shards) rather than contiguous, so per-sample cost — e.g. voice-clone "
+            "reference length — spreads evenly instead of piling into one shard."
+        ),
+    )
+    shard_index: int = Field(0, description="Which shard to keep when `num_shards` is set (0-based).")
+    exclude_sample_ids: frozenset[str] | None = Field(
+        None,
+        description=(
+            "Drop rows whose `sample_id_column` value appears in this set. Lets an interrupted "
+            "sweep resume without recomputing samples that already have results. Applied after "
+            "sharding, so shard membership does not shift as results accumulate."
+        ),
+    )
+    sample_id_column: str = Field(
+        "sample_idx",
+        description="Column holding the dataset-stable sample id matched against `exclude_sample_ids`.",
+    )
     column_mapping: Mapping[str, str] | None = Field(
         None, description="Mapping of the column names in the dataset to the expected column names in the sample class"
     )
@@ -44,6 +66,11 @@ class DatasetConfig(BaseModel):
         ),
     )
 
+    @field_serializer("exclude_sample_ids")
+    def _serialize_exclude_sample_ids(self, sample_ids: frozenset[str] | None) -> list[str] | None:
+        """Dump the id set as a sorted list, since a set is not JSON serializable."""
+        return sorted(sample_ids) if sample_ids else None
+
     def load(self) -> HfDataset:
         """Load dataset from config.
 
@@ -83,25 +110,40 @@ def _load_local(self) -> HfDataset:
             split = self.split or "test"  # Default split
             ds = load_local_dataset(dataset_dir=dataset_path, split=split)
 
-        if self.num_samples is not None:
-            ds = ds.take(self.num_samples)
-
-        if self.column_mapping is not None:
-            ds = ds.rename_columns(self.column_mapping)
-
-        if self.column_transforms is not None:
-            for col, transform in self.column_transforms.items():
-                ds = ds.map(transform)
-
-        return ds
+        return self._postprocess(ds)
 
     def _load_huggingface(self) -> HfDataset:
         """Load dataset from HuggingFace Hub."""
         # TODO: Add support for streaming datasets
         ds = load_dataset(self.dataset_id, self.subset, split=self.split)
+        return self._postprocess(ds)
+
+    def _postprocess(self, ds: HfDataset) -> HfDataset:
+        """Apply row selection (sampling, sharding, exclusions) then column fixups."""
         if self.num_samples is not None:
             ds = ds.take(self.num_samples)
 
+        if self.num_shards is not None and self.num_shards > 1:
+            if not 0 <= self.shard_index < self.num_shards:
+                raise ValueError(f"shard_index must be in [0, {self.num_shards}), got {self.shard_index}")
+            total = len(ds)
+            # contiguous=False keeps the interleaved ds[shard_index::num_shards]
+            # assignment, which balances cost across shards.
+            ds = ds.shard(num_shards=self.num_shards, index=self.shard_index, contiguous=False)
+            logger.info(f"Shard {self.shard_index}/{self.num_shards}: kept {len(ds)} of {total} rows (interleaved)")
+
+        if self.exclude_sample_ids:
+            if self.sample_id_column not in ds.column_names:
+                raise ValueError(
+                    f"exclude_sample_ids needs column {self.sample_id_column!r}, "
+                    f"but the dataset only has {ds.column_names}"
+                )
+            excluded = self.exclude_sample_ids
+            before = len(ds)
+            # input_columns keeps the filter from decoding the audio columns.
+            ds = ds.filter(lambda sample_id: str(sample_id) not in excluded, input_columns=self.sample_id_column)
+            logger.info(f"Excluded {before - len(ds)} already-completed rows, {len(ds)} left to evaluate")
+
         if self.column_mapping is not None:
             ds = ds.rename_columns(self.column_mapping)
 

src/openbench/metric/speaker_similarity/windowed_sim_metric.py

@@ -109,6 +109,14 @@ def _embed_waveform(self, wav_16k: torch.Tensor) -> torch.Tensor:
 
     def window_scores(self, generated_audio: str, reference_audio: str) -> list[float]:
         """Cosine similarity of each generated-audio window vs the whole reference."""
+        return [score for _, score in self.window_scores_with_starts(generated_audio, reference_audio)]
+
+    def window_scores_with_starts(self, generated_audio: str, reference_audio: str) -> list[tuple[float, float]]:
+        """``(window start in seconds, similarity vs the whole reference)`` per window.
+
+        The start times make a low-scoring window locatable in the generated clip,
+        which is what turns a bad score into something you can listen to.
+        """
         emb_reference = self._embed_waveform(self._load_16k(reference_audio))
         wav = self._load_16k(generated_audio)
         total = wav.shape[-1]
@@ -117,19 +125,19 @@ def window_scores(self, generated_audio: str, reference_audio: str) -> list[floa
         min_win = int(self.min_window_seconds * self.SAMPLE_RATE)
 
         starts = list(range(0, max(total - win, 0) + 1, hop))
-        windows = [wav[:, s : s + win] for s in starts]
+        windows = [(s, wav[:, s : s + win]) for s in starts]
         # Cover the tail beyond the last full window when it is long enough to
         # embed meaningfully; a clip shorter than one window is itself the only
         # window (degenerates to plain SIM).
         tail_start = starts[-1] + hop if windows else 0
         if total - tail_start >= min_win or not windows:
-            windows.append(wav[:, tail_start:])
+            windows.append((tail_start, wav[:, tail_start:]))
 
-        scores = []
-        for w in windows:
+        scored = []
+        for start, w in windows:
             emb = self._embed_waveform(w)
-            scores.append(F.cosine_similarity(emb, emb_reference).item())
-        return scores
+            scored.append((start / self.SAMPLE_RATE, F.cosine_similarity(emb, emb_reference).item()))
+        return scored
 
     # -- BaseMetric interface --------------------------------------------------
 
@@ -147,11 +155,13 @@ def compute_components(self, reference, hypothesis, **kwargs) -> Details:
                 "(or pass reference_audio=... to the metric call)."
             )
 
-        scores = self.window_scores(generated_audio, str(reference_audio))
+        scored = self.window_scores_with_starts(generated_audio, str(reference_audio))
+        scores = [score for _, score in scored]
         n = len(scores)
         mean = sum(scores) / n
         var = sum((s - mean) ** 2 for s in scores) / n
-        lo, hi = min(scores), max(scores)
+        min_start, lo = min(scored, key=lambda window: window[1])
+        hi = max(scores)
 
         self._run_window_min = lo if self._run_window_min is None else min(self._run_window_min, lo)
         self._run_window_max = hi if self._run_window_max is None else max(self._run_window_max, hi)
@@ -167,6 +177,12 @@ def compute_components(self, reference, hypothesis, **kwargs) -> Details:
             "wsim_var_sum": var,
             "window_count": float(n),
             "count": 1.0,
+            # Not declared in metric_components(), so these are never accumulated
+            # into run totals — they ride along in the per-sample detailed output
+            # (W&B task table) to make single-sample drift analyzable.
+            "wsim_min": lo,
+            "wsim_max": hi,
+            "wsim_min_start": min_start,
         }
 
     def compute_metric(self, detail: Details) -> float:

src/openbench/runner/benchmark.py

@@ -80,6 +80,7 @@ def _maybe_create_results_sink(self, pipeline: Pipeline):
         return SpeechGenerationResultSink(
             repo_id=self.config.hf_results_repo,
             flush_every=self.config.hf_results_flush_every,
+            chunk_tag=self.config.hf_results_chunk_tag,
         )
 
     def _get_metrics(self, pipeline: Pipeline) -> dict[str, BaseMetric]:
@@ -234,6 +235,9 @@ def _build_speech_generation_row(self, sample, output, task_results, sample_id)
 
         sim = wer = None
         transcription = ""
+        # Windowed SIM columns stay None when `-m sim-windowed` was not requested,
+        # so the parquet schema stays stable across runs that omit it.
+        wsim_mean = wsim_var = wsim_min = wsim_max = wsim_min_start = None
         for t in task_results:
             name = self._normalize_metric_name(t.metric_name)
             if name == "sim":
@@ -243,6 +247,16 @@ def _build_speech_generation_row(self, sample, output, task_results, sample_id)
                 # The ASR transcription used to compute WER rides along in the
                 # metric's per-sample detailed output (see speech_generation_wer).
                 transcription = (t.detailed_result or {}).get("transcription") or ""
+            elif name == "sim-windowed":
+                detail = t.detailed_result or {}
+                # Metric value is the per-sample windowed mean; within-sample
+                # variance rides as wsim_var_sum (same column names as the
+                # standalone wSIM backfill).
+                wsim_mean = t.result
+                wsim_var = detail.get("wsim_var_sum")
+                wsim_min = detail.get("wsim_min")
+                wsim_max = detail.get("wsim_max")
+                wsim_min_start = detail.get("wsim_min_start")
 
         try:
             gen_array, gen_sr = sf.read(output.prediction.audio_path, dtype="float32")
@@ -271,11 +285,16 @@ def _build_speech_generation_row(self, sample, output, task_results, sample_id)
             "prompt_audio": prompt,
             "sim_reference_audio": sim_reference,
             "generated_audio": {"array": gen_array, "sampling_rate": int(gen_sr)},
-            # prompt_text / transcription / WER / SIM kept adjacent for analysis.
+            # prompt_text / transcription / WER / SIM / wSIM kept adjacent for analysis.
             "prompt_text": sample.text,
             "transcription": transcription,
             "WER": wer,
             "SIM": sim,
+            "wsim_mean": wsim_mean,
+            "wsim_var": wsim_var,
+            "wsim_min": wsim_min,
+            "wsim_max": wsim_max,
+            "wsim_min_start": wsim_min_start,
         }
 
     def _run_pipeline_on_dataset_parallel(

src/openbench/runner/config.py

@@ -37,6 +37,14 @@ class BenchmarkConfig(BaseModel):
     hf_results_flush_every: int = Field(
         100, description="Flush buffered per-sample results to the HF repo every N samples."
     )
+    hf_results_chunk_tag: str | None = Field(
+        None,
+        description=(
+            "Suffix appended to the uploaded parquet shard names. Required when several runs "
+            "(e.g. dataset shards) push to one repo at the same time, since each picks its "
+            "starting chunk index independently and identical names would overwrite."
+        ),
+    )
     continue_on_sample_error: bool = Field(
         True,
         description=(
@@ -53,4 +61,9 @@ def get_wandb_config_to_log(self) -> dict[str, Any]:
         wandb_config: dict[str, Any] = self.model_dump()
         # Convert `metrics` that use enums to their respective values
         wandb_config["metrics"] = {metric.value: kwargs for metric, kwargs in wandb_config["metrics"].items()}
+        # A resumed sweep excludes every already-scored id; log how many were
+        # skipped instead of listing them, which would swamp the run config.
+        for dataset in wandb_config.get("datasets", {}).values():
+            if dataset.get("exclude_sample_ids") is not None:
+                dataset["exclude_sample_ids"] = f"{len(dataset['exclude_sample_ids'])} already-scored ids"
         return wandb_config