faster-whisper STT 엔진 래퍼 - 음성→텍스트 변환

"""
faster-whisper STT 엔진 래퍼 - 음성→텍스트 변환

@source voice-to-text-v2
@extracted 2026-02-15
@updated 2026-03-14 — initial_prompt, word_timestamps, 자동 언어 감지 추가
@version 1.1.0

의존성:
    - faster-whisper (pip install faster-whisper)
    - numpy
    - (선택) torch - CUDA GPU 자동 감지용

사용법:
    from ai.whisper_engine import WhisperSTTEngine, get_engine

    # 초기화 (싱글톤)
    engine = get_engine()

    # 모델 로드
    engine.load_model(model_size="base", device="auto")

    # 변환
    result = engine.transcribe(audio_float32_array, language="ko")
    print(result["text"])
    print(f"신뢰도: {result['confidence']:.2f}")

    # 청크 배치 변환
    results = engine.transcribe_streaming([chunk1, chunk2, chunk3])

지원 모델:
    tiny (~75MB), base (~142MB), small (~466MB),
    medium (~1.5GB), large-v3 (~3GB)
"""

import logging
import math
import time
from pathlib import Path

import numpy as np

logger = logging.getLogger(__name__)

# faster-whisper 선택적 import
try:
    from faster_whisper import WhisperModel
except ImportError:
    WhisperModel = None

# 모델 카탈로그
WHISPER_MODELS = {
    "tiny": {"name": "Tiny", "description": "가장 빠름, 정확도 낮음 (~75MB)"},
    "base": {"name": "Base", "description": "속도와 정확도 균형 (~142MB)"},
    "small": {"name": "Small", "description": "높은 정확도 (~466MB)"},
    "medium": {"name": "Medium", "description": "매우 높은 정확도 (~1.5GB)"},
    "large-v3": {"name": "Large V3", "description": "최고 정확도, 느림 (~3GB)"},
}

DEFAULT_SAMPLE_RATE = 16000


def detect_device() -> tuple[str, str]:
    """
    최적 디바이스 및 컴퓨트 타입 자동 감지.

    Returns:
        (device, compute_type) 튜플
        - CUDA GPU: ("cuda", "float16")
        - CPU: ("cpu", "int8")
    """
    try:
        import torch
        if torch.cuda.is_available():
            gpu_name = torch.cuda.get_device_name(0)
            vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
            logger.info(f"CUDA GPU 감지: {gpu_name} ({vram_gb:.1f}GB VRAM)")
            return ("cuda", "float16")
    except ImportError:
        pass

    try:
        import ctranslate2
        cuda_types = ctranslate2.get_supported_compute_types("cuda")
        if cuda_types:
            logger.info(f"CTranslate2 CUDA 지원: {cuda_types}")
            compute = "float16" if "float16" in cuda_types else "int8"
            return ("cuda", compute)
    except Exception:
        pass

    logger.info("GPU 미감지, CPU 사용")
    return ("cpu", "int8")


