# /// script # requires-python = ">=3.10" # dependencies = [ # "transformers", # "torch", # "numpy", # "scipy", # "safetensors", # "amfm-decompy==1.0.12.2", # "praat-parselmouth", # "datasets", # "jiwer", # "soundfile", # ] # /// """Run the pipeline over a HuggingFace speech dataset and score it. python evaluate_dataset.py --dataset espnet/yodas-granary \ --data_files data/en001/asr_only/00000000.parquet \ --split train --count 3000 --save ours_en001.jsonl The dependency block above is PEP 723 metadata, so this file also runs with no setup at all -- `uv run evaluate_dataset.py ...` installs what it needs into a throwaway environment, and `hf jobs uv run ` does the same on HuggingFace's own GPUs, downloading nothing locally. Each utterance goes audio -> f0 + intensity -> conformer -> mel -> Whisper, and the hypothesis is scored against the dataset's own transcript. A microphone tells you the pipeline runs; a dataset tells you whether it is right, because the reference is there to compare against. WER is reported over the corpus rather than averaged per utterance: it is total errors over total reference words, so averaging would weight a three-word reference the same as a thirty-word one and give a different, wrong number. The normalisation matches the published inference runs -- punctuation and case alone would otherwise count as errors. """ from __future__ import annotations import argparse import json import os import sys from pathlib import Path import numpy as np # Columns HuggingFace speech datasets use for the reference transcript, in # the order worth trying. No single name is standard: yodas uses "text", # common_voice "sentence", and several others "transcription". TEXT_COLUMNS = ( "text", "sentence", "transcription", "transcript", "normalized_text", ) def find_text_column(columns) -> str | None: for name in TEXT_COLUMNS: if name in (columns or []): return name return None def wer_transform(): """The published runs' WER normalisation, copied exactly. The number is only comparable with them if this matches: removing punctuation and folding case is the difference between a WER that measures the model and one that measures its typography. """ import jiwer return jiwer.Compose([ jiwer.RemovePunctuation(), jiwer.ToLowerCase(), jiwer.RemoveMultipleSpaces(), jiwer.Strip(), jiwer.ReduceToListOfListOfWords(), ]) def word_errors(reference: str, hypothesis: str) -> dict: """Per-utterance WER plus the substitution/deletion/insertion split. The breakdown is worth having: on this pipeline a hypothesis is usually the right *length* with the wrong words, which is substitutions, and that reads very differently from a model dropping or inventing half the utterance. """ import jiwer transform = wer_transform() out = jiwer.process_words( reference, hypothesis, reference_transform=transform, hypothesis_transform=transform, ) return { "wer": out.wer, "substitutions": out.substitutions, "deletions": out.deletions, "insertions": out.insertions, "hits": out.hits, } def normalise(text: str) -> str: """Enough to compare two transcripts without calling a near-miss a miss.""" return text.strip().lower().strip(".,!?;:\"'") def audio_from_row(row, column: str = "audio", audio_dir=None): """(audio, sample_rate) from a datasets row, as mono float64. Handles the shapes the column arrives in: the {"array", "sampling_rate"} dict older `datasets` gives, the torchcodec AudioDecoder 5.x returns, and a bare filename -- which is how the yodas splits store it, with the audio in a sibling directory rather than inside the dataset. """ item = row[column] # Audio(decode=False) hands over {"bytes": ..., "path": ...}: the file's # own bytes, which soundfile reads without any system codec. if isinstance(item, dict) and ("bytes" in item or "path" in item): import io import soundfile source = (io.BytesIO(item["bytes"]) if item.get("bytes") else str(item["path"])) audio, sr = soundfile.read(source, dtype="float64", always_2d=False) if getattr(audio, "ndim", 1) > 1: audio = audio.mean(axis=1) return np.asarray(audio, dtype=np.float64), int(sr) if isinstance(item, str): path = Path(item) if not path.is_absolute() and audio_dir is not None: path = Path(audio_dir) / path if not path.exists(): raise FileNotFoundError( f"{path} not found. This dataset names its audio by " f"filename; pass --audio_dir pointing at the directory " f"holding them." ) import soundfile audio, sr = soundfile.read(str(path), dtype="float64", always_2d=False) if getattr(audio, "ndim", 1) > 1: audio = audio.mean(axis=1) return np.asarray(audio, dtype=np.float64), int(sr) # datasets 5.x: a decoder object rather than the samples. if hasattr(item, "get_all_samples"): samples = item.get_all_samples() audio = np.asarray(samples.data, dtype=np.float64) if audio.ndim > 1: audio = audio.mean(axis=0 if audio.shape[0] <= 2 else 1) return audio, int(samples.sample_rate) audio = np.asarray(item["array"], dtype=np.float64) if audio.ndim > 1: audio = audio.mean(axis=1) return audio, int(item["sampling_rate"]) def load(args): """The dataset, however it was asked for.""" from datasets import Audio, load_dataset, load_from_disk print(f"loading {args.dataset} (split={args.split})...", flush=True) kwargs = {"split": args.split} if args.config: kwargs["name"] = args.config if args.data_files: kwargs["data_files"] = args.data_files # Streaming by default: these sets run to hundreds of gigabytes, and # scoring a few thousand utterances needs the first few thousand. if args.stream: kwargs["streaming"] = True local = Path(args.dataset) # A save_to_disk tree has to go through load_from_disk; load_dataset # globs the directory and tries to merge the .indices files beside the # data, which fails on mismatched schemas rather than saying what is # wrong. if (local / "dataset_info.json").exists(): data = load_from_disk(str(local)) print(f" from disk: {len(data)} rows") else: data = load_dataset(args.dataset, **kwargs) # Decode the audio here rather than letting `datasets` do it. # # Its 5.x decoder goes through torchcodec, which dlopens system FFmpeg # shared libraries -- so a plain `pip install` leaves it raising # "libavutil.so not found" on any machine without FFmpeg development # libraries, which is most of them. soundfile reads the bytes directly # and ships its own libsndfile, so it works from a wheel alone. # # Decoding at the file's own rate also matters: asking for 16 kHz hands # the resampling to whichever backend is active, and which filter runs # changes the samples enough to move yaapt's voicing decisions at gap # edges. extract_features resamples itself. columns = data.column_names or [] if args.audio_column in columns: try: data = data.cast_column(args.audio_column, Audio(decode=False)) except Exception: # An older `datasets`, or a column that is already undecoded. pass return data def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument("--dataset", required=True, help="HuggingFace dataset id, or a local " "save_to_disk directory") parser.add_argument("--config", default=None, help="dataset config name") parser.add_argument("--split", default="train") parser.add_argument("--data_files", default=None, help="a single file inside the repo, e.g. " "data/en001/asr_only/00000000.parquet") parser.add_argument("--audio_column", default="audio") parser.add_argument("--text_column", default=None, help=f"reference transcript; tried in order from " f"{', '.join(TEXT_COLUMNS)}") parser.add_argument("--id_column", default=None, help="carried into --save so a diff is a join on id") parser.add_argument("--audio_dir", type=Path, default=None, help="where the audio lives, when the column holds " "filenames rather than samples") parser.add_argument("--skip", type=int, default=0) parser.add_argument("--count", type=int, default=0, help="stop after this many (0 = the whole split)") parser.add_argument("--no_stream", dest="stream", action="store_false", help="download the split instead of streaming it") parser.add_argument("--save", type=Path, default=None, help="write one JSON object per utterance here") parser.add_argument("--model", default="Daporte/prosody2text-weights", help="model repo id, or a local directory holding a " "checkout of one") parser.add_argument("--fp16", action="store_true", help="half precision: ~1.8x faster on a GPU, and it " "transcribed every template identically") parser.add_argument("--cpu", action="store_true", help="force the conformer onto CPU") parser.add_argument("--beams", type=int, default=None, help="beam width (default 10, as published)") args = parser.parse_args(argv) try: import jiwer # noqa: F401 except ImportError: print("jiwer is needed for scoring: pip install jiwer", file=sys.stderr) return 1 from transformers import AutoModel try: data = load(args) except Exception as error: print(f"could not load {args.dataset}: {error}", file=sys.stderr) return 1 columns = data.column_names or [] text_column = args.text_column or find_text_column(columns) print(f" columns: {', '.join(columns) if columns else '(streaming)'}") if text_column: print(f" reference from {text_column!r}") else: print(" no transcript column found -- predictions will be unscored") if args.skip: print(f" skipping the first {args.skip}") data = (data.skip(args.skip) if args.stream else data.select(range(args.skip, len(data)))) print(f"loading {args.model}...", flush=True) model = AutoModel.from_pretrained( args.model, trust_remote_code=True, half=args.fp16, cpu=args.cpu ) if args.beams: model.config.num_beams = args.beams print(f" {model.device_summary}", flush=True) transform = wer_transform() references: list[str] = [] predictions: list[str] = [] exact = scored = 0 handle = None if args.save: # Make the directory rather than failing on it. This opens *after* # the models have loaded, so a missing parent threw away a 3 GB load # and several seconds for a mkdir. args.save.parent.mkdir(parents=True, exist_ok=True) handle = open(args.save, "w", encoding="utf-8") # Kept so it can be closed on the way out. Breaking out of a streaming # dataset leaves its prefetch thread mid-request, and if the process then # exits, that thread touches the GIL while the interpreter is finalizing: # "Fatal Python error: PyGILState_Release". The scores are already # printed by then, so it is a crash banner after a successful run -- # harmless, and alarming in a way that reads as broken. rows = iter(data) try: for index, row in enumerate(rows, start=args.skip): if args.audio_column not in row: print(f"no {args.audio_column!r} column; available: " f"{', '.join(row)}", file=sys.stderr) return 1 reference = row.get(text_column) if text_column else None try: audio, sr = audio_from_row( row, args.audio_column, row.get("_audio_base_path") or args.audio_dir, ) except Exception as error: print(f"[{index:>5}] skipped: {error}", file=sys.stderr) continue result = model.transcribe(audio, sr) if reference and result.ranked: scored += 1 exact += normalise(result.text) == normalise(reference) references.append(reference) predictions.append(result.text) import jiwer running = 100 * jiwer.wer( references, predictions, reference_transform=transform, hypothesis_transform=transform, ) stats = word_errors(reference, result.text) # Reference, hypothesis, then the score: the two transcripts # adjacent and aligned so they can be read against each # other, and the number last, where it reads as the verdict # on what is just above it rather than a header for it. # # A WER on its own says how wrong; only the two lines # together say in what way -- and on this pipeline that is # the interesting part, since the words are usually the # right shape and the wrong content. print(f"[{index:>5}] ref {reference}\n" f" ours {result.text}\n" f" WER {100 * stats['wer']:5.1f}% " f"running {running:5.1f}%", flush=True) else: print(f"[{index:>5}] {result.text}", flush=True) if handle is not None: # Same field names and order as the published runs' # predictions jsonl, so a comparison is a join on id rather # than a translation. handle.write(json.dumps({ "prediction": result.text, "candidates": [text for text, _ in result.ranked], "reference": reference or "", "id": (str(row[args.id_column]) if args.id_column else str(index)), }) + "\n") handle.flush() if args.count and index + 1 - args.skip >= args.count: break except KeyboardInterrupt: print("\ninterrupted", file=sys.stderr) finally: if handle is not None: handle.close() # Shut the prefetch thread down while the interpreter is still alive. close = getattr(rows, "close", None) if close is not None: try: close() except Exception: pass if scored: import jiwer final = 100 * jiwer.wer( references, predictions, reference_transform=transform, hypothesis_transform=transform, ) print(f"\n{scored} utterance(s) scored") print(f" WER {final:.2f}%") print(f" exact match {exact}/{scored} ({100 * exact / scored:.1f}%)") else: print("\nnothing scored") if args.save: print(f" written to {args.save}") return 0 def _exit(code: int) -> None: """Leave without running interpreter finalization. A streaming dataset keeps a prefetch thread (threading's "Dummy-1") alive, and closing the iterator does not retire it -- measured: the thread is still in threading.enumerate() at atexit. When the interpreter then finalizes, that thread touches the GIL and the process dies with "Fatal Python error: PyGILState_Release" and SIGABRT, status 134. That happens after the scores are printed and the file is flushed, so the run has in fact succeeded -- but a 134 makes every `&&`, Makefile and CI job treat it as a failure. os._exit skips finalization entirely, so the status is the one this function was given. Everything that needs flushing has been flushed by here. """ sys.stdout.flush() sys.stderr.flush() os._exit(code) if __name__ == "__main__": _exit(main())