run_eval.py

"""
Open ASR Leaderboard entry: MOSS-Transcribe-preview-2B.

Pipeline:
  - librosa.load @ 16 kHz
  - WhisperFeatureExtractor log-mel (n_fft=400, hop=160, dim=128)
  - MossProcessor with a single default chat template (identical across all
    datasets, no per-dataset style switching)
  - MossForCausalLM.generate(greedy, num_beams=1, eos=processor.end_token_id)
"""
import argparse
import os

import librosa
import numpy as np
import torch
from tqdm import tqdm
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BatchEncoding,
)
from transformers.dynamic_module_utils import get_class_from_dynamic_module
from transformers.models.whisper.feature_extraction_whisper import WhisperFeatureExtractor

import evaluate
from normalizer import data_utils

wer_metric = evaluate.load("wer")

SAMPLE_RATE = 16000
MEL_DIM = 128
N_FFT = 400
HOP = 160


def build_processor(model_id, tokenizer, revision=None):
    """Reproduce WhisperLogMelMossProcessor.__call__ behavior."""
    MossProcessor = get_class_from_dynamic_module(
        "processing_Moss.MossProcessor", model_id, revision=revision
    )
    MelConfig = get_class_from_dynamic_module(
        "processing_Moss.MelConfig", model_id, revision=revision
    )
    mel_cfg = MelConfig(mel_sr=SAMPLE_RATE, mel_dim=MEL_DIM, mel_n_fft=N_FFT, mel_hop_length=HOP)
    processor = MossProcessor(tokenizer, config=mel_cfg, enable_time_marker=False)
    template_path = os.path.join(os.path.dirname(__file__), "chat_template_default.py")
    processor.load_template(template_path)

    fe = WhisperFeatureExtractor(
        feature_size=MEL_DIM, sampling_rate=SAMPLE_RATE, hop_length=HOP, n_fft=N_FFT
    )
    return processor, fe


def encode_one(audio_np: np.ndarray, processor, fe):
    """Single-sample encode: returns (input_ids[L], audio_mask[L], mel[128,T], T)."""
    wav = audio_np.astype(np.float32)
    try:
        mel = fe._np_extract_fbank_features(wav[None, ...], device="cpu")[0]
    except TypeError:
        mel = fe._np_extract_fbank_features(wav[None, ...])[0]
    mel = torch.from_numpy(mel).to(processor.config.mel_dtype)
    if mel.dim() == 3:
        mel = mel.squeeze(0)

    T = mel.shape[-1]
    num_audio_tokens = processor._get_feat_extract_output_lengths(T)
    if processor.chat_template is not None:
        ids, mask = processor._build_input_from_template(num_audio_tokens)
    else:
        ids, mask = processor._build_input_legacy(num_audio_tokens)
    return ids, mask, mel, T


def encode_batch(audios, processor, fe, device, model_dtype) -> BatchEncoding:
    """Concat mel along time dim, left-pad input_ids per sample for batched generate."""
    encs = [encode_one(a, processor, fe) for a in audios]
    L = max(len(e[0]) for e in encs)
    bs = len(encs)
    pad_id = processor.tokenizer.pad_token_id or processor.end_token_id

    input_ids = torch.full((bs, L), pad_id, dtype=torch.long)
    attn_mask = torch.zeros((bs, L), dtype=torch.long)
    audio_mask = torch.zeros((bs, L), dtype=torch.bool)
    seq_lens = torch.zeros((bs,), dtype=torch.long)
    mel_parts = []
    for i, (ids, m, mel, T) in enumerate(encs):
        li = len(ids)
        input_ids[i, -li:]  = torch.tensor(ids, dtype=torch.long)
        attn_mask[i, -li:]  = 1
        audio_mask[i, -li:] = torch.tensor(m, dtype=torch.bool)
        seq_lens[i] = T
        mel_parts.append(mel)
    audio_data = torch.cat(mel_parts, dim=-1)

    return BatchEncoding(data={
        "input_ids": input_ids.to(device),
        "attention_mask": attn_mask.to(device),
        "audio_data": audio_data.to(device).to(model_dtype),
        "audio_data_seqlens": seq_lens.to(device),
        "audio_input_mask": audio_mask.to(device),
    })


