Update run_eval.py

run_eval.py

@@ -1,276 +1,100 @@
+"""Evaluation script for Higgs Audio v3 models on ESB benchmark datasets.
+
+No external dependencies beyond transformers + torch. The model bundles
+its own audio preprocessing via trust_remote_code=True.
+
+Usage:
+    python run_eval_higgs_audio.py \
+        --model_id bosonai/higgs-audio-v3-8b-stt \
+        --dataset_path hf-audio/esb-datasets-test-only-sorted \
+        --dataset ami --split test --device 0 --batch_size 4
+"""
+
 import argparse
 import os
-import re
+import sys
+import time
+import runpy
+
 import torch
-from torch.nn.attention import sdpa_kernel, SDPBackend
-from transformers import AutoConfig, AutoModelForSpeechSeq2Seq, AutoModelForMultimodalLM, AutoModelForCTC, AutoProcessor, MODEL_FOR_MULTIMODAL_LM_MAPPING, MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING, MODEL_FOR_CTC_MAPPING, CompileConfig
 import evaluate
 from normalizer import data_utils
+from transformers import AutoModel, AutoTokenizer
 from tqdm import tqdm
-import random
-import numpy as np
 
 wer_metric = evaluate.load("wer")
-torch.set_float32_matmul_precision('high')
 
 
-def remove_brackets(text):
-    """
-    Remove parentheses from text, replacing them with spaces.
+def load_transcribe_fn(model_id):
+    """Load the bundled transcribe_batch function from the model repo.
 
