dudu/eval/asr_eval/report.py

"""汇总每条样本结果 → 语料级指标 + 延迟分位 + 资源，产出 JSON / Markdown / HTML。

错误率在语料级聚合（累加 S/D/I/N 再求率）。延迟取 p50/p95。资源取均值。
"""

from __future__ import annotations

import json
import os
from datetime import datetime

import numpy as np

from .metrics.err import Counts
from .metrics.resource import hardware_info


def _pct(values: list[float], p: float) -> float | None:
    vals = [v for v in values if v is not None]
    if not vals:
        return None
    return float(np.percentile(vals, p))


def _mean(values: list[float]) -> float | None:
    vals = [v for v in values if v is not None]
    if not vals:
        return None
    return float(np.mean(vals))


def _group_stats(rows: list[dict], cost_per_min: float | None) -> dict:
    acc = {"cer": Counts(), "wer": Counts(), "mer": Counts()}
    for r in rows:
        for k in acc:
            acc[k] = acc[k] + r["counts"][k]
    audio_total = sum(r["audio_sec"] for r in rows)
    return {
        "samples": len(rows),
        "errors": sum(1 for r in rows if r["error"]),
        "cer": acc["cer"].rate,
        "wer": acc["wer"].rate,
        "mer": acc["mer"].rate,
        "rtf_mean": _mean([r["rtf"] for r in rows]),
        "first_partial_p50": _pct([r["first_partial_sec"] for r in rows], 50),
        "first_partial_p95": _pct([r["first_partial_sec"] for r in rows], 95),
        "finalize_p50": _pct([r["finalize_sec"] for r in rows], 50),
        "finalize_p95": _pct([r["finalize_sec"] for r in rows], 95),
        "peak_rss_mb_mean": _mean([r["peak_rss_mb"] for r in rows]),
        "avg_cpu_mean": _mean([r["avg_cpu"] for r in rows]),
        "audio_sec_total": audio_total,
        "cost_total": (audio_total / 60.0 * cost_per_min) if cost_per_min else None,
    }


def aggregate(rows: list[dict], engines_meta: dict) -> dict:
    report = {
        "generated_at": datetime.now().isoformat(timespec="seconds"),
        "hardware": hardware_info(),
        "engines": {},
    }
    by_engine: dict[str, list[dict]] = {}
    for r in rows:
        by_engine.setdefault(r["engine"], []).append(r)

    for eng, erows in by_engine.items():
        meta = engines_meta.get(eng, {})
        cost = meta.get("cost_per_min")
        cats: dict[str, list[dict]] = {}
        dsets: dict[str, list[dict]] = {}
        for r in erows:
            cats.setdefault(r["category"], []).append(r)
            dsets.setdefault(r["dataset"], []).append(r)
        report["engines"][eng] = {
            "is_local": meta.get("is_local"),
            "kind": meta.get("kind"),
            "model_size_mb": meta.get("model_size_mb"),
            "cost_per_min": cost,
            "overall": _group_stats(erows, cost),
            "by_category": {c: _group_stats(rs, cost) for c, rs in sorted(cats.items())},
            "by_dataset": {d: _group_stats(rs, cost) for d, rs in sorted(dsets.items())},
        }
    return report


# ---------- 渲染 ----------

def _pct_str(v: float | None) -> str:
    return f"{v * 100:.2f}%" if v is not None else "—"


def _num(v: float | None, fmt: str = "{:.2f}") -> str:
    return fmt.format(v) if v is not None else "—"


def render_markdown(report: dict) -> str:
    hw = report["hardware"]
    lines = [
        "# ASR 模型评估报告",
        "",
        f"- 生成时间：{report['generated_at']}",
        f"- 硬件：{hw['platform']} · {hw['machine']} · {hw['cpu_count']} vCPU · {hw['total_mem_gb']} GB",
        "",
        "## 总览（各引擎全量）",
        "",
        "| 引擎 | 本地 | 样本 | 错误 | MER | CER | WER | RTF | 首包 p50 | 定稿 p50 | 峰值内存MB | 模型大小MB | 成本 |",
        "|---|---|--:|--:|--:|--:|--:|--:|--:|--:|--:|--:|--:|",
    ]
    for eng, e in report["engines"].items():
        o = e["overall"]
        lines.append(
            "| {eng} | {loc} | {n} | {err} | {mer} | {cer} | {wer} | {rtf} | {fp} | {fin} | {mem} | {sz} | {cost} |".format(
                eng=eng,
                loc="是" if e["is_local"] else "云",
                n=o["samples"],
                err=o["errors"],
                mer=_pct_str(o["mer"]),
                cer=_pct_str(o["cer"]),
                wer=_pct_str(o["wer"]),
                rtf=_num(o["rtf_mean"]),
                fp=_num(o["first_partial_p50"], "{:.2f}s") if o["first_partial_p50"] is not None else "—",
                fin=_num(o["finalize_p50"], "{:.2f}s") if o["finalize_p50"] is not None else "—",
                mem=_num(o["peak_rss_mb_mean"], "{:.0f}"),
                sz=_num(e["model_size_mb"], "{:.0f}"),
                cost=_num(o["cost_total"], "¥{:.4f}") if o["cost_total"] is not None else "—",
            )
        )

    # 中英混说专项
    lines += ["", "## 中英混说（code-switch）专项 MER", "", "| 引擎 | 样本 | MER | CER | WER |", "|---|--:|--:|--:|--:|"]
    for eng, e in report["engines"].items():
        cs = e["by_category"].get("code-switch")
        if not cs:
            continue
        lines.append(
            f"| {eng} | {cs['samples']} | {_pct_str(cs['mer'])} | {_pct_str(cs['cer'])} | {_pct_str(cs['wer'])} |"
        )

    # 按数据集
    lines += ["", "## 按数据集 MER", "", "| 引擎 | 数据集 | 样本 | MER | CER | WER |", "|---|---|--:|--:|--:|--:|"]
    for eng, e in report["engines"].items():
        for ds, g in e["by_dataset"].items():
            lines.append(
                f"| {eng} | {ds} | {g['samples']} | {_pct_str(g['mer'])} | {_pct_str(g['cer'])} | {_pct_str(g['wer'])} |"
            )
    lines.append("")
    lines.append(
        "> 注：流式(gummy)与离线(本地)延迟口径不同——gummy 的 RTF/延迟含 2x 推流与网络，"
        "非纯算力；本地引擎为整段离线处理。资源(内存/CPU)仅本地引擎测量，含 Python 进程基底。"
    )
    return "\n".join(lines)


