speech/app/asr/workers/qwen3_worker.py

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)