captcha_solver/engines/whisper_engine.py

"""faster-whisper engine for audio captcha transcription.

Uses CTranslate2 under the hood. CPU-friendly, very small models
(tiny=75MB, base=150MB). For captcha audio (digits/phrases with light
noise) 'tiny' is sufficient and 10x faster than 'base'.
"""

from __future__ import annotations

import os
import tempfile
from pathlib import Path
from threading import Lock
from typing import Optional

from captcha_solver.engines.base import BaseEngine
from captcha_solver.config import get_settings


class WhisperEngine(BaseEngine):
    name = "whisper"

    def __init__(self) -> None:
        super().__init__()
        self._model = None
        self._lock = Lock()

    def _do_load(self) -> None:
        from faster_whisper import WhisperModel

        s = get_settings()
        cache = s.cache_dir / "whisper"
        cache.mkdir(parents=True, exist_ok=True)
        os.environ.setdefault("HF_HOME", str(s.cache_dir / "hf"))

        self._model = WhisperModel(
            s.whisper_model,
            device=s.whisper_device,
            compute_type=s.whisper_compute_type,
            download_root=str(cache),
        )

    def _do_unload(self) -> None:
        self._model = None

    def transcribe(
        self,
        audio_bytes: bytes,
        language: Optional[str] = "en",
        beam_size: int = 1,
    ) -> tuple[str, float]:
        """Transcribe audio bytes. Returns (text, confidence).

        Confidence is the average segment log-prob (mapped to 0-1).
        """
        if not self._loaded:
            self.load()
        assert self._model is not None

        with tempfile.NamedTemporaryFile(suffix=".audio", delete=False) as tmp:
            tmp.write(audio_bytes)
            tmp_path = tmp.name
        try:
            segments, info = self._model.transcribe(
                tmp_path,
                language=language,
                beam_size=beam_size,
                vad_filter=False,
            )
            texts: list[str] = []
            confs: list[float] = []
            for seg in segments:
                texts.append(seg.text.strip())
                if seg.avg_logprob is not None:
                    import math
                    p = math.exp(seg.avg_logprob)
                    confs.append(max(0.0, min(1.0, p)))
            text = " ".join(t for t in texts if t).strip()
            conf = (sum(confs) / len(confs)) if confs else 0.0
            return text, conf
        finally:
            try:
                Path(tmp_path).unlink(missing_ok=True)
            except Exception:
                pass