Add custom endpoint handler with micro-batching

README.md

@@ -169,6 +169,31 @@ The 10-second chunking eliminates the long-ayah drift problem entirely — every
 
 ## Usage
 
+### Hugging Face Inference Endpoint (custom handler, recommended)
+
+This repo now includes a custom `handler.py` for HF Inference Endpoints that supports:
+
+- micro-batching concurrent requests on one GPU worker
+- low-latency queue window (`ASR_BATCH_WINDOW_MS`)
+- bounded batching (`ASR_MAX_BATCH_SIZE`)
+- output shape compatible with clients expecting `{ text, chunks }`
+
+#### Endpoint configuration
+
+Set your endpoint to use this repository revision and custom handler.
+
+Optional environment variables:
+
+| Variable | Default | Description |
+|---|---|---|
+| `ASR_BATCH_WINDOW_MS` | `35` | Queue window to coalesce near-simultaneous requests into one forward pass |
+| `ASR_MAX_BATCH_SIZE` | `4` | Max requests per micro-batch |
+| `ASR_REQUEST_TIMEOUT_S` | `45` | Per-request queue wait timeout |
+
+The handler accepts raw `audio/wav` bytes or JSON payloads with:
+- `inputs` (base64/string/bytes audio), and optional
+- `parameters` (`language`, `task`, `return_timestamps`, `chunk_length_s`, `temperature`)
+
 ### Basic Transcription
 
 ```python

handler.py

@@ -0,0 +1,202 @@
+import base64
+import binascii
+import os
+import queue
+import threading
+import time
+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
+
+
+@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._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}"
+        )
+
+    @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)
+
+        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
+            inputs = [{"array": r.audio, "sampling_rate": SAMPLE_RATE} for r in group]
+
+            try:
+                outputs = self._pipe(
+                    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)}")
+
+    @staticmethod
+    def _decode_audio(payload: Any) -> np.ndarray:
+        if isinstance(payload, dict) and "array" in payload:
+            arr = np.asarray(payload["array"], dtype=np.float32)
+            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)}")
+
+        audio = ffmpeg_read(audio_bytes, SAMPLE_RATE)
+        return np.asarray(audio, dtype=np.float32)