-    Some models (e.g. Cohere ASR) output parentheses that would cause the
-    normalizer to delete the enclosed text entirely, leading to false
-    deletion errors in the predictions.
+    Downloads all Python files needed by transcribe.py, then loads it via
+    runpy with the download directory on sys.path so plain (non-relative)
+    imports resolve to sibling files.
     """
-    text = text.replace("(", " ").replace(")", " ")
-    text = re.sub(r'\s+', ' ', text)
-    return text
+    from transformers.utils import cached_file
+
+    for filename in [
+        "transcribe.py",
+        "higgs_audio_collator.py",
+        "modeling_higgs_audio_xcodec.py",
+        "utils.py",
+        "common.py",
+        "configuration_higgs_audio.py",
+    ]:
+        cached_file(model_id, filename)
+
+    path = cached_file(model_id, "transcribe.py")
+    module_dir = os.path.dirname(path)
+
+    sys.path.insert(0, module_dir)
+    try:
+        module_globals = runpy.run_path(path)
+    finally:
+        sys.path.pop(0)
+
+    return module_globals["transcribe_batch"]
 
 
 def main(args):
+    device = f"cuda:{args.device}" if args.device >= 0 else "cpu"
 
-    # Set seed due to randomness in some models (e.g. VibeVoice's acoustic tokenizer sampling)
-    seed = 42
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    torch.backends.cudnn.deterministic = True
-
-    torch_dtype = getattr(torch, args.dtype)
-
-    config = AutoConfig.from_pretrained(args.model_id, revision=args.revision)
-    if type(config) in MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING:
-        cls_model = AutoModelForSpeechSeq2Seq
-    elif type(config) in MODEL_FOR_MULTIMODAL_LM_MAPPING:
-        cls_model = AutoModelForMultimodalLM
-    elif type(config) in MODEL_FOR_CTC_MAPPING:
-        cls_model = AutoModelForCTC
-    else:
-        raise ValueError(f"Model config of type {type(config)} not recognized in Transformers mappings.")
-    is_ctc = cls_model == AutoModelForCTC
-
-    if "vibevoice" in args.model_id.lower():
-        model = cls_model.from_pretrained(
-            args.model_id, 
-            dtype=torch_dtype, 
-            attn_implementation={
-                "acoustic_tokenizer_encoder_config": "eager",
-                "semantic_tokenizer_encoder_config": "eager",
-                "text_config": "sdpa",
-            }
-        )
-    else:
-        model = cls_model.from_pretrained(
-            args.model_id, 
-            dtype=torch_dtype, 
-            revision=args.revision,
-            attn_implementation=args.attn_implementation,
-        )
-    model.to(args.device)
+    model = AutoModel.from_pretrained(
+        args.model_id,
+        torch_dtype=torch.bfloat16,
+        trust_remote_code=True,
+        attn_implementation="eager",
+        device_map=device,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(args.model_id, trust_remote_code=True)
     model.eval()
     print(f"Model size: {sum(p.numel() for p in model.parameters()) / 1e9:.2f}B parameters")
-    processor = AutoProcessor.from_pretrained(args.model_id, revision=args.revision)
-    has_transcription_processor = hasattr(processor, "apply_transcription_request")
-    is_cohere = "cohere" in args.model_id.lower() and "transcribe" in args.model_id.lower()
-
-    # Optional prompt for audio language models, newer models should use `apply_transcription_request`
-    text = None
-    if "granite-speech-3.3" in args.model_id.lower():
-        # create text prompt
-        chat = [
-            {
-                "role": "system",
-                "content": "Knowledge Cutoff Date: April 2024.\nToday's Date: December 19, 2024.\nYou are Granite, developed by IBM. You are a helpful AI assistant",
-            },
-            {
-                "role": "user",
-                "content": "<|audio|>can you transcribe the speech into a written format?",
-            }
-        ]
 
-        text = processor.apply_chat_template(
-            chat, tokenize=False, add_generation_prompt=True
-        )
+    # Required for generation stop conditions
+    model.audio_out_bos_token_id = tokenizer.convert_tokens_to_ids("<|audio_out_bos|>")
+    model.audio_eos_token_id = tokenizer.convert_tokens_to_ids("<|audio_eos|>")
 
-    # Extract sampling rate
-    if hasattr(processor, "feature_extractor") and processor.feature_extractor is not None:
-        sampling_rate = processor.feature_extractor.sampling_rate
-    elif hasattr(processor, "audio_processor") and processor.audio_processor is not None:
-        sampling_rate = processor.audio_processor.sampling_rate
-    else:
-        sampling_rate = 16_000
-
-    # Set generate arguments
-    if model.can_generate():
-        gen_kwargs = {"max_new_tokens": args.max_new_tokens}
-        if getattr(model.generation_config, "is_multilingual", False):
-            gen_kwargs["language"] = "en"
-            gen_kwargs["task"] = "transcribe"
-        # Clear deprecated Whisper generation config fields to suppress warnings
-        if hasattr(model.generation_config, "forced_decoder_ids"):
-            model.generation_config.forced_decoder_ids = None
-        if hasattr(model.generation_config, "suppress_tokens"):
-            model.generation_config.suppress_tokens = []
-        if hasattr(model.generation_config, "begin_suppress_tokens"):
-            model.generation_config.begin_suppress_tokens = []
-        if "granite-speech-3.3" in args.model_id.lower():
-            gen_kwargs["repetition_penalty"] = 1.0
-    elif args.max_new_tokens:
-        raise ValueError("`max_new_tokens` should only be set for auto-regressive models, but got a CTC model.")
-
-    if args.torch_compile is not None:
-        if model.can_generate():
-            gen_kwargs["compile_config"] = CompileConfig(mode=args.torch_compile, fullgraph=args.compile_fullgraph)
-            # enable static k/v cache for autoregressive models
-            model.generation_config.cache_implementation = "static"
-        else:
-            model = torch.compile(model, mode=args.torch_compile, fullgraph=args.compile_fullgraph)     
-
-        # Ensure warm-up runs when using torch.compile
-        if args.warmup_steps is None or args.warmup_steps < 1:
-            print("`--torch_compile` is enabled; forcing `--warmup_steps=10` to trigger compilation before timed runs.")
-            args.warmup_steps = 10
-    
-    def benchmark(batch, min_new_tokens=None):
+    transcribe_batch = load_transcribe_fn(args.model_id)
+
+    def benchmark(batch):
         # Load audio inputs
         audios = [audio["array"] for audio in batch["audio"]]
+        batch["audio_length_s"] = [
+            len(audio) / batch["audio"][0]["sampling_rate"] for audio in audios
+        ]
         minibatch_size = len(audios)
-        sampling_rate = batch["audio"][0]["sampling_rate"]
-        batch["audio_length_s"] = [len(audio) / sampling_rate for audio in audios]
-        batch["audio_filepath"] = data_utils.extract_audio_filepaths_from_batch(batch, minibatch_size)
-        if text is not None:
-            texts=[text] * minibatch_size
-        else:
-            texts = None
 
         # START TIMING
-        torch.cuda.synchronize(device=args.device)
-        start_event = torch.cuda.Event(enable_timing=True)
-        end_event = torch.cuda.Event(enable_timing=True)
-        start_event.record()
-
-        # 1. Pre-Processing
-        # 1.1 Pad audios to max batch size if using torch compile to prevent re-compilations
-        padding_size = None
-        if minibatch_size != args.batch_size and args.torch_compile is not None:
-            padding_size = args.batch_size - minibatch_size
-            padding_audios = [audios[-1] for _ in range(padding_size)]
-            audios.extend(padding_audios)
-
-        if is_cohere:
-            inputs = processor(
-                audios,
-                sampling_rate=sampling_rate,
-                return_tensors="pt",
-                language="en",
-                punctuation=False,
-            )
-        elif has_transcription_processor:
-            if "voxtral" in args.model_id.lower():
-                inputs = processor.apply_transcription_request(
-                    language="en",  # English for benchmark consistency
-                    audio=audios,
-                    model_id=args.model_id,
-                    sampling_rate=sampling_rate,
-                    format=["wav"] * len(audios),
-                )
-            else:
-                inputs = processor.apply_transcription_request(audios)
-            prompt_len = inputs["input_ids"].shape[1]
-        elif texts is not None:
-            inputs = processor(
-                texts,
-                audios,
-                device=args.device, # Computation device; returned tensors are put on CPU
-                return_tensors="pt",
-            ).to(args.device)
-            prompt_len = inputs["input_ids"].shape[1]
-        elif not model.can_generate(): #or len(audios[0]) > processor.feature_extractor.n_samples:
-            # 1.2 Either CTC pre-processing (normalize to mean 0, std 1), or long-form Whisper processing
-            inputs = processor(
-                audios,
-                sampling_rate=sampling_rate,
-                truncation=False,
-                padding="longest",
-                return_tensors="pt",
-                return_attention_mask=True,
-            )
-        else:
-            # 1.3 Standard Whisper processing: pad audios to 30-seconds and converted to log-mel
-            if args.longform:
-                inputs = processor(
-                    audios,
-                    sampling_rate=sampling_rate,
-                    return_tensors="pt",
-                    truncation=False,
-                    padding="longest",
-                    return_attention_mask=True,
-                )
-            else:
-                inputs = processor(audios, sampling_rate=sampling_rate, return_tensors="pt", padding="longest", return_attention_mask=True, device=args.device)
-
-        inputs = inputs.to(args.device, dtype=torch_dtype)
-
-        # 2. Model Inference
-        if args.torch_compile is not None:
-            sdpa_backends = [SDPBackend.MATH]
-        else:
-            sdpa_backends = [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
-        with sdpa_kernel(sdpa_backends):
-            if model.can_generate():
-                # 2.1 Auto-regressive generation for LM-based models
-                if args.longform:
-                    pred_ids = model.generate(**inputs, **gen_kwargs, return_timestamps=True)
-                else:   
-                    pred_ids = model.generate(**inputs, **gen_kwargs, min_new_tokens=min_new_tokens)
-            else:
-                # 2.2. Single forward pass for CTC
-                with torch.no_grad():
-                    logits = model(**inputs).logits
-                    pred_ids = logits.argmax(-1)
-
-        # 3. Post-processing
-        # 3.1 Strip padded ids from predictions
-        if padding_size is not None:
-            pred_ids = pred_ids[:-padding_size, ...]
-
-        # 3.2 Convert token ids to text transcription
-        if is_cohere:
-            audio_chunk_index = inputs.get("audio_chunk_index")
-            pred_text = processor.decode(
-                pred_ids, skip_special_tokens=True,
-                audio_chunk_index=audio_chunk_index, language="en",
-            )
-            pred_text = [remove_brackets(t) for t in pred_text]
-        elif "vibevoice" in args.model_id.lower():
-            # VibeVoice: strip the input prompt tokens then use the model's own decode API
-            generated_ids = pred_ids[:, prompt_len:]
-            try:
-                pred_text = processor.decode(generated_ids, return_format="transcription_only")
-            except Exception as e:
-                print(f"Batch decoding failed with error: {e}. Falling back to individual sample decoding.")
-                pred_text = []
-                for i, sample_ids in enumerate(generated_ids):
-                    try:
-                        decoded = processor.decode(sample_ids.unsqueeze(0), return_format="transcription_only")
-                        pred_text.append(decoded[0] if isinstance(decoded, list) else decoded)
-                    except Exception as sample_error:
-                        print(f"Sample {i} decoding failed with error: {sample_error}. Setting to empty transcript.")
-                        pred_text.append("")
-        elif has_transcription_processor or texts is not None:
-            # Strip input prompt tokens
-            pred_text = processor.decode(pred_ids[:, prompt_len:], skip_special_tokens=True)
-        elif is_ctc:
-            # don't use skip_special_tokens as it collapses double letters
-            pred_text = processor.batch_decode(pred_ids)
-        else:
-            pred_text = processor.decode(pred_ids, skip_special_tokens=True)
+        start_time = time.time()
+
+        # INFERENCE
+        pred_text = transcribe_batch(
+            model, tokenizer, audios, sample_rates=16000,
+            max_new_tokens=args.max_new_tokens,
+        )
 
         # END TIMING
-        end_event.record()
-        torch.cuda.synchronize(device=args.device)
-        runtime = start_event.elapsed_time(end_event) / 1000.0
+        runtime = time.time() - start_time
 
         # normalize by minibatch size since we want the per-sample time
         batch["transcription_time_s"] = minibatch_size * [runtime / minibatch_size]
@@ -281,30 +105,38 @@ def main(args):
         return batch
 
     if args.warmup_steps is not None:
-        dataset = data_utils.load_data(args)
-        dataset = data_utils.prepare_data(dataset, sampling_rate=sampling_rate)
+        warmup_dataset = data_utils.load_data(args)
+        warmup_dataset = data_utils.prepare_data(warmup_dataset)
 
         num_warmup_samples = args.warmup_steps * args.batch_size
         if args.streaming:
-            warmup_dataset = dataset.take(num_warmup_samples)
+            warmup_dataset = warmup_dataset.take(num_warmup_samples)
         else:
-            warmup_dataset = dataset.select(range(min(num_warmup_samples, len(dataset))))
-        warmup_dataset = iter(warmup_dataset.map(benchmark, batch_size=args.batch_size, batched=True, fn_kwargs={"min_new_tokens": args.max_new_tokens}))
+            warmup_dataset = warmup_dataset.select(
+                range(min(num_warmup_samples, len(warmup_dataset)))
+            )
+        warmup_dataset = iter(
+            warmup_dataset.map(benchmark, batch_size=args.batch_size, batched=True)
+        )
 
         for _ in tqdm(warmup_dataset, desc="Warming up..."):
             continue
 
     dataset = data_utils.load_data(args)
+    dataset = data_utils.prepare_data(dataset)
+
     if args.max_eval_samples is not None and args.max_eval_samples > 0:
         print(f"Subsampling dataset to first {args.max_eval_samples} samples!")
         if args.streaming:
             dataset = dataset.take(args.max_eval_samples)
         else:
-            dataset = dataset.select(range(min(args.max_eval_samples, len(dataset))))
-    dataset = data_utils.prepare_data(dataset, sampling_rate=sampling_rate)
+            dataset = dataset.select(
+                range(min(args.max_eval_samples, len(dataset)))
+            )
 
     dataset = dataset.map(
-        benchmark, batch_size=args.batch_size, batched=True, remove_columns=["audio"],
+        benchmark, batch_size=args.batch_size, batched=True,
+        remove_columns=["audio"],
     )
 
     all_results = {
@@ -312,7 +144,6 @@ def main(args):
         "transcription_time_s": [],
         "predictions": [],
         "references": [],
-        "audio_filepath": [],
     }
     result_iter = iter(dataset)
     for result in tqdm(result_iter, desc="Samples..."):
@@ -320,7 +151,6 @@ def main(args):
             all_results[key].append(result[key])
 
     # Write manifest results (WER and RTFX)
-    # Filtering of empty references is handled inside write_manifest.
     manifest_path = data_utils.write_manifest(
         all_results["references"],
         all_results["predictions"],
@@ -330,7 +160,6 @@ def main(args):
         args.split,
         audio_length=all_results["audio_length_s"],
         transcription_time=all_results["transcription_time_s"],
-        audio_filepaths=all_results["audio_filepath"],
     )
     print("Results saved at path:", os.path.abspath(manifest_path))
 
@@ -338,7 +167,9 @@ def main(args):
         references=all_results["references"], predictions=all_results["predictions"]
     )
     wer = round(100 * wer, 2)
-    rtfx = round(sum(all_results["audio_length_s"]) / sum(all_results["transcription_time_s"]), 2)
+    rtfx = round(
+        sum(all_results["audio_length_s"]) / sum(all_results["transcription_time_s"]), 2
+    )
     print("WER:", wer, "%", "RTFx:", rtfx)
 
 
@@ -349,26 +180,25 @@ if __name__ == "__main__":
         "--model_id",
         type=str,
         required=True,
-        help="Model identifier. Should be loadable with 🤗 Transformers",
+        help="Model identifier. Should be a HiggsAudio3 checkpoint on the HF Hub.",
     )
     parser.add_argument(
         "--dataset_path",
         type=str,
-        default="hf-audio/open-asr-leaderboard",
-        help="Dataset path. By default, it is `hf-audio/open-asr-leaderboard`",
+        default="esb/datasets",
+        help="Dataset path. By default, it is `esb/datasets`.",
     )
     parser.add_argument(
         "--dataset",
         type=str,
         required=True,
-        help="Dataset name. *E.g.* `'librispeech_asr` for the LibriSpeech ASR dataset, or `'common_voice'` for Common Voice. The full list of dataset names "
-        "can be found at `https://huggingface.co/datasets/hf-audio/open-asr-leaderboard`",
+        help="Dataset name.",
     )
     parser.add_argument(
         "--split",
         type=str,
         default="test",
-        help="Split of the dataset. *E.g.* `'validation`' for the dev split, or `'test'` for the test split.",
+        help="Split of the dataset.",
     )
     parser.add_argument(
         "--device",
@@ -379,7 +209,7 @@ if __name__ == "__main__":
     parser.add_argument(
         "--batch_size",
         type=int,
-        default=16,
+        default=4,
         help="Number of samples to go through each streamed batch.",
     )
     parser.add_argument(
@@ -389,43 +219,16 @@ if __name__ == "__main__":
         help="Number of samples to be evaluated. Put a lower number e.g. 64 for testing this script.",
     )
     parser.add_argument(
-        "--streaming",
-        action="store_true",
-        help="Stream the dataset lazily over the network instead of downloading it in full before the evaluation. Off by default for reproducible benchmark timings.",
+        "--no-streaming",
+        dest="streaming",
+        action="store_false",
+        help="Choose whether you'd like to download the entire dataset or stream it during the evaluation.",
     )
     parser.add_argument(
         "--max_new_tokens",
         type=int,
-        default=None,
-        help="Maximum number of tokens to generate (for auto-regressive models).",
-    )
-    parser.add_argument(
-        "--longform",
-        action="store_true",
-        help="Whether to use longform mode.",
-    )
-    parser.add_argument(
-        "--torch_compile",
-        type=str,
-        default=None,
-        help="Mode for torch compiling model forward pass. Can be either 'default', 'reduce-overhead', 'max-autotune' or 'max-autotune-no-cudagraphs'.",
-    )
-    parser.add_argument(
-        "--compile_fullgraph",
-        action="store_true",
-        help="Whether to do full graph compilation.",
-    )
-    parser.add_argument(
-        "--dtype",
-        type=str,
-        default="bfloat16",
-        help="The dtype to use for model loading and inference. E.g. 'bfloat16', 'float16', 'float32'.",
-    )
-    parser.add_argument(
-        "--attn_implementation",
-        type=str,
-        default="sdpa",
-        help="Attention implementation to use for model loading (e.g. 'sdpa', 'eager', 'flash_attention_2').",
+        default=1024,
+        help="Maximum number of tokens to generate (includes the chain-of-thought block).",
     )
     parser.add_argument(
         "--warmup_steps",
@@ -433,16 +236,7 @@ if __name__ == "__main__":
         default=10,
         help="Number of warm-up steps to run before launching the timed runs.",
     )
-    parser.add_argument(
-        "--revision",
-        type=str,
-        default=None,
-        help="Model revision to use (e.g. 'refs/pr/11' for a PR branch). Defaults to the main branch.",
-    )
     args = parser.parse_args()
+    parser.set_defaults(streaming=False)
 
-    print("*" * 100)
-    print(f"Evaluating {args.model_id} on {args.dataset_path} / {args.dataset} / {args.split}")
-    print("*" * 100)
-
-    main(args)
+    main(args)