def main(args):
    device = f"cuda:{args.device}" if args.device >= 0 else "cpu"

    model = AutoModelForCausalLM.from_pretrained(
        args.model_id, dtype="auto", trust_remote_code=True, revision=args.model_revision
    ).to(device)
    model.eval()
    print(f"Model size: {sum(p.numel() for p in model.parameters()) / 1e9:.2f}B parameters")

    tokenizer = AutoTokenizer.from_pretrained(
        args.model_id, trust_remote_code=True, revision=args.model_revision
    )
    processor, fe = build_processor(args.model_id, tokenizer, revision=args.model_revision)
    eos_ids = [processor.end_token_id]

    def benchmark(batch):
        # Open ASR Leaderboard hands us audio dicts at 16 kHz already; resample defensively.
        sr_in = batch["audio"][0]["sampling_rate"]
        audios = []
        for a in batch["audio"]:
            wav = a["array"]
            if sr_in != SAMPLE_RATE:
                wav = librosa.resample(wav.astype(np.float32), orig_sr=sr_in, target_sr=SAMPLE_RATE)
            audios.append(wav)
        batch["audio_length_s"] = [len(w) / SAMPLE_RATE for w in audios]
        mb = len(audios)
        batch["audio_filepath"] = data_utils.extract_audio_filepaths_from_batch(batch, mb)

        torch.cuda.synchronize(device=device)
        start_event = torch.cuda.Event(enable_timing=True)
        end_event = torch.cuda.Event(enable_timing=True)
        start_event.record()

        inputs = encode_batch(audios, processor, fe, device, model.dtype)
        with torch.no_grad():
            gen_ids = model.generate(
                **inputs,
                max_new_tokens=args.max_new_tokens,
                do_sample=False,
                num_beams=1,
                use_cache=True,
                eos_token_id=eos_ids,
            )
        new_ids = gen_ids[:, inputs["input_ids"].shape[1]:]
        preds = [t.strip() for t in processor.batch_decode(new_ids, skip_special_tokens=True)]

        end_event.record()
        torch.cuda.synchronize(device=device)
        runtime = start_event.elapsed_time(end_event) / 1000.0
        batch["transcription_time_s"] = mb * [runtime / mb]
        # Save raw (unnormalized) outputs; normalization is applied at scoring time.
        batch["predictions"] = preds
        batch["references"] = batch["original_text"]
        return batch

    # ---- Warm-up ----
    if args.warmup_steps:
        warmup = data_utils.load_data(args)
        warmup = data_utils.prepare_data(warmup)
        n = args.warmup_steps * args.batch_size
        warmup = warmup.take(n) if args.streaming else warmup.select(range(min(n, len(warmup))))
        for _ in tqdm(iter(warmup.map(benchmark, batch_size=args.batch_size, batched=True)), desc="Warmup"):
            pass

    # ---- Eval ----
    ds = data_utils.load_data(args)
    ds = data_utils.prepare_data(ds)
    if args.max_eval_samples and args.max_eval_samples > 0:
        ds = ds.take(args.max_eval_samples) if args.streaming else ds.select(
            range(min(args.max_eval_samples, len(ds)))
        )
    ds = ds.map(benchmark, batch_size=args.batch_size, batched=True, remove_columns=["audio"])

    res = {"audio_length_s": [], "transcription_time_s": [], "predictions": [], "references": [], "audio_filepath": []}
    for r in tqdm(iter(ds), desc="Samples"):
        for k in res:
            res[k].append(r[k])

    manifest = data_utils.write_manifest(
        res["references"], res["predictions"], args.model_id,
        args.dataset_path, args.dataset, args.split,
        audio_length=res["audio_length_s"], transcription_time=res["transcription_time_s"],
        audio_filepaths=res["audio_filepath"],
    )
    print("Manifest:", os.path.abspath(manifest))

    # Normalize raw references/predictions at scoring time.
    norm_refs = [data_utils.normalizer(r) for r in res["references"]]
    norm_preds = [data_utils.normalizer(p) for p in res["predictions"]]
    wer = round(100 * wer_metric.compute(references=norm_refs, predictions=norm_preds), 2)
    rtfx = round(sum(res["audio_length_s"]) / sum(res["transcription_time_s"]), 2)
    print(f"WER: {wer}%  RTFx: {rtfx}")


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--model_id", type=str, required=True)
    p.add_argument("--model_revision", type=str, default="c4b3988677df13c14e79d9db59f356ed761db366",
                   help="Pin the model repo revision (commit) for trust_remote_code stability.")
    p.add_argument("--dataset_path", type=str, default="hf-audio/esb-datasets-test-only-sorted")
    p.add_argument("--dataset", type=str, required=True)
    p.add_argument("--split", type=str, default="test")
    p.add_argument("--device", type=int, default=0)
    p.add_argument("--batch_size", type=int, default=1)
    p.add_argument("--max_new_tokens", type=int, default=1024)
    p.add_argument("--max_eval_samples", type=int, default=None)
    p.add_argument("--streaming", action="store_true", help="Stream the dataset instead of downloading it.")
    p.add_argument("--warmup_steps", type=int, default=10)
    args = p.parse_args()
    main(args)