app.py

"""한국어 ASR 모델 비교 데모 (ZeroGPU).

Open Ko-S2S 리더보드의 CER 숫자가 실제 음성에서 어떻게 들리는지 직접 확인하는 데모.
"""

from __future__ import annotations

import html
import time
import traceback

import gradio as gr
import librosa
import numpy as np
import spaces
import torch
from transformers import (
    Wav2Vec2ForCTC,
    Wav2Vec2Processor,
    WhisperForConditionalGeneration,
    WhisperProcessor,
)

from models_meta import LEADERBOARD_URL, MODELS

SR = 16_000
MAX_SEC = 30
DEVICE = "cuda"
DTYPE = torch.float16

# ---------------------------------------------------------------------------
# 모델 로딩
#
# ZeroGPU 규칙: 실제 GPU는 @spaces.GPU 함수 안에서만 잡히지만, 모델은 반드시
# **모듈 레벨에서 .to('cuda')** 로 올려야 한다. 모듈 레벨에서는 PyTorch CUDA
# 에뮬레이션이 동작하고, 데코레이트된 함수 안에서 실제 CUDA로 승격된다.
# (lazy-load / 함수 안에서 .to('cuda') 는 공식적으로 비권장)
# https://huggingface.co/docs/hub/spaces-zerogpu#model-loading
# ---------------------------------------------------------------------------
LOADED: dict[str, dict] = {}
LOAD_ERRORS: dict[str, str] = {}

for _m in MODELS:
    _repo = _m["repo"]
    try:
        if _m["arch"] == "whisper":
            _proc = WhisperProcessor.from_pretrained(_repo)
            _model = WhisperForConditionalGeneration.from_pretrained(
                _repo, torch_dtype=DTYPE
            )
            # 일부 파인튜닝 체크포인트는 generation_config 에 forced_decoder_ids 가
            # 박혀 있어서 generate(language=...) 와 충돌한다. 비워두고 매 호출마다 지정.
            if getattr(_model.generation_config, "forced_decoder_ids", None):
                _model.generation_config.forced_decoder_ids = None
            _model.config.forced_decoder_ids = None
            _model.to(DEVICE).eval()
        else:  # CTC
            _proc = Wav2Vec2Processor.from_pretrained(_repo)
            _model = Wav2Vec2ForCTC.from_pretrained(_repo)  # CTC는 fp32 유지
            _model.to(DEVICE).eval()
        LOADED[_m["key"]] = {"processor": _proc, "model": _model}
        print(f"[load] OK  {_repo}", flush=True)
    except Exception as exc:  # noqa: BLE001
        LOAD_ERRORS[_m["key"]] = f"{type(exc).__name__}: {exc}"
        print(f"[load] FAIL {_repo} -> {exc}", flush=True)
        traceback.print_exc()


# ---------------------------------------------------------------------------
# 추론
# ---------------------------------------------------------------------------
def _load_audio(path: str) -> tuple[np.ndarray, float, bool]:
    wav, _ = librosa.load(path, sr=SR, mono=True)
    dur = len(wav) / SR
    truncated = False
    if dur > MAX_SEC:
        wav = wav[: MAX_SEC * SR]
        truncated = True
        dur = MAX_SEC
    return wav.astype(np.float32), dur, truncated


def _sync() -> None:
    if torch.cuda.is_available():
        torch.cuda.synchronize()


@torch.inference_mode()
def _run_whisper(bundle: dict, wav: np.ndarray) -> str:
    proc, model = bundle["processor"], bundle["model"]
    feats = proc(wav, sampling_rate=SR, return_tensors="pt").input_features
    feats = feats.to(DEVICE, DTYPE)
    kwargs = dict(max_new_tokens=256, num_beams=1, do_sample=False)
    try:
        ids = model.generate(feats, language="ko", task="transcribe", **kwargs)
    except Exception:  # noqa: BLE001 - 언어 강제가 안 먹는 체크포인트 대비
        traceback.print_exc()
        ids = model.generate(feats, **kwargs)
    return proc.batch_decode(ids, skip_special_tokens=True)[0].strip()


@torch.inference_mode()
def _run_ctc(bundle: dict, wav: np.ndarray) -> str:
    proc, model = bundle["processor"], bundle["model"]
    inputs = proc(wav, sampling_rate=SR, return_tensors="pt", padding=True)
    logits = model(inputs.input_values.to(DEVICE)).logits
    pred = torch.argmax(logits, dim=-1)
    return proc.batch_decode(pred)[0].strip()


