dudu/eval/asr_eval/engines/gummy.py

"""云端 gummy 引擎：阿里云百炼 DashScope 流式 ASR（gummy-realtime-v1）。

协议对照 server/internal/asr/gummy.go：
  run-task → 等 task-started → 推二进制 PCM(3200B/16k/pcm) → finish-task → 收 task-finished。
下行解析 payload.output.transcription(gummy) 或 sentence(paraformer)，
sentence_end 兼容 bool 与字符串 "true"/"false"（gummy.go 的 flexBool 怪癖）。
"""

from __future__ import annotations

import json
import threading
import time
import uuid

import websocket  # websocket-client

from ..audio import load_pcm16
from .base import Engine, Transcript

WS_URL = "wss://dashscope.aliyuncs.com/api-ws/v1/inference"
FRAME = 3200  # 1600 samples = 100ms @16k


def _flex_bool(v) -> bool:
    return v is True or v == "true"


class GummyEngine(Engine):
    name = "gummy"
    kind = "streaming"
    is_local = False

    def __init__(
        self,
        name: str = "gummy",
        api_key: str = "",
        model: str = "gummy-realtime-v1",
        cost_per_min: float | None = None,
        realtime_factor: float = 2.0,
    ):
        self.name = name
        self.api_key = api_key
        self.model = model
        self._cost = cost_per_min
        self.rt = realtime_factor  # 推流倍速；2.0 = 2x 实时（同 gummycheck）

    def cost_per_min(self) -> float | None:
        return self._cost

    def _run_task_msg(self, task_id: str, sample_rate: int) -> str:
        return json.dumps(
            {
                "header": {"action": "run-task", "task_id": task_id, "streaming": "duplex"},
                "payload": {
                    "task_group": "audio",
                    "task": "asr",
                    "function": "recognition",
                    "model": self.model,
                    "parameters": {
                        "sample_rate": sample_rate,
                        "format": "pcm",
                        "transcription_enabled": True,
                        "translation_enabled": False,
                    },
                    "input": {},
                },
            }
        )

    def _finish_task_msg(self, task_id: str) -> str:
        return json.dumps(
            {
                "header": {"action": "finish-task", "task_id": task_id, "streaming": "duplex"},
                "payload": {"input": {}},
            }
        )

    def transcribe(self, audio_path: str) -> Transcript:
        pcm, sr = load_pcm16(audio_path, 16000)
        audio = pcm.tobytes()
        audio_sec = len(pcm) / sr
        task_id = uuid.uuid4().hex

        started = threading.Event()
        finished = threading.Event()
        results: list[tuple[float, str, bool]] = []  # (rel_t, text, is_final)
        err: list[str | None] = [None]
        first_partial: list[float | None] = [None]
        t0 = time.monotonic()

        try:
            ws = websocket.create_connection(
                WS_URL,
                header=[
                    f"Authorization: bearer {self.api_key}",
                    "X-DashScope-DataInspection: enable",
                ],
                timeout=20,
            )
        except Exception as ex:  # 连接失败
            return Transcript(text="", audio_sec=audio_sec, proc_sec=time.monotonic() - t0, error=f"dial: {ex}")

        def reader() -> None:
            try:
                while True:
                    msg = ws.recv()
                    if not msg:
                        continue
                    ev = json.loads(msg)
                    h = ev.get("header", {})
                    event = h.get("event")
                    if event == "task-started":
                        started.set()
                    elif event == "result-generated":
                        out = ev.get("payload", {}).get("output", {})
                        sen = out.get("transcription") or out.get("sentence")
                        if not sen or not sen.get("text"):
                            continue
                        is_final = _flex_bool(sen.get("sentence_end")) or _flex_bool(sen.get("is_sentence_end"))
                        if first_partial[0] is None:
                            first_partial[0] = time.monotonic() - t0
                        results.append((time.monotonic() - t0, sen["text"], is_final))
                    elif event == "task-finished":
                        finished.set()
                        break
                    elif event == "task-failed":
                        err[0] = h.get("error_message", "task-failed")
                        finished.set()
                        break
            except Exception as ex:
                if not finished.is_set():
                    err[0] = str(ex)
                finished.set()

        rt = threading.Thread(target=reader, daemon=True)
        rt.start()

        try:
            ws.send(self._run_task_msg(task_id, sr))
            if not started.wait(timeout=10):
                ws.close()
                return Transcript(
                    text="", audio_sec=audio_sec, proc_sec=time.monotonic() - t0, error="task-started timeout"
                )

            # 推流：每帧 100ms 音频，按倍速 sleep（2x → 50ms）
            per_frame_sleep = 0.1 / self.rt
            for off in range(0, len(audio), FRAME):
                ws.send_binary(audio[off : off + FRAME])
                time.sleep(per_frame_sleep)
            last_audio_t = time.monotonic()

            ws.send(self._finish_task_msg(task_id))
            finished.wait(timeout=30)
        except Exception as ex:
            err[0] = err[0] or str(ex)
        finally:
            try:
                ws.close()
            except Exception:
                pass

        # gummy 按句给 final；拼接所有 final，无 final 则回退最后一条 partial
        finals = [t for _, t, f in results if f]
        text = "".join(finals) if finals else (results[-1][1] if results else "")

        finalize_sec = None
        final_rel_times = [rel for rel, _, f in results if f]
        if final_rel_times:
            finalize_sec = max(0.0, (t0 + final_rel_times[-1]) - last_audio_t)

        return Transcript(
            text=text,
            audio_sec=audio_sec,
            proc_sec=time.monotonic() - t0,
            first_partial_sec=first_partial[0],
            finalize_sec=finalize_sec,
            error=err[0],
        )