File size: 4,556 Bytes
bc8e2a9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | #!/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()
|