dudu/eval/asr_eval/cli.py

"""命令行入口：python -m asr_eval run --config config.yaml

流程：载 config → 物料化样本 → 逐引擎逐样本识别(本地引擎包资源探针) → 算分 → 出报告。
gummy 的 API key：优先 env DASHSCOPE_API_KEY，否则读 stdin 首行
（遵全局规则：rbw get dashscope-api-key | python -m asr_eval run ...）。
"""

from __future__ import annotations

import argparse
import os
import sys
from datetime import datetime

import yaml
from rich.console import Console
from rich.progress import BarColumn, Progress, TextColumn, TimeElapsedColumn

from .datasets import build_dataset
from .engines import Transcript, build_engine
from .metrics.err import empty_counts, score_sample
from .metrics.resource import ResourceProbe
from .report import aggregate, render_markdown, write_reports

console = Console()


def _resolve_api_key() -> str:
    key = os.environ.get("DASHSCOPE_API_KEY", "").strip()
    if key:
        return key
    if not sys.stdin.isatty():
        line = sys.stdin.readline().strip()
        if line:
            return line
    return ""


def _load_samples(cfg: dict, only: set[str] | None, limit: int | None) -> list:
    audio_cache = cfg.get("audio_cache", "datasets/_audio_cache")
    samples = []
    for dcfg in cfg.get("datasets", []):
        if only and dcfg["name"] not in only:
            continue
        ds = build_dataset(dcfg, audio_cache)
        console.print(f"[cyan]载入数据集[/] {dcfg['name']} (type={dcfg['type']}) ...")
        try:
            got = list(ds.samples(limit=limit))
        except Exception as ex:
            console.print(f"[red]数据集 {dcfg['name']} 载入失败，跳过：{ex}[/]")
            continue
        console.print(f"  → {len(got)} 条")
        samples.extend(got)
    return samples


def _enabled_engines(cfg: dict, only: set[str] | None, api_key: str) -> list[dict]:
    out = []
    for ecfg in cfg.get("engines", []):
        if not ecfg.get("enabled", True):
            continue
        if only and ecfg["name"] not in only:
            continue
        if ecfg["type"] == "gummy":
            ecfg = {**ecfg, "api_key": api_key}
        out.append(ecfg)
    return out


def cmd_run(args: argparse.Namespace) -> int:
    with open(args.config, encoding="utf-8") as f:
        cfg = yaml.safe_load(f)

    limit = args.sample_limit if args.sample_limit is not None else cfg.get("sample_limit")
    only_engines = set(args.engines.split(",")) if args.engines else None
    only_datasets = set(args.datasets.split(",")) if args.datasets else None

    engine_cfgs = _enabled_engines(cfg, only_engines, "")
    need_key = any(e["type"] == "gummy" for e in engine_cfgs)
    api_key = _resolve_api_key() if need_key else ""
    if need_key and not api_key:
        console.print("[red]gummy 已启用但未取到 API key（设 DASHSCOPE_API_KEY 或 stdin 传入）[/]")
        return 2
    engine_cfgs = _enabled_engines(cfg, only_engines, api_key)

    samples = _load_samples(cfg, only_datasets, limit)
    if not samples:
        console.print("[red]没有可评估的样本，检查 datasets 配置或 manifest[/]")
        return 1

    rows: list[dict] = []
    engines_meta: dict[str, dict] = {}

    for ecfg in engine_cfgs:
        name = ecfg["name"]
        console.print(f"\n[bold magenta]引擎 {name}[/] (type={ecfg['type']}) 加载中 ...")
        engine = build_engine(ecfg)
        try:
            engine.load()
        except Exception as ex:
            console.print(f"[red]引擎 {name} 加载失败，跳过：{ex}[/]")
            continue

        engines_meta[name] = {
            "is_local": engine.is_local,
            "kind": engine.kind,
            "model_size_mb": engine.model_size_mb(),
            "cost_per_min": engine.cost_per_min(),
        }

        with Progress(
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TextColumn("{task.completed}/{task.total}"),
            TimeElapsedColumn(),
            console=console,
        ) as prog:
            task = prog.add_task(f"{name}", total=len(samples))
            for s in samples:
                try:
                    tr = _transcribe_one(engine, s.audio_path)
                except Exception as ex:  # 单条样本失败不拖垮整轮评估
                    tr = Transcript(text="", audio_sec=0.0, proc_sec=0.0, error=str(ex))
                # 失败样本不计入准确率（counts 置空），仅在 errors 列计数
                sc = empty_counts() if tr.error else score_sample(s.ref_text, tr.text)
                rows.append({
                    "engine": name,
                    "dataset": s.dataset,
                    "category": s.category,
                    "lang": s.lang,
                    "domain": s.domain,
                    "sample_id": s.id,
                    "ref": s.ref_text,
                    "hyp": tr.text,
                    "counts": sc,
                    "audio_sec": tr.audio_sec,
                    "proc_sec": tr.proc_sec,
                    "rtf": tr.rtf,
                    "first_partial_sec": tr.first_partial_sec,
                    "finalize_sec": tr.finalize_sec,
                    "peak_rss_mb": tr.peak_rss_mb,
                    "avg_cpu": tr.avg_cpu,
                    "error": tr.error,
                })
                prog.advance(task)
        engine.unload()

    if not rows:
        console.print("[red]无结果（所有引擎都失败了？）[/]")
        return 1

    report = aggregate(rows, engines_meta)
    out_dir = os.path.join(cfg.get("output_dir", "results"), datetime.now().strftime("%Y%m%d-%H%M%S"))
    paths = write_reports(report, out_dir)
    _dump_raw(rows, out_dir)

    console.print("\n" + render_markdown(report))
    console.print(f"\n[green]报告已写入[/] {out_dir}/  (report.json / report.md / report.html)")
    return 0


def _transcribe_one(engine, audio_path: str):
    """本地引擎包资源探针；云端引擎直接调用。"""
    if engine.is_local:
        with ResourceProbe() as probe:
            tr = engine.transcribe(audio_path)
        if tr.peak_rss_mb is None:
            tr.peak_rss_mb = probe.peak_rss_mb
        if tr.avg_cpu is None:
            tr.avg_cpu = probe.avg_cpu
        return tr
    return engine.transcribe(audio_path)


def _dump_raw(rows: list[dict], out_dir: str) -> None:
    import json

    path = os.path.join(out_dir, "results.jsonl")
    with open(path, "w", encoding="utf-8") as f:
        for r in rows:
            rec = {k: v for k, v in r.items() if k != "counts"}
            rec["cer"] = r["counts"]["cer"].rate
            rec["wer"] = r["counts"]["wer"].rate
            rec["mer"] = r["counts"]["mer"].rate
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")


def main() -> None:
    ap = argparse.ArgumentParser(prog="asr_eval", description="dudu ASR 模型评估框架")
    sub = ap.add_subparsers(dest="cmd", required=True)

    run = sub.add_parser("run", help="跑评估")
    run.add_argument("--config", default="config.yaml")
    run.add_argument("--engines", help="只跑这些引擎（逗号分隔，覆盖 enabled）")
    run.add_argument("--datasets", help="只跑这些数据集（逗号分隔）")
    run.add_argument("--sample-limit", type=int, default=None, help="每个数据集采样上限（覆盖 config）")
    run.set_defaults(func=cmd_run)

    args = ap.parse_args()
    sys.exit(args.func(args))


if __name__ == "__main__":
    main()