Upload evaluate_dataset.py with huggingface_hub

77506e9 verified 6 days ago

14.5 kB

"""Run the pipeline over a HuggingFace speech dataset and score it.
    python scripts/evaluate_dataset.py --dataset espnet/yodas-granary \
        --data_files data/en001/asr_only/00000000.parquet \
        --split train --count 3000 --save ours_en001.jsonl
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 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)")
    parser.add_argument("--no_trim", action="store_true",
                        help="keep the quiet lead-in and tail. Dataset clips "
                             "are usually trimmed already, so this is the "
                             "closer match to the published runs")
    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")

    try:
        for index, row in enumerate(data, 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, trim=not args.no_trim
            )

            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)
                print(f"[{index:>5}] WER {100 * stats['wer']:5.1f}%  "
                      f"running {running:5.1f}%  | {result.text[:50]}",
                      flush=True)
            else:
                print(f"[{index:>5}] {result.text[:60]}", 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()

    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


if __name__ == "__main__":
    raise SystemExit(main())

	"""Run the pipeline over a HuggingFace speech dataset and score it.

	python scripts/evaluate_dataset.py --dataset espnet/yodas-granary \
	--data_files data/en001/asr_only/00000000.parquet \
	--split train --count 3000 --save ours_en001.jsonl

	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 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)")
	parser.add_argument("--no_trim", action="store_true",
	help="keep the quiet lead-in and tail. Dataset clips "
	"are usually trimmed already, so this is the "
	"closer match to the published runs")
	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")

	try:
	for index, row in enumerate(data, 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, trim=not args.no_trim
	)

	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)
	print(f"[{index:>5}] WER {100 * stats['wer']:5.1f}% "
	f"running {running:5.1f}% \| {result.text[:50]}",
	flush=True)
	else:
	print(f"[{index:>5}] {result.text[:60]}", 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()

	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


	if __name__ == "__main__":
	raise SystemExit(main())