Restructure into asr/tts packages and add TTS tab skeleton
Split the single ASR gateway into app/main.py (entrypoint) + app/core (shared env helpers) + app/asr (all existing ASR logic, unchanged behavior) + app/tts (placeholder router/config, no engine yet), so a real TTS backend can be dropped in later without reshuffling ASR code. API moved under /asr/* and /tts/* prefixes; UI gained a top-level ASR/TTS tab switcher built on a reusable nested .tab-group mechanism, with JS split into common.js/asr.js/tts.js. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
du5t
0
app/asr/__init__.py
Normal file
0
app/asr/__init__.py
Normal file
40
app/asr/config.py
Normal file
40
app/asr/config.py
Normal file
@@ -0,0 +1,40 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from core.config import env_str
|
||||
|
||||
ASR_BASE_DIR = Path(env_str("ASR_BASE_DIR", "/srv/asr"))
|
||||
MODEL_CACHE = Path(env_str("WHISPER_CACHE_DIR", str(ASR_BASE_DIR / "models-cache")))
|
||||
UPLOAD_DIR = ASR_BASE_DIR / "uploads"
|
||||
RESULT_DIR = ASR_BASE_DIR / "results"
|
||||
CUSTOM_MODEL_DIR = ASR_BASE_DIR / "custom-models"
|
||||
|
||||
DEVICE = env_str("ASR_DEVICE", "cuda")
|
||||
COMPUTE_TYPE = env_str("ASR_COMPUTE_TYPE", "float16")
|
||||
DEFAULT_BACKEND = env_str("DEFAULT_BACKEND", "faster-whisper")
|
||||
DEFAULT_MODEL = env_str("DEFAULT_MODEL", "large-v3")
|
||||
DEFAULT_LANGUAGE = env_str("DEFAULT_LANGUAGE", "ko")
|
||||
|
||||
FASTER_WHISPER_URL = env_str("FASTER_WHISPER_URL", "http://127.0.0.1:8001")
|
||||
QWEN3_URL = env_str("QWEN3_URL", "http://127.0.0.1:8004")
|
||||
|
||||
PYANNOTE_HF_TOKEN = env_str("PYANNOTE_HF_TOKEN", "")
|
||||
|
||||
|
||||
def ensure_runtime_dirs() -> None:
|
||||
for p in [ASR_BASE_DIR, MODEL_CACHE, UPLOAD_DIR, RESULT_DIR, CUSTOM_MODEL_DIR]:
|
||||
p.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
def resolve_custom_model_path(custom_model_path: Optional[str]) -> Optional[str]:
|
||||
if not custom_model_path:
|
||||
return None
|
||||
raw = custom_model_path.strip()
|
||||
if not raw:
|
||||
return None
|
||||
candidate = Path(raw)
|
||||
if candidate.is_absolute():
|
||||
return str(candidate)
|
||||
return str((CUSTOM_MODEL_DIR / candidate).resolve())
|
||||
354
app/asr/router.py
Normal file
354
app/asr/router.py
Normal file
@@ -0,0 +1,354 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import json
|
||||
import struct
|
||||
import wave
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import httpx
|
||||
from fastapi import APIRouter, File, Form, HTTPException, UploadFile, WebSocket, WebSocketDisconnect
|
||||
from fastapi.responses import FileResponse
|
||||
|
||||
from asr.config import (
|
||||
DEFAULT_BACKEND,
|
||||
DEFAULT_LANGUAGE,
|
||||
DEFAULT_MODEL,
|
||||
FASTER_WHISPER_URL,
|
||||
QWEN3_URL,
|
||||
RESULT_DIR,
|
||||
UPLOAD_DIR,
|
||||
ensure_runtime_dirs,
|
||||
)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
BACKENDS = {
|
||||
"faster-whisper": FASTER_WHISPER_URL,
|
||||
"qwen3": QWEN3_URL,
|
||||
}
|
||||
|
||||
REALTIME_SAMPLE_RATE = 16000
|
||||
REALTIME_CHUNK_SECONDS = 3
|
||||
|
||||
|
||||
# ─── Config ────────────────────────────────────────────────────────────────────
|
||||
|
||||
@router.get("/config")
|
||||
def config() -> Dict[str, Any]:
|
||||
return {
|
||||
"default_backend": DEFAULT_BACKEND,
|
||||
"default_model": DEFAULT_MODEL,
|
||||
"default_language": DEFAULT_LANGUAGE,
|
||||
"backends": {
|
||||
"faster-whisper": {
|
||||
"models": ["tiny", "base", "small", "medium", "large-v3", "large-v2", "turbo"],
|
||||
},
|
||||
"qwen3": {
|
||||
"models": ["Qwen/Qwen3-ASR-2B", "Qwen/Qwen3-ASR-8B"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# ─── History ──────────────────────────────────────────────────────────────────
|
||||
|
||||
@router.get("/history")
|
||||
def list_history() -> List[Dict[str, Any]]:
|
||||
ensure_runtime_dirs()
|
||||
records = []
|
||||
for p in sorted(RESULT_DIR.glob("*.json"), reverse=True):
|
||||
try:
|
||||
data = json.loads(p.read_text(encoding="utf-8"))
|
||||
upload_path = data.get("upload_file", "")
|
||||
upload_name = Path(upload_path).name if upload_path else ""
|
||||
records.append({
|
||||
"id": p.stem,
|
||||
"filename": data.get("_filename", upload_name),
|
||||
"backend": data.get("gateway_backend", data.get("backend", "")),
|
||||
"model": data.get("model", ""),
|
||||
"language": data.get("language", ""),
|
||||
"duration": data.get("duration"),
|
||||
"text_preview": (data.get("text", "") or "")[:120],
|
||||
"upload_file": upload_name,
|
||||
"diarized": data.get("diarized", False),
|
||||
})
|
||||
except Exception:
|
||||
continue
|
||||
return records
|
||||
|
||||
|
||||
@router.get("/history/{record_id}")
|
||||
def get_history(record_id: str) -> Dict[str, Any]:
|
||||
safe_id = Path(record_id).name
|
||||
p = RESULT_DIR / f"{safe_id}.json"
|
||||
if not p.exists():
|
||||
raise HTTPException(status_code=404, detail="Record not found")
|
||||
return json.loads(p.read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
@router.delete("/history/{record_id}")
|
||||
def delete_history(record_id: str) -> Dict[str, str]:
|
||||
safe_id = Path(record_id).name
|
||||
p = RESULT_DIR / f"{safe_id}.json"
|
||||
if not p.exists():
|
||||
raise HTTPException(status_code=404, detail="Record not found")
|
||||
|
||||
# Also remove upload file if referenced
|
||||
try:
|
||||
data = json.loads(p.read_text(encoding="utf-8"))
|
||||
upload_path = data.get("upload_file", "")
|
||||
if upload_path:
|
||||
up = Path(upload_path)
|
||||
if up.exists() and up.is_relative_to(UPLOAD_DIR):
|
||||
up.unlink(missing_ok=True)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
p.unlink(missing_ok=True)
|
||||
return {"status": "deleted", "id": safe_id}
|
||||
|
||||
|
||||
@router.get("/uploads/{filename}")
|
||||
def serve_upload(filename: str) -> FileResponse:
|
||||
safe_name = Path(filename).name
|
||||
p = UPLOAD_DIR / safe_name
|
||||
if not p.exists():
|
||||
raise HTTPException(status_code=404, detail="File not found")
|
||||
return FileResponse(p)
|
||||
|
||||
|
||||
# ─── File Transcription ───────────────────────────────────────────────────────
|
||||
|
||||
@router.post("/transcribe")
|
||||
async def transcribe(
|
||||
file: UploadFile = File(...),
|
||||
backend: str = Form(DEFAULT_BACKEND),
|
||||
model: str = Form(DEFAULT_MODEL),
|
||||
custom_model_path: Optional[str] = Form(None),
|
||||
language: Optional[str] = Form(None),
|
||||
task: str = Form("transcribe"),
|
||||
# faster-whisper options
|
||||
beam_size: int = Form(5),
|
||||
temperature: float = Form(0.0),
|
||||
word_timestamps: bool = Form(False),
|
||||
diarize: bool = Form(False),
|
||||
num_speakers: Optional[int] = Form(None),
|
||||
min_speakers: Optional[int] = Form(None),
|
||||
max_speakers: Optional[int] = Form(None),
|
||||
no_repeat_ngram_size: int = Form(0),
|
||||
repetition_penalty: float = Form(1.0),
|
||||
compression_ratio_threshold: float = Form(2.4),
|
||||
log_prob_threshold: float = Form(-1.0),
|
||||
no_speech_threshold: float = Form(0.6),
|
||||
condition_on_previous_text: bool = Form(True),
|
||||
) -> Dict[str, Any]:
|
||||
if backend not in BACKENDS:
|
||||
raise HTTPException(status_code=400, detail=f"Unsupported backend: {backend}")
|
||||
|
||||
ensure_runtime_dirs()
|
||||
original_filename = file.filename or "upload.bin"
|
||||
suffix = Path(original_filename).suffix or ".bin"
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
|
||||
saved_upload = UPLOAD_DIR / f"{timestamp}{suffix}"
|
||||
|
||||
content = await file.read()
|
||||
saved_upload.write_bytes(content)
|
||||
|
||||
try:
|
||||
worker_url = BACKENDS[backend]
|
||||
effective_language = language or DEFAULT_LANGUAGE
|
||||
|
||||
if backend == "qwen3":
|
||||
data: Dict[str, str] = {
|
||||
"model": model,
|
||||
"language": effective_language,
|
||||
"task": task,
|
||||
}
|
||||
else:
|
||||
# faster-whisper
|
||||
data = {
|
||||
"model": model,
|
||||
"custom_model_path": custom_model_path or "",
|
||||
"language": effective_language,
|
||||
"task": task,
|
||||
"beam_size": str(beam_size),
|
||||
"temperature": str(temperature),
|
||||
"word_timestamps": str(word_timestamps).lower(),
|
||||
"diarize": str(diarize).lower(),
|
||||
"no_repeat_ngram_size": str(no_repeat_ngram_size),
|
||||
"repetition_penalty": str(repetition_penalty),
|
||||
"compression_ratio_threshold": str(compression_ratio_threshold),
|
||||
"log_prob_threshold": str(log_prob_threshold),
|
||||
"no_speech_threshold": str(no_speech_threshold),
|
||||
"condition_on_previous_text": str(condition_on_previous_text).lower(),
|
||||
}
|
||||
if num_speakers is not None:
|
||||
data["num_speakers"] = str(num_speakers)
|
||||
if min_speakers is not None:
|
||||
data["min_speakers"] = str(min_speakers)
|
||||
if max_speakers is not None:
|
||||
data["max_speakers"] = str(max_speakers)
|
||||
|
||||
files_payload = {
|
||||
"file": (original_filename, content, file.content_type or "application/octet-stream")
|
||||
}
|
||||
async with httpx.AsyncClient(timeout=httpx.Timeout(3600.0, connect=60.0)) as client:
|
||||
response = await client.post(f"{worker_url}/transcribe", data=data, files=files_payload)
|
||||
|
||||
if response.status_code >= 400:
|
||||
raise HTTPException(status_code=response.status_code, detail=f"Backend error: {response.text}")
|
||||
|
||||
payload = response.json()
|
||||
payload["gateway_backend"] = backend
|
||||
payload["_filename"] = original_filename
|
||||
payload["_timestamp"] = timestamp
|
||||
payload["upload_file"] = str(saved_upload)
|
||||
|
||||
result_path = RESULT_DIR / f"{timestamp}.json"
|
||||
result_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
payload["result_file"] = str(result_path)
|
||||
return payload
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=500, detail=f"Gateway error: {exc}") from exc
|
||||
|
||||
|
||||
# ─── Real-time WebSocket (faster-whisper only) ────────────────────────────────
|
||||
|
||||
def _pcm_float32_to_wav(pcm_bytes: bytes, sample_rate: int = 16000) -> bytes:
|
||||
n = len(pcm_bytes) // 4
|
||||
try:
|
||||
import numpy as np
|
||||
arr = np.frombuffer(pcm_bytes, dtype=np.float32)
|
||||
arr_i16 = np.clip(arr * 32767.0, -32768, 32767).astype(np.int16)
|
||||
raw16 = arr_i16.tobytes()
|
||||
except ImportError:
|
||||
raw16 = b"".join(
|
||||
struct.pack("<h", max(-32768, min(32767, int(struct.unpack("<f", pcm_bytes[i*4:(i+1)*4])[0] * 32767))))
|
||||
for i in range(n)
|
||||
)
|
||||
buf = io.BytesIO()
|
||||
with wave.open(buf, "wb") as wf:
|
||||
wf.setnchannels(1)
|
||||
wf.setsampwidth(2)
|
||||
wf.setframerate(sample_rate)
|
||||
wf.writeframes(raw16)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
async def _send_chunk(wav_bytes: bytes, *, model: str, language: str, beam_size: int,
|
||||
no_repeat_ngram_size: int, repetition_penalty: float,
|
||||
compression_ratio_threshold: float, log_prob_threshold: float,
|
||||
no_speech_threshold: float, condition_on_previous_text: bool) -> Dict[str, Any]:
|
||||
data = {
|
||||
"model": model,
|
||||
"language": language,
|
||||
"task": "transcribe",
|
||||
"beam_size": str(beam_size),
|
||||
"temperature": "0",
|
||||
"word_timestamps": "false",
|
||||
"diarize": "false",
|
||||
"no_repeat_ngram_size": str(no_repeat_ngram_size),
|
||||
"repetition_penalty": str(repetition_penalty),
|
||||
"compression_ratio_threshold": str(compression_ratio_threshold),
|
||||
"log_prob_threshold": str(log_prob_threshold),
|
||||
"no_speech_threshold": str(no_speech_threshold),
|
||||
"condition_on_previous_text": str(condition_on_previous_text).lower(),
|
||||
}
|
||||
files_payload = {"file": ("chunk.wav", wav_bytes, "audio/wav")}
|
||||
async with httpx.AsyncClient(timeout=httpx.Timeout(300.0, connect=30.0)) as client:
|
||||
resp = await client.post(f"{FASTER_WHISPER_URL}/transcribe", data=data, files=files_payload)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
|
||||
@router.websocket("/ws/realtime")
|
||||
async def realtime_ws(websocket: WebSocket) -> None:
|
||||
await websocket.accept()
|
||||
|
||||
audio_buf = bytearray()
|
||||
partial_texts: List[str] = []
|
||||
all_segments: List[Dict[str, Any]] = []
|
||||
time_offset = 0.0
|
||||
|
||||
try:
|
||||
cfg = json.loads(await websocket.receive_text())
|
||||
|
||||
model = cfg.get("model", DEFAULT_MODEL)
|
||||
language = cfg.get("language") or DEFAULT_LANGUAGE
|
||||
beam_size = int(cfg.get("beam_size", 5))
|
||||
no_repeat_ngram_size = int(cfg.get("no_repeat_ngram_size", 0))
|
||||
repetition_penalty = float(cfg.get("repetition_penalty", 1.0))
|
||||
compression_ratio_threshold = float(cfg.get("compression_ratio_threshold", 2.4))
|
||||
log_prob_threshold = float(cfg.get("log_prob_threshold", -1.0))
|
||||
no_speech_threshold = float(cfg.get("no_speech_threshold", 0.6))
|
||||
condition_on_previous_text = bool(cfg.get("condition_on_previous_text", True))
|
||||
chunk_seconds = int(cfg.get("chunk_seconds", REALTIME_CHUNK_SECONDS))
|
||||
chunk_bytes = chunk_seconds * REALTIME_SAMPLE_RATE * 4
|
||||
|
||||
await websocket.send_json({"type": "ready"})
|
||||
|
||||
async def process_buffer(buf: bytearray) -> None:
|
||||
nonlocal time_offset
|
||||
if len(buf) < REALTIME_SAMPLE_RATE * 4 * 0.3:
|
||||
return
|
||||
wav = _pcm_float32_to_wav(bytes(buf))
|
||||
result = await _send_chunk(
|
||||
wav, model=model, language=language, beam_size=beam_size,
|
||||
no_repeat_ngram_size=no_repeat_ngram_size,
|
||||
repetition_penalty=repetition_penalty,
|
||||
compression_ratio_threshold=compression_ratio_threshold,
|
||||
log_prob_threshold=log_prob_threshold,
|
||||
no_speech_threshold=no_speech_threshold,
|
||||
condition_on_previous_text=condition_on_previous_text,
|
||||
)
|
||||
text = result.get("text", "").strip()
|
||||
segs = result.get("segments", [])
|
||||
for seg in segs:
|
||||
seg["start"] = round(seg.get("start", 0) + time_offset, 3)
|
||||
seg["end"] = round(seg.get("end", 0) + time_offset, 3)
|
||||
time_offset += len(buf) / 4 / REALTIME_SAMPLE_RATE
|
||||
if text:
|
||||
partial_texts.append(text)
|
||||
all_segments.extend(segs)
|
||||
await websocket.send_json({
|
||||
"type": "partial",
|
||||
"text": text,
|
||||
"segments": segs,
|
||||
"full_text": " ".join(x for x in partial_texts if x),
|
||||
})
|
||||
|
||||
while True:
|
||||
data = await websocket.receive()
|
||||
if data.get("type") == "websocket.disconnect":
|
||||
break
|
||||
if "bytes" in data:
|
||||
audio_buf.extend(data["bytes"])
|
||||
while len(audio_buf) >= chunk_bytes:
|
||||
await process_buffer(bytearray(audio_buf[:chunk_bytes]))
|
||||
audio_buf = audio_buf[chunk_bytes:]
|
||||
elif "text" in data:
|
||||
msg = json.loads(data["text"])
|
||||
if msg.get("cmd") == "stop":
|
||||
if audio_buf:
|
||||
await process_buffer(audio_buf)
|
||||
await websocket.send_json({
|
||||
"type": "final",
|
||||
"text": " ".join(x for x in partial_texts if x),
|
||||
"segments": all_segments,
|
||||
"diarized": False,
|
||||
})
|
||||
break
|
||||
|
||||
except WebSocketDisconnect:
|
||||
pass
|
||||
except Exception as e:
|
||||
try:
|
||||
await websocket.send_json({"type": "error", "detail": str(e)})
|
||||
except Exception:
|
||||
pass
|
||||
273
app/asr/workers/faster_whisper_worker.py
Normal file
273
app/asr/workers/faster_whisper_worker.py
Normal file
@@ -0,0 +1,273 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
|
||||
from faster_whisper import WhisperModel
|
||||
|
||||
import sys
|
||||
sys.path.insert(0, "/app")
|
||||
from asr.config import (
|
||||
COMPUTE_TYPE,
|
||||
CUSTOM_MODEL_DIR,
|
||||
DEFAULT_MODEL,
|
||||
DEVICE,
|
||||
MODEL_CACHE,
|
||||
PYANNOTE_HF_TOKEN,
|
||||
ensure_runtime_dirs,
|
||||
resolve_custom_model_path,
|
||||
)
|
||||
|
||||
app = FastAPI(title="ASR Faster-Whisper Worker")
|
||||
_MODEL_CACHE: Dict[str, WhisperModel] = {}
|
||||
_DIARIZATION_PIPELINE: Any = None
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
def startup() -> None:
|
||||
ensure_runtime_dirs()
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health() -> Dict[str, Any]:
|
||||
return {
|
||||
"status": "ok",
|
||||
"device": DEVICE,
|
||||
"compute_type": COMPUTE_TYPE,
|
||||
"diarization_available": bool(PYANNOTE_HF_TOKEN),
|
||||
}
|
||||
|
||||
|
||||
@app.post("/transcribe")
|
||||
async def transcribe(
|
||||
file: UploadFile = File(...),
|
||||
model: str = Form(DEFAULT_MODEL),
|
||||
custom_model_path: Optional[str] = Form(None),
|
||||
language: Optional[str] = Form(None),
|
||||
task: str = Form("transcribe"),
|
||||
beam_size: int = Form(5),
|
||||
temperature: float = Form(0.0),
|
||||
word_timestamps: bool = Form(False),
|
||||
diarize: bool = Form(False),
|
||||
num_speakers: Optional[int] = Form(None),
|
||||
min_speakers: Optional[int] = Form(None),
|
||||
max_speakers: Optional[int] = Form(None),
|
||||
no_repeat_ngram_size: int = Form(0),
|
||||
repetition_penalty: float = Form(1.0),
|
||||
compression_ratio_threshold: float = Form(2.4),
|
||||
log_prob_threshold: float = Form(-1.0),
|
||||
no_speech_threshold: float = Form(0.6),
|
||||
condition_on_previous_text: bool = Form(True),
|
||||
) -> Dict[str, Any]:
|
||||
ensure_runtime_dirs()
|
||||
suffix = Path(file.filename or "upload.bin").suffix or ".bin"
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
|
||||
tmp.write(await file.read())
|
||||
tmp_path = Path(tmp.name)
|
||||
|
||||
try:
|
||||
resolved_model = resolve_custom_model_path(custom_model_path) or model
|
||||
result = _transcribe(
|
||||
audio_path=tmp_path,
|
||||
model_name=resolved_model,
|
||||
language=language or None,
|
||||
task=task,
|
||||
beam_size=beam_size,
|
||||
temperature=temperature,
|
||||
word_timestamps=word_timestamps,
|
||||
no_repeat_ngram_size=no_repeat_ngram_size,
|
||||
repetition_penalty=repetition_penalty,
|
||||
compression_ratio_threshold=compression_ratio_threshold,
|
||||
log_prob_threshold=log_prob_threshold,
|
||||
no_speech_threshold=no_speech_threshold,
|
||||
condition_on_previous_text=condition_on_previous_text,
|
||||
)
|
||||
if diarize:
|
||||
result = _apply_diarization(
|
||||
audio_path=tmp_path,
|
||||
result=result,
|
||||
num_speakers=num_speakers,
|
||||
min_speakers=min_speakers,
|
||||
max_speakers=max_speakers,
|
||||
)
|
||||
return result
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=500, detail=f"Worker failure: {exc}") from exc
|
||||
finally:
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
|
||||
|
||||
def _load_model(model_name: str) -> WhisperModel:
|
||||
if model_name not in _MODEL_CACHE:
|
||||
_MODEL_CACHE[model_name] = WhisperModel(
|
||||
model_name,
|
||||
device=DEVICE,
|
||||
compute_type=COMPUTE_TYPE,
|
||||
download_root=str(MODEL_CACHE),
|
||||
)
|
||||
return _MODEL_CACHE[model_name]
|
||||
|
||||
|
||||
def _transcribe(
|
||||
audio_path: Path,
|
||||
model_name: str,
|
||||
language: Optional[str],
|
||||
task: str,
|
||||
beam_size: int,
|
||||
temperature: float,
|
||||
word_timestamps: bool,
|
||||
no_repeat_ngram_size: int,
|
||||
repetition_penalty: float,
|
||||
compression_ratio_threshold: float,
|
||||
log_prob_threshold: float,
|
||||
no_speech_threshold: float,
|
||||
condition_on_previous_text: bool,
|
||||
) -> Dict[str, Any]:
|
||||
model = _load_model(model_name)
|
||||
|
||||
temperatures: Union[float, List[float]] = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0] if temperature == 0.0 else temperature
|
||||
|
||||
segments_iter, info = model.transcribe(
|
||||
str(audio_path),
|
||||
language=language,
|
||||
task=task,
|
||||
beam_size=beam_size,
|
||||
temperature=temperatures,
|
||||
word_timestamps=word_timestamps,
|
||||
no_repeat_ngram_size=no_repeat_ngram_size,
|
||||
repetition_penalty=repetition_penalty,
|
||||
compression_ratio_threshold=compression_ratio_threshold,
|
||||
log_prob_threshold=log_prob_threshold,
|
||||
no_speech_threshold=no_speech_threshold,
|
||||
condition_on_previous_text=condition_on_previous_text,
|
||||
)
|
||||
|
||||
segment_list = []
|
||||
full_text = []
|
||||
for seg in segments_iter:
|
||||
seg_dict: Dict[str, Any] = {
|
||||
"id": seg.id,
|
||||
"start": round(float(seg.start), 3),
|
||||
"end": round(float(seg.end), 3),
|
||||
"text": seg.text,
|
||||
"avg_logprob": round(float(seg.avg_logprob), 4) if seg.avg_logprob is not None else None,
|
||||
"compression_ratio": round(float(seg.compression_ratio), 4) if seg.compression_ratio is not None else None,
|
||||
"no_speech_prob": round(float(seg.no_speech_prob), 4) if seg.no_speech_prob is not None else None,
|
||||
}
|
||||
if word_timestamps and getattr(seg, "words", None):
|
||||
seg_dict["words"] = [
|
||||
{
|
||||
"start": round(float(w.start), 3),
|
||||
"end": round(float(w.end), 3),
|
||||
"word": w.word,
|
||||
"probability": round(float(w.probability), 4),
|
||||
}
|
||||
for w in seg.words
|
||||
]
|
||||
segment_list.append(seg_dict)
|
||||
full_text.append(seg.text.strip())
|
||||
|
||||
duration = None
|
||||
try:
|
||||
duration = round(float(info.duration), 3)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return {
|
||||
"backend": "faster-whisper",
|
||||
"model": model_name,
|
||||
"language": getattr(info, "language", language),
|
||||
"language_probability": round(float(getattr(info, "language_probability", 0) or 0), 4),
|
||||
"duration": duration,
|
||||
"text": " ".join(x for x in full_text if x),
|
||||
"segments": segment_list,
|
||||
"diarized": False,
|
||||
}
|
||||
|
||||
|
||||
def _get_diarization_pipeline():
|
||||
global _DIARIZATION_PIPELINE
|
||||
if _DIARIZATION_PIPELINE is None:
|
||||
try:
|
||||
from pyannote.audio import Pipeline
|
||||
import torch
|
||||
except ImportError as exc:
|
||||
raise RuntimeError("pyannote.audio가 설치되지 않았습니다.") from exc
|
||||
if not PYANNOTE_HF_TOKEN:
|
||||
raise RuntimeError("화자 분리를 사용하려면 PYANNOTE_HF_TOKEN 환경변수를 설정하세요.")
|
||||
_DIARIZATION_PIPELINE = Pipeline.from_pretrained(
|
||||
"pyannote/speaker-diarization-3.1",
|
||||
use_auth_token=PYANNOTE_HF_TOKEN,
|
||||
)
|
||||
if DEVICE == "cuda":
|
||||
import torch
|
||||
_DIARIZATION_PIPELINE = _DIARIZATION_PIPELINE.to(torch.device("cuda"))
|
||||
return _DIARIZATION_PIPELINE
|
||||
|
||||
|
||||
def _apply_diarization(
|
||||
audio_path: Path,
|
||||
result: Dict[str, Any],
|
||||
num_speakers: Optional[int],
|
||||
min_speakers: Optional[int],
|
||||
max_speakers: Optional[int],
|
||||
) -> Dict[str, Any]:
|
||||
pipeline = _get_diarization_pipeline()
|
||||
kwargs: Dict[str, Any] = {}
|
||||
if num_speakers is not None:
|
||||
kwargs["num_speakers"] = num_speakers
|
||||
else:
|
||||
if min_speakers is not None:
|
||||
kwargs["min_speakers"] = min_speakers
|
||||
if max_speakers is not None:
|
||||
kwargs["max_speakers"] = max_speakers
|
||||
|
||||
diarization = pipeline(str(audio_path), **kwargs)
|
||||
segments = result.get("segments", [])
|
||||
for seg in segments:
|
||||
seg_start = float(seg.get("start", 0))
|
||||
seg_end = float(seg.get("end", 0))
|
||||
speaker_times: Dict[str, float] = {}
|
||||
for turn, _, speaker in diarization.itertracks(yield_label=True):
|
||||
overlap_start = max(seg_start, turn.start)
|
||||
overlap_end = min(seg_end, turn.end)
|
||||
if overlap_end > overlap_start:
|
||||
speaker_times[speaker] = speaker_times.get(speaker, 0.0) + (overlap_end - overlap_start)
|
||||
seg["speaker"] = max(speaker_times, key=speaker_times.get) if speaker_times else "UNKNOWN"
|
||||
|
||||
lines = []
|
||||
current_speaker: Optional[str] = None
|
||||
current_texts: list = []
|
||||
for seg in segments:
|
||||
speaker = seg.get("speaker", "UNKNOWN")
|
||||
text = seg.get("text", "").strip()
|
||||
if not text:
|
||||
continue
|
||||
if speaker != current_speaker:
|
||||
if current_texts and current_speaker is not None:
|
||||
lines.append(f"[{current_speaker}]: {' '.join(current_texts)}")
|
||||
current_speaker = speaker
|
||||
current_texts = [text]
|
||||
else:
|
||||
current_texts.append(text)
|
||||
if current_texts and current_speaker is not None:
|
||||
lines.append(f"[{current_speaker}]: {' '.join(current_texts)}")
|
||||
|
||||
result["text"] = "\n".join(lines)
|
||||
result["diarized"] = True
|
||||
return result
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", default="0.0.0.0")
|
||||
parser.add_argument("--port", type=int, default=8001)
|
||||
args = parser.parse_args()
|
||||
ensure_runtime_dirs()
|
||||
uvicorn.run(app, host=args.host, port=args.port)
|
||||
178
app/asr/workers/qwen3_worker.py
Normal file
178
app/asr/workers/qwen3_worker.py
Normal file
@@ -0,0 +1,178 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
import sys
|
||||
sys.path.insert(0, "/app")
|
||||
from asr.config import DEVICE, MODEL_CACHE, ensure_runtime_dirs
|
||||
|
||||
app = FastAPI(title="ASR Qwen3 Worker")
|
||||
|
||||
_PIPE_CACHE: Dict[str, Any] = {}
|
||||
|
||||
LANG_MAP = {
|
||||
"ko": "korean", "en": "english", "ja": "japanese", "zh": "chinese",
|
||||
"fr": "french", "de": "german", "es": "spanish", "ru": "russian",
|
||||
"vi": "vietnamese", "th": "thai", "ar": "arabic", "pt": "portuguese",
|
||||
"it": "italian", "nl": "dutch", "pl": "polish", "tr": "turkish",
|
||||
}
|
||||
|
||||
|
||||
def _log(msg: str) -> None:
|
||||
print(f"[qwen3] {msg}", flush=True)
|
||||
|
||||
|
||||
def _free_gpu(*objs: Any) -> None:
|
||||
for obj in objs:
|
||||
try:
|
||||
del obj
|
||||
except Exception:
|
||||
pass
|
||||
gc.collect()
|
||||
try:
|
||||
import torch
|
||||
if DEVICE == "cuda" and torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _load_pipe(model_id: str) -> Any:
|
||||
if model_id in _PIPE_CACHE:
|
||||
return _PIPE_CACHE[model_id]
|
||||
|
||||
import torch
|
||||
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
|
||||
|
||||
_log(f"loading model {model_id}")
|
||||
dtype = torch.float16 if DEVICE == "cuda" else torch.float32
|
||||
|
||||
try:
|
||||
model = AutoModelForSpeechSeq2Seq.from_pretrained(
|
||||
model_id,
|
||||
torch_dtype=dtype,
|
||||
low_cpu_mem_usage=True,
|
||||
cache_dir=str(MODEL_CACHE),
|
||||
)
|
||||
model.to(DEVICE)
|
||||
processor = AutoProcessor.from_pretrained(model_id, cache_dir=str(MODEL_CACHE))
|
||||
pipe = pipeline(
|
||||
"automatic-speech-recognition",
|
||||
model=model,
|
||||
tokenizer=processor.tokenizer,
|
||||
feature_extractor=processor.feature_extractor,
|
||||
torch_dtype=dtype,
|
||||
device=DEVICE,
|
||||
)
|
||||
except Exception:
|
||||
# Fallback: let pipeline handle loading (uses device_map)
|
||||
_log("direct load failed, trying pipeline auto-load")
|
||||
pipe = pipeline(
|
||||
"automatic-speech-recognition",
|
||||
model=model_id,
|
||||
torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
|
||||
device_map="auto" if DEVICE == "cuda" else None,
|
||||
model_kwargs={"cache_dir": str(MODEL_CACHE)},
|
||||
)
|
||||
|
||||
_PIPE_CACHE[model_id] = pipe
|
||||
_log(f"model {model_id} loaded")
|
||||
return pipe
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
def startup() -> None:
|
||||
ensure_runtime_dirs()
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health() -> Dict[str, Any]:
|
||||
return {"status": "ok", "device": DEVICE, "loaded_models": list(_PIPE_CACHE.keys())}
|
||||
|
||||
|
||||
@app.post("/transcribe")
|
||||
async def transcribe(
|
||||
file: UploadFile = File(...),
|
||||
model: str = Form("Qwen/Qwen3-ASR-2B"),
|
||||
language: Optional[str] = Form("ko"),
|
||||
task: str = Form("transcribe"),
|
||||
) -> JSONResponse:
|
||||
ensure_runtime_dirs()
|
||||
suffix = Path(file.filename or "audio.bin").suffix or ".wav"
|
||||
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
|
||||
tmp.write(await file.read())
|
||||
tmp_path = Path(tmp.name)
|
||||
|
||||
try:
|
||||
pipe = _load_pipe(model)
|
||||
|
||||
lang_code = (language or "ko").strip().lower()
|
||||
lang_name = LANG_MAP.get(lang_code, lang_code)
|
||||
|
||||
generate_kwargs: Dict[str, Any] = {"task": task}
|
||||
if lang_code:
|
||||
generate_kwargs["language"] = lang_name
|
||||
|
||||
_log(f"transcribing language={lang_code} model={model}")
|
||||
raw = pipe(
|
||||
str(tmp_path),
|
||||
return_timestamps=True,
|
||||
generate_kwargs=generate_kwargs,
|
||||
)
|
||||
|
||||
# raw may be {"text": "...", "chunks": [...]} or {"text": "..."}
|
||||
full_text: str = raw.get("text", "").strip()
|
||||
chunks: List[Dict[str, Any]] = raw.get("chunks", [])
|
||||
|
||||
segments = []
|
||||
for i, chunk in enumerate(chunks):
|
||||
ts = chunk.get("timestamp") or (None, None)
|
||||
start = float(ts[0]) if ts[0] is not None else None
|
||||
end = float(ts[1]) if ts[1] is not None else None
|
||||
segments.append({
|
||||
"id": i,
|
||||
"start": round(start, 3) if start is not None else None,
|
||||
"end": round(end, 3) if end is not None else None,
|
||||
"text": chunk.get("text", "").strip(),
|
||||
})
|
||||
|
||||
# Estimate duration from last segment end
|
||||
duration = None
|
||||
if segments and segments[-1]["end"] is not None:
|
||||
duration = segments[-1]["end"]
|
||||
|
||||
return JSONResponse({
|
||||
"backend": "qwen3",
|
||||
"model": model,
|
||||
"language": lang_code,
|
||||
"duration": duration,
|
||||
"text": full_text,
|
||||
"segments": segments,
|
||||
"diarized": False,
|
||||
})
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
_log(f"error: {type(e).__name__}: {e}")
|
||||
raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")
|
||||
finally:
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", default="0.0.0.0")
|
||||
parser.add_argument("--port", type=int, default=8004)
|
||||
args = parser.parse_args()
|
||||
_log(f"starting host={args.host} port={args.port} device={DEVICE}")
|
||||
uvicorn.run(app, host=args.host, port=args.port)
|
||||
Reference in New Issue
Block a user