def _cer(v):
    return "—" if v is None else f"{v:.2f}%"


def _render(rows: list[dict], dur: float, truncated: bool) -> str:
    head = (
        f"<p style='margin:0 0 12px'><b>입력 길이</b> {dur:.1f}초"
        + (
            " <span style='color:#c0392b'>(30초 초과분은 잘렸습니다)</span>"
            if truncated
            else ""
        )
        + " · 시간은 GPU 추론 시간만 측정(모델 로딩 제외)</p>"
    )
    cards = []
    for r in rows:
        badge = (
            "<span style='background:#c0392b;color:#fff;border-radius:4px;"
            "padding:1px 6px;font-size:11px;margin-left:6px'>in-domain 오염</span>"
            if r["in_domain"]
            else ""
        )
        text = r["text"] or "<i style='color:#888'>(빈 결과)</i>"
        cards.append(
            f"""
<div style="border:1px solid #ddd;border-radius:10px;padding:12px 14px;margin-bottom:10px">
  <div style="font-weight:700;font-size:15px">
    <a href="https://huggingface.co/{r['repo']}" target="_blank">{html.escape(r['label'])}</a>
    <span style="color:#666;font-weight:400;font-size:12px"> · {r['params']}</span>{badge}
  </div>
  <div style="font-size:12px;color:#444;margin:6px 0 8px">
    리더보드 CER — 낭독체(Zeroth) <b>{_cer(r['cer_read'])}</b>
    · 자유발화(KsponSpeech) <b>{_cer(r['cer_spont'])}</b>
    &nbsp;|&nbsp; 이번 추론 <b>{r['sec']}</b>
  </div>
  <div style="font-size:16px;line-height:1.6;white-space:pre-wrap">{text}</div>
  <div style="font-size:12px;color:#777;margin-top:8px">{r['note']}</div>
</div>"""
        )
    return head + "".join(cards)


@spaces.GPU(duration=90)
def transcribe(audio_path: str | None):
    if not audio_path:
        return "<p>먼저 마이크로 녹음하거나 오디오 파일을 업로드해 주세요.</p>"

    wav, dur, truncated = _load_audio(audio_path)
    if dur < 0.2:
        return "<p>오디오가 너무 짧습니다. 1초 이상 말해 주세요.</p>"

    rows = []
    for meta in MODELS:
        key = meta["key"]
        row = {
            "repo": meta["repo"],
            "label": meta["label"],
            "params": meta["params"],
            "cer_read": meta["cer_read"],
            "cer_spont": meta["cer_spont"],
            "in_domain": meta["in_domain"],
            "note": meta["note"],
        }
        bundle = LOADED.get(key)
        if bundle is None:
            row["text"] = (
                f"<span style='color:#c0392b'>모델 로딩 실패: "
                f"{html.escape(LOAD_ERRORS.get(key, 'unknown'))}</span>"
            )
            row["sec"] = "—"
            rows.append(row)
            continue
        try:
            _sync()
            t0 = time.perf_counter()
            if meta["arch"] == "whisper":
                text = _run_whisper(bundle, wav)
            else:
                text = _run_ctc(bundle, wav)
            _sync()
            elapsed = time.perf_counter() - t0
            row["text"] = html.escape(text)
            row["sec"] = f"{elapsed:.2f}초 (RTF {elapsed / dur:.3f})"
        except Exception as exc:  # noqa: BLE001
            traceback.print_exc()
            row["text"] = (
                f"<span style='color:#c0392b'>추론 오류: "
                f"{html.escape(type(exc).__name__ + ': ' + str(exc))}</span>"
            )
            row["sec"] = "—"
        rows.append(row)

    return _render(rows, dur, truncated)


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
INTRO = f"""
# 🎙️ 한국어 ASR 모델 비교 — CER 숫자를 귀로 확인하기

**[Open Ko-S2S 리더보드]({LEADERBOARD_URL})** 의 CER 수치가 실제 내 목소리에서
어떤 차이로 나타나는지 직접 들어보는 데모입니다. 같은 오디오를 4개 한국어 ASR 모델에
동시에 넣고, 전사 결과와 추론 시간을 나란히 보여줍니다.

**이 데모가 보여주려는 것**

1. **낭독체 점수가 좋다고 자유발화도 잘하는 게 아닙니다.**
   `whisper-large-v3-turbo`는 낭독체(Zeroth) CER **5.47%**로 이 중 최고지만,
   자유발화(KsponSpeech)에서는 **12.53%**로 무너집니다.
   반대로 73M짜리 `whisper-base-komixv2`는 낭독체 8.01% / 자유발화 **7.71%** 로
   11배 큰 모델을 자유발화에서 이깁니다. 즉 **10배 작은 모델이 실제 대화에서는 더 낫습니다.**
2. **in-domain 점수는 점수가 아닙니다.**
   `wav2vec2-large-xlsr-korean`은 Zeroth 낭독체에서 CER 1.78%를 자칭하지만,
   그 Zeroth train으로 학습된 모델입니다(<span style="color:#c0392b">in-domain 오염</span>).
   빨간 배지가 붙은 모델의 점수는 다른 모델과 같은 선에서 비교하면 안 됩니다.

**해보면 좋은 것**: 원고를 또박또박 읽어보고 → 그다음 아무 준비 없이 말을 더듬으며
"어, 그러니까 그게…" 처럼 자유발화로 말해보세요. 두 경우의 모델 순위가 뒤집히는 걸 볼 수 있습니다.

> ⏱️ ZeroGPU 공용 GPU에서 돌아갑니다. 입력은 **최대 30초**까지만 처리하며,
> 첫 실행은 GPU 할당 대기 때문에 몇 초 더 걸릴 수 있습니다.
"""

