transcribe_file.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
from typing import Any

from asr_onnx_runtime import OnnxAsrEngine, OnnxCacheAsrEngine


APP_DIR = Path(__file__).resolve().parent
BUNDLE_CANDIDATES = (
    "model_bundle",
)


def default_bundle_dir() -> Path:
    env_value = os.environ.get("AUDIO8_ASR_BUNDLE")
    if env_value:
        return Path(env_value).expanduser()
    for name in BUNDLE_CANDIDATES:
        candidate = APP_DIR / name
        if (candidate / "metadata.json").exists():
            return candidate
    return APP_DIR / "model_bundle"


def transcribe_file(
    audio_path: str | Path,
    *,
    bundle_dir: str | Path | None = None,
    backend: str = "onnx_cache",
    cache_precision: str = "int8",
    audio_precision: str = "int8",
    language: str | None = "Chinese",
    max_new_tokens: int = 128,
    provider: str = "CPUExecutionProvider",
    threads: int | None = None,
    hotwords: str | list[str] | None = None,
    hotword_topk: int = 50,
    hotword_start_boost: float = 6.0,
    hotword_continuation_boost: float = 8.0,
) -> dict[str, Any]:
    bundle_path = Path(bundle_dir).expanduser() if bundle_dir is not None else default_bundle_dir()
    audio_file = Path(audio_path).expanduser()
    if backend == "onnx":
        engine = OnnxAsrEngine(
            bundle_path,
            provider=provider,
            intra_op_num_threads=threads,
            audio_precision=audio_precision,
        )
    elif backend == "onnx_cache":
        engine = OnnxCacheAsrEngine(
            bundle_path,
            provider=provider,
            intra_op_num_threads=threads,
            cache_precision=cache_precision,
            audio_precision=audio_precision,
        )
    else:
        raise ValueError(f"Unsupported backend: {backend}")
    return engine.transcribe(
        audio_file.read_bytes(),
        language=language,
        max_new_tokens=max_new_tokens,
        hotwords=hotwords,
        hotword_topk=hotword_topk,
        hotword_start_boost=hotword_start_boost,
        hotword_continuation_boost=hotword_continuation_boost,
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Transcribe one audio file with the Audio8 ASR ONNX runtime.")
    parser.add_argument("audio", type=Path, help="Audio file path. WAV is recommended; librosa/soundfile handle common formats.")
    parser.add_argument("--bundle_dir", type=Path, default=default_bundle_dir())
    parser.add_argument("--backend", choices=["onnx_cache", "onnx"], default=os.environ.get("ASR_BACKEND", "onnx_cache"))
    parser.add_argument("--cache_precision", choices=["fp32", "int8", "int4", "auto"], default=os.environ.get("ASR_CACHE_PRECISION", "int8"))
    parser.add_argument("--audio_precision", choices=["fp32", "int8", "auto"], default=os.environ.get("ASR_AUDIO_PRECISION", "int8"))
    parser.add_argument("--language", default="Chinese")
    parser.add_argument("--max_new_tokens", type=int, default=128)
    parser.add_argument("--provider", default=os.environ.get("ORT_PROVIDER", "CPUExecutionProvider"))
    parser.add_argument("--threads", type=int, default=int(os.environ.get("ORT_THREADS", "0")))
    parser.add_argument("--hotwords", default=None, help="Comma or Chinese-comma separated hotwords, disabled when omitted.")
    parser.add_argument("--hotword_topk", type=int, default=50)
    parser.add_argument("--hotword_start_boost", type=float, default=6.0)
    parser.add_argument("--hotword_continuation_boost", type=float, default=8.0)
    parser.add_argument("--json", action="store_true", help="Print the full result JSON instead of only transcript text.")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    result = transcribe_file(
        args.audio,
        bundle_dir=args.bundle_dir,
        backend=args.backend,
        cache_precision=args.cache_precision,
        audio_precision=args.audio_precision,
        language=args.language or None,
        max_new_tokens=args.max_new_tokens,
        provider=args.provider,
        threads=args.threads if args.threads > 0 else None,
        hotwords=args.hotwords or None,
        hotword_topk=args.hotword_topk,
        hotword_start_boost=args.hotword_start_boost,
        hotword_continuation_boost=args.hotword_continuation_boost,
    )
    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        print(result.get("text", ""))


if __name__ == "__main__":
    main()