From 2d317001decc4dae202840d6eefd604ce2ed26c3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ay=C5=9Feg=C3=BCl=20Orhon?= Date: Sun, 26 Jul 2026 01:18:24 +0300 Subject: [PATCH 1/3] Split a speech-generation sweep into independently retryable shards MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The reference-length voice-clone sweep (~718 samples) takes ~26h on one Mac Studio, which no single CI job can hold: GitHub cancels a job at the default 360-minute timeout, and a runner disconnect anywhere in those hours loses the whole run. Nothing in the pipeline could evaluate part of a dataset — only num_samples, which always takes the first N. Add --dataset-shard INDEX/TOTAL so one alias can be swept as several jobs. Shards are interleaved rather than contiguous so that per-sample cost (here reference length, which spans 5s to 550s) spreads evenly instead of piling into the last shard. Membership depends only on INDEX/TOTAL, so a shard that fails can be retried by itself. Two things were needed to make those jobs safe to run against one results repo: - --hf-results-chunk-tag, because the sink resolves its next chunk index once at startup. Concurrent writers would all pick the same index and silently overwrite each other's uploads; distinct tags keep filenames disjoint. - --skip-completed-in, to resume past samples that already have results instead of recomputing them. Ids are read column-only over HTTP, so scanning existing results costs megabytes rather than downloading the embedded audio. Also write one row group per sample, so a reader can range-request a single sample's audio out of a shard instead of pulling the whole file — one long reference clip is already tens of MB. And carry per-sample wsim min / max / min-start in the windowed-SIM detail (keys outside metric_components(), so run totals are untouched), which locates drift in a clip without a separate backfill pass. Co-authored-by: Cursor --- src/openbench/cli/commands/evaluate.py | 79 +++++++++++++++++++ src/openbench/dataset/dataset_base.py | 59 +++++++++++--- .../speaker_similarity/windowed_sim_metric.py | 32 ++++++-- src/openbench/runner/benchmark.py | 1 + src/openbench/runner/config.py | 8 ++ .../runner/speech_generation_sink.py | 74 +++++++++++++++-- 6 files changed, 226 insertions(+), 27 deletions(-) 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..1e9d4ee 100644 --- a/src/openbench/dataset/dataset_base.py +++ b/src/openbench/dataset/dataset_base.py @@ -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" ) @@ -83,25 +105,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..28d8b2b 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]: diff --git a/src/openbench/runner/config.py b/src/openbench/runner/config.py index 1384038..7d1b8ab 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=( diff --git a/src/openbench/runner/speech_generation_sink.py b/src/openbench/runner/speech_generation_sink.py index 9998bba..c6b5c17 100644 --- a/src/openbench/runner/speech_generation_sink.py +++ b/src/openbench/runner/speech_generation_sink.py @@ -17,6 +17,7 @@ import re import tempfile +from collections.abc import Iterable from pathlib import Path from argmaxtools.utils import get_logger @@ -28,15 +29,30 @@ # 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._-]+") 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 +62,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: @@ -95,12 +112,16 @@ def flush(self) -> None: 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 +135,40 @@ 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. + """ + import pyarrow.parquet as pq + from huggingface_hub import HfFileSystem + from huggingface_hub.utils import HfHubHTTPError + + fs = HfFileSystem() + completed: set[str] = set() + for repo_id in repo_ids: + repo_id = repo_id.strip() + if not repo_id: + continue + try: + shards = fs.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 + for shard in shards: + with fs.open(shard, "rb") as handle: + table = pq.read_table(handle, columns=[column]) + ids = {str(v) for v in table.column(column).to_pylist() if v is not None} + 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 From c6cd5a7d304506240f001a01165788e866a36f5c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ay=C5=9Feg=C3=BCl=20Orhon?= Date: Sun, 26 Jul 2026 01:32:23 +0300 Subject: [PATCH 2/3] Keep a resumed sweep's config JSON serializable Passing --skip-completed-in put a frozenset of already-scored ids into the dataset config, and wandb serializes that config at init, so every sharded run died with "Object of type frozenset is not JSON serializable" after it had already loaded the dataset and read the existing results. Serialize the id set as a sorted list, and log only how many ids were skipped rather than the ids themselves, which would otherwise swamp the run config. The in-memory value stays a set so per-row membership checks stay O(1). Co-authored-by: Cursor --- src/openbench/dataset/dataset_base.py | 7 ++++++- src/openbench/runner/config.py | 5 +++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/src/openbench/dataset/dataset_base.py b/src/openbench/dataset/dataset_base.py index 1e9d4ee..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 @@ -66,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. diff --git a/src/openbench/runner/config.py b/src/openbench/runner/config.py index 7d1b8ab..a28538d 100644 --- a/src/openbench/runner/config.py +++ b/src/openbench/runner/config.py @@ -61,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 From fa878e46e8ef0b07359fa8cc626ed69bf93f5264 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ay=C5=9Feg=C3=BCl=20Orhon?= Date: Tue, 28 Jul 2026 00:28:27 +0300 Subject: [PATCH 3/3] Write windowed SIM columns into the speech-generation HF results sink. Previously `-m sim-windowed` was computed but only SIM/WER landed in parquet, so every sweep still needed a local wSIM backfill. Co-authored-by: Cursor --- src/openbench/runner/benchmark.py | 20 +++++++++- .../runner/speech_generation_sink.py | 39 +++++++++++++------ 2 files changed, 47 insertions(+), 12 deletions(-) diff --git a/src/openbench/runner/benchmark.py b/src/openbench/runner/benchmark.py index 28d8b2b..2a09a03 100644 --- a/src/openbench/runner/benchmark.py +++ b/src/openbench/runner/benchmark.py @@ -235,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": @@ -244,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") @@ -272,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/speech_generation_sink.py b/src/openbench/runner/speech_generation_sink.py index c6b5c17..bd7881a 100644 --- a/src/openbench/runner/speech_generation_sink.py +++ b/src/openbench/runner/speech_generation_sink.py @@ -11,13 +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 @@ -31,6 +33,7 @@ _DATA_DIR = "data" _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: @@ -102,11 +105,16 @@ 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, [] @@ -144,30 +152,39 @@ def completed_sample_ids(repo_ids: Iterable[str], column: str = "sample_idx") -> 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 - fs = HfFileSystem() 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 = fs.glob(f"datasets/{repo_id}/**/*.parquet") + 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 - for shard in shards: - with fs.open(shard, "rb") as handle: - table = pq.read_table(handle, columns=[column]) - ids = {str(v) for v in table.column(column).to_pylist() if v is not None} - found += len(ids) - completed |= ids + 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")