FOOT = f"""
---
### 수치 출처
- **낭독체(Zeroth)**: [Open Ko-S2S 리더보드]({LEADERBOARD_URL}) · Zeroth-Korean test (n=457) ·
  CER(space-free) · RTX 4090 실측. `wav2vec2-large-xlsr-korean`만 모델 카드 자체 보고치(리더보드 미수록).
- **자유발화**: 동일 하네스의 KsponSpeech 평가. CTC 모델은 미측정(`—`).
- 이 데모의 "이번 추론" 시간은 ZeroGPU(NVIDIA RTX Pro 6000 Blackwell) 위에서 측정한 값이라
  리더보드의 RTX 4090 RTF와 직접 비교하면 안 됩니다. **모델 간 상대 비교용**으로만 보세요.
- CTC 모델은 띄어쓰기·구두점을 생성하지 않습니다. 전사가 붙어 나오는 건 버그가 아닙니다.
"""

with gr.Blocks(title="한국어 ASR 모델 비교", theme=gr.themes.Soft()) as demo:
    gr.Markdown(INTRO)

    with gr.Row():
        with gr.Column(scale=1):
            audio_in = gr.Audio(
                sources=["microphone", "upload"],
                type="filepath",
                label="음성 입력 (녹음 또는 업로드 · 최대 30초)",
            )
            run_btn = gr.Button("4개 모델로 동시에 전사하기", variant="primary")
            gr.Examples(
                examples=[["examples/zeroth_read_sample.wav"]],
                inputs=audio_in,
                label="낭독체 예시 (Zeroth-Korean test · CC-BY-4.0) — 마이크가 없으면 이걸로 먼저",
            )
            gr.Markdown(
                "예시 정답: **몬터규는 자녀들이 사랑을 제대로 못 받고 크면 "
                "매우 심각한 결과가 초래된다는 결론을 내렸습니다**\n\n"
                "이 문장은 Zeroth-Korean **test** 셋 문장입니다. "
                "`wav2vec2-large-xlsr-korean`은 이 데이터셋 train으로 학습돼 "
                "고유명사 '몬터규'까지 정확히 맞히지만, 학습에서 이 도메인을 보지 않은 "
                "범용 모델들은 '몬토규'처럼 흔들립니다. **이게 in-domain 점수의 정체입니다.**"
            )
            gr.Markdown(
                "**모델 구성**\n\n"
                + "\n".join(
                    f"- `{m['repo']}` ({m['params']})"
                    + ("  ⚠️ in-domain" if m["in_domain"] else "")
                    for m in MODELS
                )
            )
        with gr.Column(scale=2):
            out = gr.HTML(label="전사 결과")

    gr.Markdown(FOOT)

    run_btn.click(transcribe, inputs=audio_in, outputs=out, api_name="transcribe")
    # 녹음을 멈추면 자동 실행 (같은 함수라 별도 API 엔드포인트는 만들지 않음)
    audio_in.stop_recording(transcribe, inputs=audio_in, outputs=out, api_name=False)

if __name__ == "__main__":
    demo.queue(max_size=20).launch()