encoach_be_odoo19/encoach_ai_media/services/whisper_service.py

import logging
from concurrent.futures import ThreadPoolExecutor

import numpy as np
from tenacity import retry, stop_after_attempt, wait_exponential

from odoo.addons.encoach_ai.models.constants import GPT_MODELS, TEMPERATURE
from odoo.addons.encoach_ai.services.openai_service import EncoachOpenAIService

_logger = logging.getLogger(__name__)

SAMPLE_RATE = 16000
CHUNK_SAMPLES = 480000  # 30 seconds at 16 kHz
OVERLAP_RATIO = 0.25
WHISPER_OPTIONS = {
    "fp16": False,
    "language": "English",
    "verbose": False,
}

OVERLAP_CLEANUP_PROMPT = (
    "The following are two transcribed segments from the same audio. "
    "They overlap at the boundary. Remove any duplicated words at the "
    "junction and return the cleaned, combined text as JSON: "
    '{"text": "cleaned text"}'
)

DIALOG_DETECTION_PROMPT = (
    "You are a helpful assistant designed to output JSON on either one of "
    "these formats:\n"
    '1 - {"dialog": [{"name": "name", "gender": "gender", "text": "text"}]}\n'
    '2 - {"dialog": "text"}\n\n'
    "A transcription of an audio file will be provided to you. Based on that "
    "transcription you will need to determine whether the transcription is a "
    "conversation or a monologue. If it is a conversation, output format 1. "
    "If it is a monologue, output format 2."
)

_model_pool = []
_pool_executor = None
_pool_lock = None


def _get_pool(num_workers=4):
    global _model_pool, _pool_executor, _pool_lock
    import threading
    if _pool_lock is None:
        _pool_lock = threading.Lock()

    with _pool_lock:
        if _pool_executor is None:
            import whisper
            _logger.info("Loading %d Whisper model instances...", num_workers)
            _model_pool = [whisper.load_model("base") for _ in range(num_workers)]
            _pool_executor = ThreadPoolExecutor(max_workers=num_workers)
            _logger.info("Whisper pool ready with %d workers", num_workers)
    return _pool_executor, _model_pool


class EncoachWhisperService:

    def __init__(self, env):
        self.env = env
        self.ai = EncoachOpenAIService(env)
        self._num_workers = int(
            env["ir.config_parameter"].sudo().get_param("encoach.whisper_workers", "4")
        )
        self._worker_idx = 0
        import threading
        self._idx_lock = threading.Lock()

    def _get_model(self):
        executor, pool = _get_pool(self._num_workers)
        with self._idx_lock:
            idx = self._worker_idx % len(pool)
            self._worker_idx += 1
        return pool[idx]

    def transcribe(self, audio_data):
        if not isinstance(audio_data, np.ndarray):
            audio_data = np.frombuffer(audio_data, dtype=np.float32)

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

        chunks = self._chunk_audio(audio_data)

        if len(chunks) == 1:
            return self._transcribe_chunk(chunks[0])

        executor, pool = _get_pool(self._num_workers)
        futures = []
        for i, chunk in enumerate(chunks):
            model = pool[i % len(pool)]
            future = executor.submit(self._transcribe_chunk_with_model, chunk, model)
            futures.append(future)

        segments = [f.result() for f in futures]
        return self._merge_segments(segments)

    @retry(stop=stop_after_attempt(3), wait=wait_exponential(min=1, max=10))
    def _transcribe_chunk(self, chunk):
        model = self._get_model()
        result = model.transcribe(chunk, **WHISPER_OPTIONS)
        return result.get("text", "").strip()

    @staticmethod
    @retry(stop=stop_after_attempt(3), wait=wait_exponential(min=1, max=10))
    def _transcribe_chunk_with_model(chunk, model):
        result = model.transcribe(chunk, **WHISPER_OPTIONS)
        return result.get("text", "").strip()

    def _merge_segments(self, segments):
        if len(segments) <= 1:
            return segments[0] if segments else ""

        merged = segments[0]
        for i in range(1, len(segments)):
            result = self.ai.prediction(
                model=GPT_MODELS["grading"],
                messages=[
                    {"role": "system", "content": OVERLAP_CLEANUP_PROMPT},
                    {
                        "role": "user",
                        "content": (
                            f'Segment A (end): "...{merged[-200:]}"\n'
                            f'Segment B (start): "{segments[i][:200]}..."'
                        ),
                    },
                ],
                temperature=TEMPERATURE["grading"],
                check_blacklisted=False,
            )
            if result and "text" in result:
                merged = result["text"]
            else:
                merged = f"{merged} {segments[i]}"
        return merged

    def detect_dialog(self, transcript):
        return self.ai.prediction(
            model=GPT_MODELS["grading"],
            messages=[
                {"role": "system", "content": DIALOG_DETECTION_PROMPT},
                {"role": "user", "content": transcript},
            ],
            temperature=TEMPERATURE["grading"],
            check_blacklisted=False,
        )

    @staticmethod
    def _chunk_audio(audio, chunk_size=CHUNK_SAMPLES, overlap=OVERLAP_RATIO):
        total = len(audio)
        if total <= chunk_size:
            return [audio]

        step = int(chunk_size * (1 - overlap))
        chunks = []
        start = 0
        while start < total:
            end = min(start + chunk_size, total)
            chunks.append(audio[start:end])
            if end >= total:
                break
            start += step
        return chunks