handler.py

"""HuggingFace Inference Endpoint custom handler.

Deploy this by creating a HF Endpoint from a repo that contains:
    handler.py        (this file)
    requirements.txt
    pipeline/         (the long_form package)
    inference/        (only needed if you re-import evaluate_ctc directly)

Endpoint URL convention:
    POST  https://<id>.endpoints.huggingface.cloud
    Headers:
        Authorization: Bearer <HF_TOKEN>
        Content-Type:  audio/wav  (or audio/mp3, audio/m4a, audio/flac)
    Body:
        raw audio bytes
    OR JSON body:
        {"inputs": "<base64-encoded-audio>", "parameters": {"language_hint": "amh"}}

Response: JSON matching the TranscribeResponse contract in MOBILE_APP_DESIGN_PROMPT.md
"""
from __future__ import annotations

import base64
import json
import logging
import os
import sys
import tempfile
from pathlib import Path
from typing import Any, Dict

HERE = Path(__file__).resolve().parent
if str(HERE) not in sys.path:
    sys.path.insert(0, str(HERE))

from pipeline.long_form import LongFormPipeline  # noqa: E402

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s endpoint.handler %(levelname)s | %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("endpoint.handler")

ASR_REPO = os.environ.get("ASR_REPO", "boazsew/Ethio-ASR-w2v-bert-2.0-uf")
DIAR_MODEL = os.environ.get("DIAR_MODEL", "pyannote/speaker-diarization-3.1")


class EndpointHandler:
    def __init__(self, path: str = ""):
        # `path` is the local checkout of the HF repo when HF builds the container.
        # We don't use it directly because the ASR model lives in a different repo;
        # the LongFormPipeline downloads from `ASR_REPO`.
        log.info(f"init: ASR={ASR_REPO}  diarizer={DIAR_MODEL}")
        token = (os.environ.get("HF_TOKEN") or os.environ.get("HF_API_KEY") or "").strip()
        if not token:
            log.warning("HF_TOKEN not set in env — pyannote diarizer download will fail")
        self.pipe = LongFormPipeline(
            model_dir=ASR_REPO,
            hf_token=token,
            diar_model=DIAR_MODEL,
        )
        log.info("init: pipeline ready")

    def _extract_audio_bytes(self, data: Dict[str, Any]) -> bytes:
        """HF Endpoints may deliver audio in several shapes:
        - data["inputs"] = bytes               (raw upload, Content-Type audio/*)
        - data["inputs"] = str (base64)         (JSON body)
        - data["inputs"] = {"data": "..."}      (legacy serializer)
        """
        inputs = data.get("inputs")
        if inputs is None:
            raise ValueError("missing 'inputs' field")
        if isinstance(inputs, bytes):
            return inputs
        if isinstance(inputs, str):
            try:
                return base64.b64decode(inputs)
            except Exception as e:
                raise ValueError(f"could not base64-decode inputs: {e}") from e
        if isinstance(inputs, dict) and "data" in inputs:
            return self._extract_audio_bytes({"inputs": inputs["data"]})
        raise ValueError(f"unsupported inputs type: {type(inputs).__name__}")

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        try:
            audio_bytes = self._extract_audio_bytes(data)
        except ValueError as e:
            return {"error": str(e)}

        if len(audio_bytes) < 100:
            return {"error": "audio too small"}
        if len(audio_bytes) > 50 * 1024 * 1024:
            return {"error": f"audio too large ({len(audio_bytes)} bytes, limit 50 MB)"}

        # librosa can decode bytes via a temp file (most reliable across formats)
        suffix = (data.get("parameters") or {}).get("format", "wav")
        if not suffix.startswith("."):
            suffix = "." + suffix
        fd, path = tempfile.mkstemp(suffix=suffix, prefix="hf_endpoint_")
        os.write(fd, audio_bytes)
        os.close(fd)
        try:
            result = self.pipe.transcribe(path)
            return result.to_dict()
        except Exception as e:
            log.exception("transcription failed")
            return {"error": f"transcription failed: {e!s}"}
        finally:
            try:
                os.unlink(path)
            except OSError:
                pass