_HTML_TMPL = """<!DOCTYPE html><html lang="zh-CN"><head><meta charset="utf-8">
<title>ASR 评估报告</title>
<style>
body{{font-family:-apple-system,"PingFang SC",sans-serif;max-width:1100px;margin:32px auto;padding:0 20px;color:#1B1E26;background:#F6F7FA}}
h1{{font-size:24px}} h2{{font-size:18px;margin-top:32px;border-left:3px solid #4F6EF7;padding-left:10px}}
table{{border-collapse:collapse;width:100%;background:#fff;border:1px solid #E4E6EB;border-radius:10px;overflow:hidden;margin:12px 0;font-size:13px}}
th,td{{padding:8px 10px;border-bottom:1px solid #F0F1F4;text-align:right}}
th{{background:#FCFCFD;color:#5A6072;font-weight:600;text-transform:uppercase;font-size:11px;letter-spacing:.04em}}
td:first-child,th:first-child{{text-align:left;font-weight:600}}
.meta{{color:#5A6072;font-size:13px}} .note{{color:#7A8090;font-size:12px;margin-top:8px}}
tr:hover td{{background:#F6F7FA}}
</style></head><body>
<h1>ASR 模型评估报告</h1>
<p class="meta">生成时间 {generated_at} · {hw}</p>
{body}
<p class="note">流式(gummy)与离线(本地)延迟口径不同；资源仅本地引擎测量，含 Python 进程基底。</p>
</body></html>"""


def render_html(report: dict) -> str:
    hw = report["hardware"]

    def table(headers, rows):
        h = "".join(f"<th>{x}</th>" for x in headers)
        body = ""
        for row in rows:
            body += "<tr>" + "".join(f"<td>{x}</td>" for x in row) + "</tr>"
        return f"<table><thead><tr>{h}</tr></thead><tbody>{body}</tbody></table>"

    parts = ["<h2>总览（各引擎全量）</h2>"]
    rows = []
    for eng, e in report["engines"].items():
        o = e["overall"]
        rows.append([
            eng, "是" if e["is_local"] else "云", o["samples"], o["errors"],
            _pct_str(o["mer"]), _pct_str(o["cer"]), _pct_str(o["wer"]),
            _num(o["rtf_mean"]),
            _num(o["first_partial_p50"], "{:.2f}s") if o["first_partial_p50"] is not None else "—",
            _num(o["peak_rss_mb_mean"], "{:.0f}"), _num(e["model_size_mb"], "{:.0f}"),
            _num(o["cost_total"], "¥{:.4f}") if o["cost_total"] is not None else "—",
        ])
    parts.append(table(
        ["引擎", "本地", "样本", "错误", "MER", "CER", "WER", "RTF", "首包p50", "峰值内存MB", "模型MB", "成本"], rows
    ))

    parts.append("<h2>中英混说专项 MER</h2>")
    cs_rows = []
    for eng, e in report["engines"].items():
        cs = e["by_category"].get("code-switch")
        if cs:
            cs_rows.append([eng, cs["samples"], _pct_str(cs["mer"]), _pct_str(cs["cer"]), _pct_str(cs["wer"])])
    parts.append(table(["引擎", "样本", "MER", "CER", "WER"], cs_rows))

    parts.append("<h2>按数据集</h2>")
    ds_rows = []
    for eng, e in report["engines"].items():
        for ds, g in e["by_dataset"].items():
            ds_rows.append([eng, ds, g["samples"], _pct_str(g["mer"]), _pct_str(g["cer"]), _pct_str(g["wer"])])
    parts.append(table(["引擎", "数据集", "样本", "MER", "CER", "WER"], ds_rows))

    return _HTML_TMPL.format(
        generated_at=report["generated_at"],
        hw=f"{hw['platform']} · {hw['cpu_count']} vCPU · {hw['total_mem_gb']}GB",
        body="\n".join(parts),
    )


def write_reports(report: dict, out_dir: str) -> dict:
    os.makedirs(out_dir, exist_ok=True)
    paths = {
        "json": os.path.join(out_dir, "report.json"),
        "md": os.path.join(out_dir, "report.md"),
        "html": os.path.join(out_dir, "report.html"),
    }
    with open(paths["json"], "w", encoding="utf-8") as f:
        json.dump(report, f, ensure_ascii=False, indent=2)
    with open(paths["md"], "w", encoding="utf-8") as f:
        f.write(render_markdown(report))
    with open(paths["html"], "w", encoding="utf-8") as f:
        f.write(render_html(report))
    return paths