diff --git a/src/openbench/cli/commands/evaluate.py b/src/openbench/cli/commands/evaluate.py index 492d0d8..ca3e1b0 100644 --- a/src/openbench/cli/commands/evaluate.py +++ b/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, ) diff --git a/src/openbench/dataset/dataset_base.py b/src/openbench/dataset/dataset_base.py index 2c1c418..83e6d06 100644 --- a/src/openbench/dataset/dataset_base.py +++ b/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) diff --git a/src/openbench/metric/speaker_similarity/windowed_sim_metric.py b/src/openbench/metric/speaker_similarity/windowed_sim_metric.py index efcdad7..67049ef 100644 --- a/src/openbench/metric/speaker_similarity/windowed_sim_metric.py +++ b/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: diff --git a/src/openbench/runner/benchmark.py b/src/openbench/runner/benchmark.py index dfd99d7..2a09a03 100644 --- a/src/openbench/runner/benchmark.py +++ b/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( diff --git a/src/openbench/runner/config.py b/src/openbench/runner/config.py index 1384038..a28538d 100644 --- a/src/openbench/runner/config.py +++ b/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 diff --git a/src/openbench/runner/speech_generation_sink.py b/src/openbench/runner/speech_generation_sink.py index 9998bba..bd7881a 100644 --- a/src/openbench/runner/speech_generation_sink.py +++ b/src/openbench/runner/speech_generation_sink.py @@ -11,12 +11,15 @@ Each row carries three playable Audio columns (embedded in the parquet) for listening-based debugging — ``prompt_audio`` (the clone prompt), ``sim_reference_audio`` (the clip SIM compared the generation against), and -``generated_audio`` — plus ``prompt_text``, the ASR ``transcription``, and -per-sample ``SIM`` / ``WER``. +``generated_audio`` — plus ``prompt_text``, the ASR ``transcription``, +per-sample ``SIM`` / ``WER``, and (when ``-m sim-windowed`` is enabled) +``wsim_mean`` / ``wsim_var`` / ``wsim_min`` / ``wsim_max`` / ``wsim_min_start``. """ import re import tempfile +from collections.abc import Iterable +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from argmaxtools.utils import get_logger @@ -28,15 +31,31 @@ # Where shards live inside the repo. The HF datasets viewer auto-loads every # parquet file under `data/`, concatenating them into a single table. _DATA_DIR = "data" -_CHUNK_RE = re.compile(r"chunk-(\d+)\.parquet$") +_CHUNK_RE = re.compile(r"chunk-(\d+)(?:-[A-Za-z0-9._-]+)?\.parquet$") +_TAG_SANITIZE_RE = re.compile(r"[^A-Za-z0-9._-]+") +_RESUME_READ_WORKERS = 12 class SpeechGenerationResultSink: - """Accumulates per-sample rows and flushes parquet shards to an HF dataset.""" - - def __init__(self, repo_id: str, flush_every: int = 100, private: bool = True) -> None: + """Accumulates per-sample rows and flushes parquet shards to an HF dataset. + + The chunk counter is resolved once at startup from the shards already in the + repo, so several jobs writing to one repo concurrently would all pick the + same next index and overwrite each other's uploads. Give each concurrent + writer a distinct `chunk_tag` (e.g. its shard index) to keep filenames — and + therefore uploads — disjoint. + """ + + def __init__( + self, + repo_id: str, + flush_every: int = 100, + private: bool = True, + chunk_tag: str | None = None, + ) -> None: self.repo_id = repo_id self.flush_every = max(1, int(flush_every)) + self.chunk_tag = _TAG_SANITIZE_RE.sub("-", chunk_tag).strip("-") if chunk_tag else None self._buffer: list[dict] = [] self._api = HfApi() @@ -46,7 +65,8 @@ def __init__(self, repo_id: str, flush_every: int = 100, private: bool = True) - self._chunk_index = self._next_chunk_index() logger.info( f"Speech-generation results sink → hf://datasets/{repo_id} " - f"(flush every {self.flush_every}, starting at chunk {self._chunk_index})" + f"(flush every {self.flush_every}, starting at chunk {self._chunk_index}" + f"{f', tag {self.chunk_tag}' if self.chunk_tag else ''})" ) def _next_chunk_index(self) -> int: @@ -85,22 +105,31 @@ def flush(self) -> None: "sim_reference_audio": Audio(), "generated_audio": Audio(), # Kept adjacent for easy analysis: synthesized text, its ASR - # transcription, and the two scores. + # transcription, whole-clip SIM/WER, and windowed SIM breakdown. "prompt_text": Value("string"), "transcription": Value("string"), "WER": Value("float32"), "SIM": Value("float32"), + "wsim_mean": Value("float32"), + "wsim_var": Value("float32"), + "wsim_min": Value("float32"), + "wsim_max": Value("float32"), + "wsim_min_start": Value("float32"), } ) rows, self._buffer = self._buffer, [] dataset = Dataset.from_list(rows, features=features) - chunk_name = f"chunk-{self._chunk_index:05d}.parquet" + tag_suffix = f"-{self.chunk_tag}" if self.chunk_tag else "" + chunk_name = f"chunk-{self._chunk_index:05d}{tag_suffix}.parquet" with tempfile.TemporaryDirectory() as tmp: local_path = Path(tmp) / chunk_name # to_parquet embeds the Audio columns as bytes, so each shard is - # self-contained (no external file references). - dataset.to_parquet(str(local_path)) + # self-contained (no external file references). batch_size=1 puts + # every sample in its own row group, so a reader can range-request a + # single sample's audio instead of pulling the whole shard — one + # long reference clip is already tens of MB. + dataset.to_parquet(str(local_path), batch_size=1) self._api.upload_file( path_or_fileobj=str(local_path), path_in_repo=f"{_DATA_DIR}/{chunk_name}", @@ -114,3 +143,49 @@ def flush(self) -> None: def close(self) -> None: """Flush any remaining buffered rows.""" self.flush() + + +def completed_sample_ids(repo_ids: Iterable[str], column: str = "sample_idx") -> set[str]: + """Collect the sample ids already scored in the given HF results repos. + + Only `column` is fetched from each parquet shard, so the (much larger) + embedded audio columns are never downloaded — reading the ids of a few + hundred results costs megabytes, not gigabytes. Repos that do not exist yet + contribute nothing, which makes a first run behave like a full sweep. + + Shards are read concurrently because the cost is per-file network latency, + not bandwidth; a resumed sweep would otherwise spend a large part of its job + reading shard footers one at a time before scoring its first sample. + """ + import pyarrow.parquet as pq + from huggingface_hub import HfFileSystem + from huggingface_hub.utils import HfHubHTTPError + + completed: set[str] = set() + + def shard_ids(shard: str) -> set[str]: + # A filesystem per worker: HfFileSystem holds a session that is not + # guaranteed to be thread-safe. + with HfFileSystem().open(shard, "rb") as handle: + table = pq.read_table(handle, columns=[column]) + return {str(v) for v in table.column(column).to_pylist() if v is not None} + + for repo_id in repo_ids: + repo_id = repo_id.strip() + if not repo_id: + continue + try: + shards = HfFileSystem().glob(f"datasets/{repo_id}/**/*.parquet") + except (FileNotFoundError, HfHubHTTPError) as e: + logger.info(f"No results to resume from in {repo_id}: {e}") + continue + + found = 0 + with ThreadPoolExecutor(max_workers=_RESUME_READ_WORKERS) as pool: + for ids in pool.map(shard_ids, shards): + found += len(ids) + completed |= ids + logger.info(f"Found {found} scored samples across {len(shards)} shards in {repo_id}") + + logger.info(f"Resuming past {len(completed)} unique already-scored samples") + return completed