#!/usr/bin/env python3 from __future__ import annotations import argparse import inspect import json import platform import socket from pathlib import Path from types import SimpleNamespace from typing import Any, Sequence import numpy as np import torch from run_eval import DEFAULT_MODEL_REVISION, Audio8Evaluator, iter_local_rows SPLITS = [ ("ami_cleaned", "test", "ami_cleaned/test", "ami_cleaned_test"), ("earnings22", "test", "earnings22/test", "earnings22_test"), ("gigaspeech_cleaned", "test", "gigaspeech_cleaned/test", "gigaspeech_cleaned_test"), ("librispeech", "test.clean", "librispeech/test.clean", "librispeech_test.clean"), ("librispeech", "test.other", "librispeech/test.other", "librispeech_test.other"), ("spgispeech", "test", "spgispeech/test", "spgispeech_test"), ("voxpopuli_cleaned_aa", "test", "voxpopuli_cleaned_aa/test", "voxpopuli_cleaned_aa_test"), ] def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Audit Audio8 singleton and safe batch behavior locally.") parser.add_argument("--model_id", required=True) parser.add_argument("--model_revision", default=DEFAULT_MODEL_REVISION) parser.add_argument("--data_root", type=Path, required=True) parser.add_argument("--historical_results_dir", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--device", default="0") parser.add_argument("--dtype", default="bfloat16") parser.add_argument("--samples_per_split", type=int, default=4) parser.add_argument("--batch_sizes", default="2,4,8,16,32") parser.add_argument("--max_new_tokens", type=int, default=256) parser.add_argument("--max_audio_seconds", type=float, default=30.0) parser.add_argument("--feature_workers", type=int, default=16) parser.add_argument("--torch_cpu_threads", type=int, default=1) return parser.parse_args() def evaluator_args(args: argparse.Namespace) -> SimpleNamespace: return SimpleNamespace( model_id=args.model_id, model_revision=args.model_revision, device=args.device, dtype=args.dtype, attn_implementation="eager", max_new_tokens=args.max_new_tokens, max_audio_seconds=args.max_audio_seconds, sampling_rate=16000, model_max_length=1000, local_files_only=Path(args.model_id).exists(), torch_compile=None, compile_fullgraph=False, feature_workers=args.feature_workers, torch_cpu_threads=args.torch_cpu_threads, ) def local_source_args( args: argparse.Namespace, dataset: str, split: str, relative_path: str, max_samples: int, ) -> SimpleNamespace: return SimpleNamespace( local_parquet_file=[], local_parquet_dir=args.data_root / relative_path, skip_samples=0, max_eval_samples=max_samples, sampling_rate=16000, dataset=dataset, split=split, ) def historical_predictions(results_dir: Path, suffix: str) -> dict[str, str]: matches = sorted(results_dir.glob(f"*_{suffix}.jsonl")) if len(matches) != 1: raise FileNotFoundError(f"Expected one historical result for {suffix}, found {matches}") predictions: dict[str, str] = {} with matches[0].open(encoding="utf-8") as handle: for line in handle: if not line.strip(): continue row = json.loads(line) sample_id = row.get("sample_id") if sample_id is not None: predictions[str(sample_id)] = str(row.get("pred_text") or "") return predictions def processor_features( evaluator: Audio8Evaluator, waveforms: Sequence[np.ndarray], ) -> tuple[torch.Tensor, torch.Tensor]: batch, _ = evaluator.prepare_batch_inputs(waveforms) return batch["input_features"].cpu(), batch["feature_lens"].cpu() def feature_prefix_audit( evaluator: Audio8Evaluator, waveforms: Sequence[np.ndarray], ) -> dict[str, Any]: batch_features, feature_lengths = processor_features(evaluator, waveforms) max_abs_diff = 0.0 exact_prefixes = 0 for index, waveform in enumerate(waveforms): singleton_features, singleton_lengths = processor_features(evaluator, [waveform]) feature_length = int(feature_lengths[index].item()) if feature_length != int(singleton_lengths[0].item()): raise AssertionError("Batch and singleton feature lengths differ") left = batch_features[index, :, :feature_length].float() right = singleton_features[0, :, :feature_length].float() difference = float((left - right).abs().max().item()) if left.numel() else 0.0 max_abs_diff = max(max_abs_diff, difference) exact_prefixes += int(torch.equal(left, right)) return { "batch_feature_width": int(batch_features.shape[-1]), "feature_lengths": [int(value) for value in feature_lengths.tolist()], "has_mixed_feature_lengths": len(set(int(value) for value in feature_lengths.tolist())) > 1, "exact_valid_prefixes": exact_prefixes, "samples": len(waveforms), "max_abs_valid_prefix_diff": max_abs_diff, } def compare_outputs(left: Sequence[list[int]], right: Sequence[list[int]]) -> list[bool]: return [list(a) == list(b) for a, b in zip(left, right)] def main() -> None: args = parse_args() batch_sizes = [int(value) for value in args.batch_sizes.split(",") if value.strip()] if not batch_sizes or min(batch_sizes) < 1: raise ValueError("--batch_sizes must contain positive integers") if args.feature_workers < 1 or args.torch_cpu_threads < 1: raise ValueError("--feature_workers and --torch_cpu_threads must be positive") torch.set_num_threads(int(args.torch_cpu_threads)) evaluator = Audio8Evaluator(evaluator_args(args)) model_signature = inspect.signature(evaluator.model.forward) report: dict[str, Any] = { "status": "running", "host": socket.gethostname(), "platform": platform.platform(), "model_id": args.model_id, "model_revision": args.model_revision, "model_parameters": evaluator.model_size, "dtype": str(evaluator.dtype), "device": str(evaluator.device), "feature_lens_supported_by_model": "feature_lens" in model_signature.parameters, "padding_side": evaluator.processor.tokenizer.padding_side, "samples_per_split": args.samples_per_split, "batch_sizes": batch_sizes, "feature_workers": args.feature_workers, "torch_cpu_threads": args.torch_cpu_threads, "splits": {}, } args.output.parent.mkdir(parents=True, exist_ok=True) singleton_history_mismatches = 0 repeat_mismatches = 0 feature_prefix_failures = 0 for dataset, split, relative_path, suffix in SPLITS: max_samples = max(max(batch_sizes), int(args.samples_per_split)) rows = list(iter_local_rows(local_source_args(args, dataset, split, relative_path, max_samples))) if len(rows) < args.samples_per_split: raise ValueError(f"Not enough rows for {dataset}/{split}: {len(rows)}") audit_rows = rows[: args.samples_per_split] history = historical_predictions(args.historical_results_dir, suffix) singleton_outputs = [evaluator.transcribe_batch([row.waveform]) for row in audit_rows] singleton_ids = [output.generated_ids[0] for output in singleton_outputs] singleton_predictions = [output.predictions[0] for output in singleton_outputs] history_matches = [ history.get(row.sample_id) == prediction for row, prediction in zip(audit_rows, singleton_predictions) ] singleton_history_mismatches += sum(not value for value in history_matches) waveforms = [row.waveform for row in audit_rows] batched = evaluator.transcribe_batch(waveforms) repeated = evaluator.transcribe_batch(waveforms) repeat_matches = compare_outputs(batched.generated_ids, repeated.generated_ids) singleton_matches = compare_outputs(singleton_ids, batched.generated_ids) repeat_mismatches += sum(not value for value in repeat_matches) feature_audit = feature_prefix_audit(evaluator, waveforms) feature_prefix_failures += int(feature_audit["max_abs_valid_prefix_diff"] != 0.0) search_results = [] max_stable_batch = 0 for batch_size in batch_sizes: if len(rows) < batch_size: continue selected = rows[:batch_size] selected_waveforms = [row.waveform for row in selected] if evaluator.device.type == "cuda": torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats(evaluator.device) try: output = evaluator.transcribe_batch(selected_waveforms) peak_memory = ( int(torch.cuda.max_memory_allocated(evaluator.device)) if evaluator.device.type == "cuda" else None ) audio_seconds = sum(row.duration for row in selected) search_results.append( { "batch_size": batch_size, "status": "passed", "runtime_seconds": output.runtime_seconds, "audio_seconds": audio_seconds, "rtfx": audio_seconds / output.runtime_seconds, "peak_memory_bytes": peak_memory, "stop_hits": sum(output.stop_hits), "max_new_hits": sum(output.max_new_hits), } ) max_stable_batch = batch_size except torch.cuda.OutOfMemoryError as error: search_results.append( {"batch_size": batch_size, "status": "oom", "error": str(error)} ) torch.cuda.empty_cache() break except RuntimeError as error: search_results.append( {"batch_size": batch_size, "status": "error", "error": str(error)} ) break report["splits"][f"{dataset}:{split}"] = { "sample_ids": [row.sample_id for row in audit_rows], "durations": [row.duration for row in audit_rows], "singleton_matches_historical": history_matches, "singleton_history_match_rate": sum(history_matches) / len(history_matches), "batch_matches_singleton_tokens": singleton_matches, "batch_singleton_token_match_rate": sum(singleton_matches) / len(singleton_matches), "batch_repeat_matches": repeat_matches, "batch_repeat_match_rate": sum(repeat_matches) / len(repeat_matches), "feature_prefix_audit": feature_audit, "batch_search": search_results, "max_stable_local_batch": max_stable_batch, } args.output.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") print( f"{dataset}/{split}: history={sum(history_matches)}/{len(history_matches)} " f"batch_vs_singleton={sum(singleton_matches)}/{len(singleton_matches)} " f"repeat={sum(repeat_matches)}/{len(repeat_matches)} " f"max_batch={max_stable_batch}", flush=True, ) report["singleton_history_mismatches"] = singleton_history_mismatches report["batch_repeat_mismatches"] = repeat_mismatches report["feature_prefix_failures"] = feature_prefix_failures report["status"] = ( "passed" if singleton_history_mismatches == 0 and repeat_mismatches == 0 and feature_prefix_failures == 0 else "failed" ) args.output.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") print(json.dumps({key: report[key] for key in ("status", "singleton_history_mismatches", "batch_repeat_mismatches", "feature_prefix_failures")}, indent=2)) if report["status"] != "passed": raise SystemExit(1) if __name__ == "__main__": main()