handler.py

import base64
import binascii
import os
import queue
import threading
import time
import warnings
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import torch
from transformers import pipeline
from transformers.pipelines.audio_utils import ffmpeg_read

SAMPLE_RATE = 16000
MAX_QUEUE_DEPTH = 32
MIN_AUDIO_SECONDS = 0.2
MAX_AUDIO_SECONDS = 35.0


@dataclass
class _QueuedRequest:
    audio: np.ndarray
    params: Dict[str, Any]
    event: threading.Event = field(default_factory=threading.Event)
    result: Optional[Dict[str, Any]] = None
    error: Optional[Exception] = None


class EndpointHandler:
    """
    Custom HF Inference Endpoint handler with micro-batching.

    Why this exists:
    - default ASR handling is effectively one-request-at-a-time on many setups
    - this handler coalesces near-simultaneous requests into one GPU forward pass
    - response shape is compatible with callers expecting {text, chunks}
    """

    def __init__(self, path: str = ""):
        model_path = path or "wasimlhr/whisper-quran-v1"
        use_cuda = torch.cuda.is_available()
        torch_dtype = torch.float16 if use_cuda else torch.float32
        device = 0 if use_cuda else -1

        self._pipe = pipeline(
            task="automatic-speech-recognition",
            model=model_path,
            device=device,
            torch_dtype=torch_dtype,
        )

        self._batch_window_ms = self._read_int("ASR_BATCH_WINDOW_MS", 35, 1, 200)
        self._max_batch_size = self._read_int("ASR_MAX_BATCH_SIZE", 4, 1, 16)
        self._request_timeout_s = float(os.getenv("ASR_REQUEST_TIMEOUT_S", "45"))
        self._max_queue_depth = self._read_int("ASR_MAX_QUEUE_DEPTH", MAX_QUEUE_DEPTH, 4, 256)
        self._max_audio_seconds = float(os.getenv("ASR_MAX_AUDIO_SECONDS", str(MAX_AUDIO_SECONDS)))
        self._min_audio_seconds = float(os.getenv("ASR_MIN_AUDIO_SECONDS", str(MIN_AUDIO_SECONDS)))

        self._queue: "queue.Queue[_QueuedRequest]" = queue.Queue()
        self._worker = threading.Thread(target=self._drain_loop, daemon=True)
        self._worker.start()

        print(
            f"[handler] initialized model={model_path} device={device} "
            f"batch_window_ms={self._batch_window_ms} max_batch_size={self._max_batch_size} "
            f"max_queue_depth={self._max_queue_depth}"
        )

    @staticmethod
    def _read_int(name: str, default: int, min_v: int, max_v: int) -> int:
        try:
            value = int(os.getenv(name, str(default)))
        except ValueError:
            value = default
        return max(min_v, min(max_v, value))

    def __call__(self, data: Any) -> Dict[str, Any]:
        payload, raw_params = self._extract_payload_and_params(data)
        audio = self._decode_audio(payload)
        params = self._normalize_params(raw_params)

        if self._queue.qsize() >= self._max_queue_depth:
            raise RuntimeError("ASR endpoint busy (queue saturated), please retry shortly")

        req = _QueuedRequest(audio=audio, params=params)
        self._queue.put(req)

        if not req.event.wait(timeout=self._request_timeout_s):
            raise TimeoutError("ASR request timed out while waiting in handler queue")
        if req.error is not None:
            raise RuntimeError(f"ASR request failed: {req.error}")
        return req.result or {"text": "", "chunks": []}

    def _drain_loop(self) -> None:
        while True:
            first = self._queue.get()
            batch = [first]

            deadline = time.perf_counter() + (self._batch_window_ms / 1000.0)
            while len(batch) < self._max_batch_size:
                timeout = deadline - time.perf_counter()
                if timeout <= 0:
                    break
                try:
                    batch.append(self._queue.get(timeout=timeout))
                except queue.Empty:
                    break

            self._process_batch(batch)

    def _process_batch(self, batch: List[_QueuedRequest]) -> None:
        groups: Dict[Tuple[Any, ...], List[_QueuedRequest]] = {}
        for req in batch:
            groups.setdefault(self._group_key(req.params), []).append(req)

        for group in groups.values():
            params = group[0].params
            pipeline_inputs = [{"array": r.audio, "sampling_rate": SAMPLE_RATE} for r in group]

            try:
                # Some transformers versions still emit internal deprecation warnings
                # while this pipeline API migrates. We explicitly suppress that known
                # warning noise in endpoint logs and rely on supported kwargs.
                with warnings.catch_warnings():
                    warnings.filterwarnings(
                        "ignore",
                        message="The input name `inputs` is deprecated.*",
                        category=FutureWarning,
                    )
                    outputs = self._pipe(
                        pipeline_inputs,
                        return_timestamps=params["return_timestamps"],
                        batch_size=len(group),
                        chunk_length_s=params["chunk_length_s"],
                        generate_kwargs={
                            "language": params["language"],
                            "task": params["task"],
                            "temperature": params["temperature"],
                        },
                    )
                if isinstance(outputs, dict):
                    outputs = [outputs]
                for req, out in zip(group, outputs):
                    req.result = self._format_output(out)
                    req.event.set()
            except Exception as exc:  # noqa: BLE001
                for req in group:
                    req.error = exc
                    req.event.set()

    @staticmethod
    def _group_key(params: Dict[str, Any]) -> Tuple[Any, ...]:
        return (
            params["language"],
            params["task"],
            params["return_timestamps"],
            params["chunk_length_s"],
            params["temperature"],
        )

    @staticmethod
    def _format_output(output: Dict[str, Any]) -> Dict[str, Any]:
        text = str(output.get("text", "")).strip()
        chunks_out: List[Dict[str, Any]] = []
        for ch in output.get("chunks", []) or []:
            if not isinstance(ch, dict):
                continue
            ctext = str(ch.get("text", "")).strip()
            ts = ch.get("timestamp")
            if isinstance(ts, (list, tuple)) and len(ts) >= 2:
                start, end = ts[0], ts[1]
            else:
                start, end = None, None
            chunks_out.append({"text": ctext, "start": start, "end": end, "timestamp": [start, end]})
        return {"text": text, "chunks": chunks_out}

    @staticmethod
    def _normalize_params(params: Dict[str, Any]) -> Dict[str, Any]:
        params = params or {}
        return {
            "language": str(params.get("language", "ar")),
            "task": str(params.get("task", "transcribe")),
            "return_timestamps": bool(params.get("return_timestamps", True)),
            "chunk_length_s": float(params.get("chunk_length_s", 20)),
            "temperature": float(params.get("temperature", 0.0)),
        }

    @staticmethod
    def _extract_payload_and_params(data: Any) -> Tuple[Any, Dict[str, Any]]:
        if isinstance(data, (bytes, bytearray)):
            return bytes(data), {}
        if isinstance(data, dict):
            params = data.get("parameters", {}) or {}
            if "inputs" in data:
                return data["inputs"], params
            if "audio" in data:
                return data["audio"], params
            raise ValueError("Expected 'inputs' or 'audio' in request body")
        raise TypeError(f"Unsupported request type: {type(data)}")

    def _decode_audio(self, payload: Any) -> np.ndarray:
        if isinstance(payload, dict) and "array" in payload:
            arr = np.asarray(payload["array"], dtype=np.float32)
            if arr.ndim > 1:
                arr = np.squeeze(arr)
            if arr.ndim != 1 or arr.size == 0:
                raise ValueError("Audio array payload must be a non-empty 1D array")
            duration_s = arr.size / float(SAMPLE_RATE)
            if duration_s < self._min_audio_seconds:
                raise ValueError(
                    f"Audio too short ({duration_s:.3f}s), need >= {self._min_audio_seconds:.1f}s"
                )
            if duration_s > self._max_audio_seconds:
                max_samples = int(self._max_audio_seconds * SAMPLE_RATE)
                arr = arr[:max_samples]
            return arr

        if isinstance(payload, str):
            if payload.startswith("data:") and "," in payload:
                payload = payload.split(",", 1)[1]
            try:
                audio_bytes = base64.b64decode(payload, validate=True)
            except (binascii.Error, ValueError) as exc:
                raise ValueError("String payload must be base64-encoded audio bytes") from exc
        elif isinstance(payload, (bytes, bytearray)):
            audio_bytes = bytes(payload)
        else:
            raise TypeError(f"Unsupported audio payload type: {type(payload)}")

        if not audio_bytes:
            raise ValueError("Audio payload is empty")

        try:
            audio = ffmpeg_read(audio_bytes, SAMPLE_RATE)
        except Exception as exc:  # noqa: BLE001
            raise ValueError(
                "Audio decode failed: malformed or unsupported audio payload"
            ) from exc

        arr = np.asarray(audio, dtype=np.float32)
        if arr.ndim > 1:
            arr = np.squeeze(arr)
        if arr.ndim != 1 or arr.size == 0:
            raise ValueError("Decoded audio is empty or invalid")

        duration_s = arr.size / float(SAMPLE_RATE)
        if duration_s < self._min_audio_seconds:
            raise ValueError(
                f"Audio too short ({duration_s:.3f}s), need >= {self._min_audio_seconds:.1f}s"
            )
        if duration_s > self._max_audio_seconds:
            max_samples = int(self._max_audio_seconds * SAMPLE_RATE)
            arr = arr[:max_samples]

        return arr