"""商談文字起こし(Gradio / ZeroGPU)
HuggingFace Spaces の Gradio SDK で動かす。PRO に付いている ZeroGPU を使うと、
文字起こしと話者分離が GPU で走る。GPU が無い環境(自分のPCなど)でも
そのまま CPU で動く。
起動:
python app.py
"""
from __future__ import annotations
import os
import shutil
import tempfile
import time
import traceback
from pathlib import Path
import gradio as gr
import pipeline as pl
# ---------------------------------------------------------------- ZeroGPU
# Spaces では spaces パッケージが入っている。自分のPCには無いので、
# 無ければ「何もしない飾り」に差し替えて同じコードが動くようにする。
try:
import spaces # type: ignore
HAS_SPACES = True
except Exception: # 自分のPCには無い。Spaces でも環境次第で読めないことがある
spaces = None
HAS_SPACES = False
ON_SPACES = os.environ.get("SPACE_ID") is not None
ON_ZERO_GPU = HAS_SPACES and os.environ.get("SPACES_ZERO_GPU") is not None
GPU_DEVICE = "cuda" if ON_ZERO_GPU else ""
# GPU を確保する時間の候補。長く要求するほど空きが見つかりにくいので、
# 録音の長さに合わせていちばん短いもので頼む。
GPU_TIERS = (60, 120, 300)
# 順番待ちで取れなかったときに何回まで粘るか
GPU_RETRIES = int(os.environ.get("GPU_RETRIES", "2"))
def on_gpu(seconds: int):
"""ZeroGPU のときだけ GPU を割り当てる。それ以外は素通し。"""
def deco(fn):
if not ON_ZERO_GPU:
return fn
return spaces.GPU(duration=seconds)(fn)
return deco
def carry_errors(fn):
"""GPU側で起きた失敗を、文字にして持ち帰る。
ZeroGPU は別プロセスで動かした結果を送り返すが、例外はうまく復元できず
「'RuntimeError'」のような中身の無い形になって、原因が分からなくなる。
そこで成否と本文を組にして返し、呼び出し側で組み立て直す。
"""
def wrapped(*args, **kwargs):
try:
return True, fn(*args, **kwargs)
except Exception as exc:
traceback.print_exc()
return False, f"{type(exc).__name__}: {exc}"
return wrapped
def unwrap(result):
ok, payload = result
if not ok:
raise RuntimeError(payload)
return payload
def _transcribe(wav, model_size, prompt, fast, device):
return pl.transcribe_file(Path(wav), model_size, prompt, fast, device=device)
def _diarize(wav, hf_token, num_speakers, device):
return pl.diarize(Path(wav), hf_token, num_speakers, device=device)
# 確保時間ごとに用意しておく。ZeroGPU の指定は関数を作るときに決まるため。
TRANSCRIBE_GPU = {
seconds: on_gpu(seconds)(carry_errors(
lambda wav, model_size, prompt, fast: _transcribe(wav, model_size, prompt, fast, "cuda")
))
for seconds in GPU_TIERS
}
DIARIZE_GPU = {
seconds: on_gpu(seconds)(carry_errors(
lambda wav, hf_token, num_speakers: _diarize(wav, hf_token, num_speakers, "cuda")
))
for seconds in GPU_TIERS
}
def gpu_tier(wav, factor: float) -> int:
"""この録音に必要そうな確保時間。モデル読込ぶんを足して見積もる。"""
try:
length = pl.probe_duration(Path(wav))
except Exception:
length = 0.0
needed = 30 + length * factor
for seconds in GPU_TIERS:
if needed <= seconds:
return seconds
return GPU_TIERS[-1]
def is_busy(message: str) -> bool:
"""GPUの順番待ちで弾かれたか(故障ではなく、混んでいるだけ)。"""
text = message.lower()
return "no gpu" in text or "gpu was available" in text or "quota" in text
def run_on_gpu(table, wav, args, factor, cpu_call, note):
"""GPUで実行し、混んでいて取れなければCPUに落とす。
失敗させるより、遅くても最後まで終わったほうが役に立つ。
"""
seconds = gpu_tier(wav, factor)
last = ""
for attempt in range(GPU_RETRIES + 1):
ok, payload = table[seconds](wav, *args)
if ok:
return payload
last = str(payload)
if not is_busy(last):
raise RuntimeError(last) # 混雑ではない本当の失敗
print(f"[info] GPUの空き待ち({attempt + 1}回目): {last}")
time.sleep(5)
print(f"[info] GPUが取れないのでCPUで処理します: {last}")
note("GPUが混んでいたため、一部はCPUで処理しました(時間がかかります)")
return cpu_call()
def make_workers(note):
"""pipeline に渡す、文字起こしと話者分離の実体を作る。"""
def transcriber(wav, model_size, prompt, fast):
if not ON_ZERO_GPU:
return _transcribe(wav, model_size, prompt, fast, GPU_DEVICE)
return run_on_gpu(
TRANSCRIBE_GPU, wav, (model_size, prompt, fast), 0.15,
lambda: _transcribe(wav, model_size, prompt, fast, "cpu"), note,
)
def diarizer(wav, hf_token, num_speakers):
if not ON_ZERO_GPU:
return _diarize(wav, hf_token, num_speakers, GPU_DEVICE)
return run_on_gpu(
DIARIZE_GPU, wav, (hf_token, num_speakers), 0.20,
lambda: _diarize(wav, hf_token, num_speakers, "cpu"), note,
)
return transcriber, diarizer
# ---------------------------------------------------------------- 置き場所
WORK_ROOT = Path(tempfile.gettempdir()) / "spinthoughts"
WORK_ROOT.mkdir(parents=True, exist_ok=True)
# 商談の音声を必要以上に置いておかないため、古いものは消す
JOB_TTL_SECONDS = 6 * 3600
def sweep_old_jobs() -> None:
now = time.time()
for path in WORK_ROOT.iterdir():
try:
if path.is_dir() and now - path.stat().st_mtime > JOB_TTL_SECONDS:
shutil.rmtree(path, ignore_errors=True)
except OSError:
continue
# ---------------------------------------------------------------- 画面に出す形
STATUS_MARK = {"transcribed": "", "skipped": "除外", "error": "失敗"}
def esc(text) -> str:
return (
str(text)
.replace("&", "&")
.replace("<", "<")
.replace(">", ">")
.replace('"', """)
)
def audio_url(path: Path) -> str:
return "/gradio_api/file=" + Path(path).as_posix()
def render(results) -> str:
"""録音ごとのカード。時刻を押すとその位置から音声が鳴る。"""
if not results:
return ""
# 時刻が押せることは見ただけでは分からないので、最初に書いておく
guide = ""
if any(r.playback for r in results):
guide = (
'
発話の左にある時刻を押すと、'
'その場面から録音が流れます。聞き取りにくい行だけ確かめられます。
'
)
cards = []
for i, entry in enumerate(results, 1):
audio_id = "audio-%d" % i
tag = STATUS_MARK.get(entry.status, "")
meta = " · ".join(x for x in [
pl.hhmmss(entry.duration),
entry.recorded_at,
("%d発話" % len(entry.utterances)) if entry.status == "transcribed" else "",
entry.reason or "",
] if x)
player = ""
if entry.playback:
player = (
''
% (audio_id, audio_url(entry.playback))
)
talk = pl.speaking_time(entry)
talk_html = ""
if len(talk) > 1:
parts = [
"%s %d分%02d秒・%d回"
% (esc(name), int(secs // 60), int(secs % 60), count)
for name, secs, count in talk
]
talk_html = '
%s
' % "".join(parts)
palette = {}
lines = []
for u in entry.utterances:
if u.speaker not in palette:
palette[u.speaker] = len(palette) % 5
flag = ""
if u.needs_review:
flag = '要確認' % esc(u.review_note)
who = '%s' % esc(u.speaker) if u.speaker else ""
lines.append(
'
'
''
''
'%s%s%s'
"
"
% (audio_id, u.start, pl.hhmmss(u.start), palette[u.speaker],
who, esc(u.text), flag)
)
if lines:
body = '