class WhisperSTTEngine:
    """
    faster-whisper 기반 음성→텍스트 엔진.

    특징:
        - CUDA GPU 자동 감지 + CPU 폴백
        - 모델 캐싱 (중복 로드 방지)
        - VAD (Voice Activity Detection) 필터링
        - 신뢰도(confidence) 계산 (log probability → 0~1)
        - CUDA 런타임 에러 시 자동 CPU 폴백
    """

    def __init__(self, models_dir: str | Path | None = None):
        """
        Args:
            models_dir: 모델 다운로드/캐시 디렉토리 (None이면 기본 경로)
        """
        self._model = None
        self._model_size: str | None = None
        self._device: str | None = None
        self._is_loading: bool = False
        self._load_progress: float = 0.0
        self._models_dir = str(models_dir) if models_dir else None

    @property
    def is_loaded(self) -> bool:
        return self._model is not None

    @property
    def current_model(self) -> str | None:
        return self._model_size

    @property
    def is_loading(self) -> bool:
        return self._is_loading

    def load_model(
        self,
        model_size: str = "base",
        device: str = "auto",
        compute_type: str | None = None,
    ) -> bool:
        """
        Whisper 모델 로드.

        Args:
            model_size: 모델 크기 ("tiny", "base", "small", "medium", "large-v3")
            device: 디바이스 ("auto", "cpu", "cuda")
            compute_type: 컴퓨트 타입 ("int8", "float16", "float32", None=자동)

        Returns:
            로드 성공 여부
        """
        if WhisperModel is None:
            logger.error("faster-whisper 미설치. 설치: pip install faster-whisper")
            return False

        if self._model is not None and self._model_size == model_size:
            logger.info(f"모델 '{model_size}' 이미 로드됨")
            return True

        # 디바이스 자동 감지
        if device == "auto":
            device, auto_compute = detect_device()
            if compute_type is None:
                compute_type = auto_compute
        elif compute_type is None:
            compute_type = "float16" if device == "cuda" else "int8"

        self._is_loading = True
        self._load_progress = 0.0

        try:
            logger.info(f"Whisper 모델 로드: {model_size} (device={device}, compute={compute_type})")
            self._load_progress = 10.0

            kwargs = {
                "device": device,
                "compute_type": compute_type,
            }
            if self._models_dir:
                kwargs["download_root"] = self._models_dir

            self._model = WhisperModel(model_size, **kwargs)
            self._model_size = model_size
            self._device = device
            self._load_progress = 100.0
            logger.info(f"모델 '{model_size}' 로드 완료 (device={device})")
            return True

        except Exception as e:
            logger.error(f"모델 '{model_size}' 로드 실패 ({device}): {e}")

            # CUDA 실패 시 CPU 폴백
            if device == "cuda":
                logger.warning("CUDA 실패, CPU(int8) 폴백 시도...")
                try:
                    kwargs_fallback = {"device": "cpu", "compute_type": "int8"}
                    if self._models_dir:
                        kwargs_fallback["download_root"] = self._models_dir

                    self._model = WhisperModel(model_size, **kwargs_fallback)
                    self._model_size = model_size
                    self._device = "cpu"
                    self._load_progress = 100.0
                    logger.info(f"모델 '{model_size}' CPU 폴백 로드 완료")
                    return True
                except Exception as e2:
                    logger.error(f"CPU 폴백도 실패: {e2}")

            self._model = None
            self._model_size = None
            return False

        finally:
            self._is_loading = False

    def unload_model(self):
        """모델 언로드 (메모리 해제)"""
        self._model = None
        self._model_size = None
        self._load_progress = 0.0
        logger.info("모델 언로드됨")

    def transcribe(
        self,
        audio_data: np.ndarray,
        language: str = "ko",
        beam_size: int = 5,
        vad_filter: bool = True,
        sample_rate: int = DEFAULT_SAMPLE_RATE,
        initial_prompt: str | None = None,
        word_timestamps: bool = False,
    ) -> dict:
        """
        오디오를 텍스트로 변환.

        Args:
            audio_data: float32 numpy 배열 (모노, 16kHz 권장)
            language: 언어 코드 ("ko", "en", "ja", "zh", "auto" 등). "auto" 또는 None이면 자동 감지.
            beam_size: 빔 서치 크기 (클수록 정확, 느림)
            vad_filter: VAD(Voice Activity Detection) 필터 활성화
            sample_rate: 오디오 샘플레이트
            initial_prompt: 컨텍스트 힌트 (도메인 용어, 이전 텍스트 등)
            word_timestamps: 단어별 타임스탬프 활성화

        Returns:
            dict: text, language, detected_language, language_probability,
                  duration_seconds, confidence, word_count, segments, audio_duration

        Raises:
            RuntimeError: 모델 미로드 시
        """
        if not self.is_loaded:
            raise RuntimeError("모델 미로드. load_model()을 먼저 호출하세요.")

        start_time = time.time()

        try:
            # float32 보장 + 정규화
            if audio_data.dtype != np.float32:
                audio_data = audio_data.astype(np.float32)

            max_val = np.abs(audio_data).max()
            if max_val > 1.0:
                audio_data = audio_data / max_val

            # 언어 자동 감지: "auto" 또는 None이면 faster-whisper가 자동 감지
            effective_language = None if language in (None, "auto") else language

            # 짧은 오디오에서 자동 감지 경고
            audio_duration_sec = len(audio_data) / sample_rate
            if effective_language is None and audio_duration_sec < 5.0:
                logger.warning(
                    f"자동 언어 감지: 오디오 길이 {audio_duration_sec:.1f}초 (5초 미만 시 정확도 낮음)"
                )

            # 변환 파라미터 구성
            transcribe_kwargs = dict(
                language=effective_language,
                beam_size=beam_size,
                vad_filter=vad_filter,
                vad_parameters=dict(
                    min_silence_duration_ms=500,
                    speech_pad_ms=400,
                ),
                word_timestamps=word_timestamps,
            )
            if initial_prompt:
                transcribe_kwargs["initial_prompt"] = initial_prompt

            # 변환 실행
            segments, info = self._model.transcribe(
                audio_data,
                **transcribe_kwargs,
            )

            # 결과 수집
            all_text = []
            all_segments = []
            total_confidence = 0.0
            segment_count = 0

            for segment in segments:
                all_text.append(segment.text.strip())
                seg_data = {
                    "start": segment.start,
                    "end": segment.end,
                    "text": segment.text.strip(),
                    "avg_logprob": segment.avg_logprob,
                }
                # 단어별 타임스탬프 포함
                if word_timestamps and hasattr(segment, 'words') and segment.words:
                    seg_data["words"] = [
                        {"word": w.word, "start": w.start, "end": w.end, "probability": w.probability}
                        for w in segment.words
                    ]
                all_segments.append(seg_data)
                total_confidence += segment.avg_logprob
                segment_count += 1

            text = " ".join(all_text).strip()
            duration = time.time() - start_time
            avg_confidence = (total_confidence / segment_count) if segment_count > 0 else 0.0

            # log probability → 0~1 신뢰도 변환
            # logprob  0.0 → ~1.0 / -0.5 → ~0.61 / -1.0 → ~0.37 / -2.0 → ~0.14
            confidence = max(0.0, min(1.0, math.exp(avg_confidence)))

            detected_lang = info.language if info else language
            result = {
                "text": text,
                "language": detected_lang,
                "detected_language": detected_lang,
                "language_probability": round(info.language_probability, 3) if info and hasattr(info, 'language_probability') else None,
                "duration_seconds": round(duration, 3),
                "confidence": round(confidence, 3),
                "word_count": len(text.split()) if text else 0,
                "segments": all_segments,
                "audio_duration": round(len(audio_data) / sample_rate, 2),
            }

            logger.info(
                f"변환 완료: {result['audio_duration']}s 오디오 → "
                f"{result['word_count']}단어 ({duration:.2f}s, 신뢰도: {confidence:.2f})"
            )

            return result

        except Exception as e:
            # CUDA 런타임 에러 시 CPU 폴백 후 재시도
            error_str = str(e).lower()
            if any(kw in error_str for kw in ("cuda", "cublas", "cudnn")):
                logger.warning(f"CUDA 런타임 에러: {e}. CPU로 재로드...")
                model_size = self._model_size
                self._model = None
                self._model_size = None
                if self.load_model(model_size=model_size, device="cpu", compute_type="int8"):
                    logger.info("CPU에서 재시도...")
                    return self.transcribe(
                        audio_data, language=language, beam_size=beam_size,
                        vad_filter=vad_filter, initial_prompt=initial_prompt,
                        word_timestamps=word_timestamps,
                    )
            logger.error(f"변환 실패: {e}")
            raise

    def transcribe_streaming(
        self,
        audio_chunks: list[np.ndarray],
        language: str = "ko",
        min_chunk_seconds: float = 0.5,
        sample_rate: int = DEFAULT_SAMPLE_RATE,
    ) -> list[dict]:
        """
        오디오 청크 리스트를 순차적으로 변환.

        Args:
            audio_chunks: float32 오디오 청크 리스트
            language: 언어 코드
            min_chunk_seconds: 최소 청크 길이 (초, 이하는 스킵)
            sample_rate: 샘플레이트

        Returns:
            각 청크의 변환 결과 리스트 (빈 텍스트 결과는 제외)
        """
        if not self.is_loaded:
            raise RuntimeError("모델 미로드")

        results = []
        for chunk in audio_chunks:
            if len(chunk) < sample_rate * min_chunk_seconds:
                continue
            result = self.transcribe(chunk, language=language, beam_size=3, vad_filter=False)
            if result["text"]:
                results.append(result)

        return results


# --- 싱글톤 패턴 ---

_engine: WhisperSTTEngine | None = None


def get_engine(models_dir: str | Path | None = None) -> WhisperSTTEngine:
    """
    전역 STT 엔진 싱글톤 인스턴스 반환.

    Args:
        models_dir: 모델 캐시 디렉토리

    Returns:
        WhisperSTTEngine 인스턴스
    """
    global _engine
    if _engine is None:
        _engine = WhisperSTTEngine(models_dir=models_dir)
    return _engine