Blame for evaluate_dataset.py

"""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())