279 lines
10 KiB
Python
279 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
"""Benchmark FunASR on representative audio during ASR migration.
|
|
|
|
The script intentionally measures only the FunASR side of a migration test. Run
|
|
Whisper or a cloud ASR baseline separately, then compare transcripts with your
|
|
normal WER/CER or human-review process.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
import time
|
|
import wave
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Iterable, List, Optional
|
|
|
|
try:
|
|
from funasr import AutoModel
|
|
from funasr.utils.postprocess_utils import rich_transcription_postprocess
|
|
except Exception as exc: # pragma: no cover - import message is for users
|
|
print(
|
|
"Failed to import FunASR. Install it with `pip install -U funasr` "
|
|
"or run this script from the repository root.",
|
|
file=sys.stderr,
|
|
)
|
|
raise
|
|
|
|
AUDIO_EXTENSIONS = (
|
|
".wav",
|
|
".mp3",
|
|
".flac",
|
|
".m4a",
|
|
".aac",
|
|
".ogg",
|
|
".opus",
|
|
".wma",
|
|
)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Run FunASR over a representative audio set and write migration benchmark results."
|
|
)
|
|
parser.add_argument("--input", "-i", type=Path, required=True, help="Audio file or folder to benchmark.")
|
|
parser.add_argument(
|
|
"--output-dir",
|
|
"-o",
|
|
type=Path,
|
|
default=Path("outputs/migration_benchmark"),
|
|
help="Directory for results.jsonl and summary.md.",
|
|
)
|
|
parser.add_argument(
|
|
"--model",
|
|
"-m",
|
|
default="iic/SenseVoiceSmall",
|
|
help="FunASR model name or ModelScope/Hugging Face id.",
|
|
)
|
|
parser.add_argument("--device", "-d", default="cpu", help="Inference device: cpu, cuda, or mps.")
|
|
parser.add_argument("--vad-model", default="fsmn-vad", help="VAD model, or 'none' to disable.")
|
|
parser.add_argument("--spk-model", default="", help="Optional speaker model such as cam++.")
|
|
parser.add_argument("--language", default="auto", help="Language hint passed to model.generate.")
|
|
parser.add_argument("--batch-size", type=int, default=1, help="batch_size passed to model.generate.")
|
|
parser.add_argument("--recursive", "-r", action="store_true", help="Recursively scan input folders.")
|
|
parser.add_argument(
|
|
"--extensions",
|
|
nargs="+",
|
|
default=list(AUDIO_EXTENSIONS),
|
|
help="Audio extensions to include when --input is a folder.",
|
|
)
|
|
parser.add_argument(
|
|
"--metadata",
|
|
action="append",
|
|
default=[],
|
|
help="Free-form key=value metadata to copy into the summary, e.g. --metadata baseline=whisper-large-v3.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def normalize_extensions(values: Iterable[str]) -> List[str]:
|
|
return sorted({value.lower() if value.startswith(".") else f".{value.lower()}" for value in values})
|
|
|
|
|
|
def iter_audio_files(path: Path, extensions: List[str], recursive: bool) -> List[Path]:
|
|
if path.is_file():
|
|
return [path]
|
|
if not path.exists():
|
|
raise FileNotFoundError(f"Input path does not exist: {path}")
|
|
walker = path.rglob("*") if recursive else path.iterdir()
|
|
return sorted(p for p in walker if p.is_file() and p.suffix.lower() in extensions)
|
|
|
|
|
|
def audio_duration_seconds(path: Path) -> Optional[float]:
|
|
try:
|
|
import soundfile as sf
|
|
|
|
info = sf.info(str(path))
|
|
if info.samplerate:
|
|
return float(info.frames) / float(info.samplerate)
|
|
except Exception:
|
|
pass
|
|
|
|
if path.suffix.lower() == ".wav":
|
|
try:
|
|
with wave.open(str(path), "rb") as wav_file:
|
|
frames = wav_file.getnframes()
|
|
rate = wav_file.getframerate()
|
|
if rate:
|
|
return float(frames) / float(rate)
|
|
except Exception:
|
|
return None
|
|
return None
|
|
|
|
|
|
def metadata_dict(items: Iterable[str]) -> Dict[str, str]:
|
|
parsed: Dict[str, str] = {}
|
|
for item in items:
|
|
if "=" not in item:
|
|
parsed[item] = ""
|
|
continue
|
|
key, value = item.split("=", 1)
|
|
parsed[key.strip()] = value.strip()
|
|
return parsed
|
|
|
|
|
|
def extract_text(result: Any) -> str:
|
|
if not result:
|
|
return ""
|
|
first = result[0] if isinstance(result, list) else result
|
|
if isinstance(first, dict):
|
|
text = first.get("text", "")
|
|
else:
|
|
text = str(first)
|
|
try:
|
|
return rich_transcription_postprocess(text)
|
|
except Exception:
|
|
return text
|
|
|
|
|
|
def write_jsonl(path: Path, rows: Iterable[Dict[str, Any]]) -> None:
|
|
with path.open("w", encoding="utf-8") as handle:
|
|
for row in rows:
|
|
handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")
|
|
|
|
|
|
def markdown_summary(rows: List[Dict[str, Any]], args: argparse.Namespace, model_load_seconds: float) -> str:
|
|
successful = [row for row in rows if not row.get("error")]
|
|
failed = [row for row in rows if row.get("error")]
|
|
known_audio = [row for row in successful if row.get("duration_seconds")]
|
|
total_audio = sum(float(row["duration_seconds"]) for row in known_audio)
|
|
total_elapsed = sum(float(row["elapsed_seconds"]) for row in successful)
|
|
throughput = total_audio / total_elapsed if total_audio and total_elapsed else None
|
|
meta = metadata_dict(args.metadata)
|
|
|
|
lines = [
|
|
"# FunASR Migration Benchmark Summary",
|
|
"",
|
|
"## Run configuration",
|
|
"",
|
|
f"- Input: `{args.input}`",
|
|
f"- Model: `{args.model}`",
|
|
f"- Device: `{args.device}`",
|
|
f"- VAD model: `{args.vad_model}`",
|
|
f"- Speaker model: `{args.spk_model or 'none'}`",
|
|
f"- Language: `{args.language}`",
|
|
f"- Batch size: `{args.batch_size}`",
|
|
f"- Model load seconds: `{model_load_seconds:.3f}`",
|
|
]
|
|
for key, value in meta.items():
|
|
lines.append(f"- {key}: `{value}`")
|
|
|
|
lines.extend(
|
|
[
|
|
"",
|
|
"## Aggregate results",
|
|
"",
|
|
f"- Files: `{len(rows)}`",
|
|
f"- Successful: `{len(successful)}`",
|
|
f"- Failed: `{len(failed)}`",
|
|
f"- Known audio seconds: `{total_audio:.3f}`" if known_audio else "- Known audio seconds: `unknown`",
|
|
f"- Inference seconds: `{total_elapsed:.3f}`" if successful else "- Inference seconds: `0.000`",
|
|
f"- Aggregate realtime factor: `{throughput:.3f}x`" if throughput else "- Aggregate realtime factor: `unknown`",
|
|
"",
|
|
"## Per-file results",
|
|
"",
|
|
"| File | Audio sec | Inference sec | RTF | Status | Text preview |",
|
|
"|---|---:|---:|---:|---|---|",
|
|
]
|
|
)
|
|
for row in rows:
|
|
duration = row.get("duration_seconds")
|
|
elapsed = row.get("elapsed_seconds")
|
|
rtf = row.get("realtime_factor")
|
|
status = "error" if row.get("error") else "ok"
|
|
preview = (row.get("text") or row.get("error") or "").replace("|", "\\|").replace("\n", " ")[:120]
|
|
lines.append(
|
|
"| {file} | {duration} | {elapsed} | {rtf} | {status} | {preview} |".format(
|
|
file=row["input"],
|
|
duration=f"{duration:.3f}" if isinstance(duration, (int, float)) else "",
|
|
elapsed=f"{elapsed:.3f}" if isinstance(elapsed, (int, float)) else "",
|
|
rtf=f"{rtf:.3f}x" if isinstance(rtf, (int, float)) else "",
|
|
status=status,
|
|
preview=preview,
|
|
)
|
|
)
|
|
lines.extend(
|
|
[
|
|
"",
|
|
"## Next comparison steps",
|
|
"",
|
|
"- Run your Whisper or cloud ASR baseline on the same files.",
|
|
"- Compare transcripts with human review or your normal WER/CER workflow.",
|
|
"- Keep model download and warmup time separate from steady-state throughput.",
|
|
"- Share reproducible findings in a FunASR showcase issue when possible.",
|
|
"",
|
|
]
|
|
)
|
|
return "\n".join(lines)
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
extensions = normalize_extensions(args.extensions)
|
|
files = iter_audio_files(args.input, extensions, args.recursive)
|
|
if not files:
|
|
print(f"No audio files found under {args.input} for extensions {extensions}", file=sys.stderr)
|
|
sys.exit(2)
|
|
|
|
vad_model = None if args.vad_model.lower() == "none" else args.vad_model
|
|
model_kwargs: Dict[str, Any] = {"model": args.model, "vad_model": vad_model, "device": args.device}
|
|
if args.spk_model:
|
|
model_kwargs["spk_model"] = args.spk_model
|
|
|
|
print(f"Loading FunASR model: {args.model} on {args.device}")
|
|
load_start = time.perf_counter()
|
|
model = AutoModel(**model_kwargs)
|
|
model_load_seconds = time.perf_counter() - load_start
|
|
print(f"Model loaded in {model_load_seconds:.3f}s")
|
|
|
|
rows: List[Dict[str, Any]] = []
|
|
for index, audio_path in enumerate(files, start=1):
|
|
display = str(audio_path if args.input.is_file() else audio_path.relative_to(args.input))
|
|
print(f"[{index}/{len(files)}] {display}")
|
|
duration = audio_duration_seconds(audio_path)
|
|
start = time.perf_counter()
|
|
row: Dict[str, Any] = {
|
|
"input": display,
|
|
"path": str(audio_path),
|
|
"duration_seconds": duration,
|
|
"model": args.model,
|
|
"device": args.device,
|
|
"language": args.language,
|
|
}
|
|
try:
|
|
result = model.generate(input=str(audio_path), language=args.language, batch_size=args.batch_size)
|
|
elapsed = time.perf_counter() - start
|
|
row["elapsed_seconds"] = elapsed
|
|
row["realtime_factor"] = (duration / elapsed) if duration and elapsed else None
|
|
row["text"] = extract_text(result)
|
|
print(f" ok: {elapsed:.3f}s" + (f", {row['realtime_factor']:.3f}x" if row["realtime_factor"] else ""))
|
|
except Exception as exc: # keep benchmarking other files
|
|
elapsed = time.perf_counter() - start
|
|
row["elapsed_seconds"] = elapsed
|
|
row["realtime_factor"] = None
|
|
row["error"] = repr(exc)
|
|
print(f" error: {exc}", file=sys.stderr)
|
|
rows.append(row)
|
|
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
results_path = args.output_dir / "results.jsonl"
|
|
summary_path = args.output_dir / "summary.md"
|
|
write_jsonl(results_path, rows)
|
|
summary_path.write_text(markdown_summary(rows, args, model_load_seconds), encoding="utf-8")
|
|
print(f"\nWrote {results_path}")
|
|
print(f"Wrote {summary_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|