handler.py

import base64
import datetime
import io
import logging
import os
import numpy as np
import soundfile as sf
from faster_whisper import WhisperModel
from typing import Dict, Any, Union

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class EndpointHandler:
    def __init__(self, path=""):
        """
        Initialize the endpoint handler.

        Args:
            path: Path to the model directory. In Hugging Face Inference Endpoints,
                 this will be the directory containing the model files.
        """

        logger.info("Initializing EndpointHandler")

        if os.environ.get("LOG_DIAGNOSTICS") == "true":
            self.__log_diagnostics__()

        cudaIsAvailable = os.environ.get("CUDA_VISIBLE_DEVICES")
        if cudaIsAvailable:
            logger.info("CUDA is available")
        else:
            logger.info("CUDA is not available, using CPU")

        # Load the model with fallback to CPU if CUDA fails
        device = "cuda" if cudaIsAvailable else "cpu"
        try:
            # First attempt with CUDA if available
            if device == "cuda":
                logger.info("Attempting to load model with CUDA support")
                self.model = WhisperModel(
                    path or ".",
                    compute_type="float16",
                    device="cuda",
                )
                logger.info("Model loaded successfully with CUDA")
            else:
                # CPU fallback
                raise ValueError("CUDA not available, using CPU")
        except Exception as e:
            # Log the error and fall back to CPU
            logger.warning(f"Error loading model with CUDA: {e}")
            logger.info("Falling back to CPU model")
            try:
                self.model = WhisperModel(
                    path or ".",
                    compute_type="int8",
                    device="cpu",
                )
                logger.info("Model loaded successfully with CPU")
            except Exception as cpu_err:
                logger.error(f"Error loading CPU model: {cpu_err}")
                raise

        # Set default parameters
        self.sampling_rate = 16000

        logger.info("EndpointHandler initialized")

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Process a request.

        Args:
            data: Request data containing audio input and optional parameters.
                 Expected format:
                 - For batch processing: {"inputs": audio_data, ...parameters}
                 - For streaming: {"inputs": audio_chunk, "stream": true, "session_id": "unique_id", ...parameters}

        Returns:
            Transcription result or error message.
        """
        logger.info("Processing request")

        try:
            # Extract inputs and parameters
            if "inputs" not in data:
                return {"error": "No inputs provided"}

            inputs = data.pop("inputs")
            parameters = data.pop("parameters", {})

            # Process audio input
            audio, sampling_rate = self._process_audio_input(inputs)

            # Get transcription parameters
            language = data.pop("language", "da")  # Default to Danish
            beam_size = data.pop("beam_size", 5)

            return self.transcribe(
                audio,
                sampling_rate,
                language,
                beam_size,
                **parameters,
            )

        except Exception as e:
            logger.error(f"Error processing request: {e}")
            return {"error": str(e)}

    def _process_audio_input(self, inputs: Union[str, Dict[str, Any]]) -> tuple:
        """
        Process audio input in various formats.

        Args:
            inputs: Audio input as base64 string, URL, or numpy array.

        Returns:
            Tuple of (audio_array, sampling_rate).
        """
        # Handle different input formats
        if isinstance(inputs, str):
            logger.info("Received audio input as base64 encoded string")
            # Base64 encoded audio
            if inputs.startswith(("data:", "data%3A")):

                # Handle data URI
                if "base64," in inputs:
                    audio_b64 = inputs.split("base64,")[1]
                else:
                    audio_b64 = inputs

                # Decode base64
                audio_bytes = base64.b64decode(audio_b64)

                # Read audio data
                with io.BytesIO(audio_bytes) as audio_io:
                    audio, sampling_rate = sf.read(audio_io)

                return audio, sampling_rate

            # URL or file path (not implemented for security reasons)
            else:
                raise ValueError("URL or file path inputs are not supported")

        # Dictionary with audio data
        elif isinstance(inputs, dict) and "audio" in inputs:
            logger.info("Received audio input as dictionary")
            if isinstance(inputs["audio"], list):
                # Convert list to numpy array
                audio = np.array(inputs["audio"], dtype=np.float32)
            else:
                # Assume it's already a numpy array
                audio = inputs["audio"]

            # Get sampling rate
            sampling_rate = inputs.get("sampling_rate", self.sampling_rate)

            return audio, sampling_rate

        elif isinstance(inputs, bytes):
            logger.info("Received raw bytes input")
            # Read audio data from bytes
            with io.BytesIO(inputs) as audio_io:
                audio, sampling_rate = sf.read(audio_io)
            return audio, sampling_rate

        # Unsupported input format
        else:
            raise ValueError(f"Unsupported input format: {type(inputs)}")

    def transcribe(
        self,
        audio: np.ndarray,
        sampling_rate: int,
        language: str,
        beam_size: int,
        **kwargs,
    ) -> Dict[str, Any]:
        """
        Perform the transcription

        Args:
            audio: Audio data as numpy array.
            sampling_rate: Sampling rate of the audio.
            language: Language code.
            beam_size: Beam size for transcription.
            **kwargs: Additional parameters.

        Returns:
            Transcription result.
        """
        logger.info(f"Batch transcription: {len(audio)} samples, {sampling_rate} Hz")

        # Resample if needed
        if sampling_rate != self.sampling_rate:
            logger.warning(
                f"Sampling rate mismatch: {sampling_rate} Hz vs {self.sampling_rate} Hz"
            )

        logger.info(f"Parameters: {kwargs}")

        # Transcribe
        # see parameters here: <https://github.com/SYSTRAN/faster-whisper/blob/master/faster_whisper/transcribe.py>
        now = datetime.datetime.now()
        segments, info = self.model.transcribe(
            audio,
            language=language,
            beam_size=beam_size,
            **kwargs,
        )

        logger.info(f"Transcription info: {info}")

        # Format results
        result = {
            "text": "",  # Full text will be populated below
            "segments": [],
            "language": info.language,
            "language_probability": info.language_probability,
        }

        # Process segments
        # Segments is a generator and the actual transcription is done in the loop below,
        all_text = []
        for segment in segments:
            segment_info = {
                "id": segment.id,
                "text": segment.text,
                "start": segment.start,
                "end": segment.end,
                "temperature": segment.temperature,
                "avg_logprob": segment.avg_logprob,
                "compression_ratio": segment.compression_ratio,
                "no_speech_prob": segment.no_speech_prob,
            }
            if kwargs.get("word_timestamps", False):
                # Add word timestamps if requested
                segment_info["words"] = [
                    {
                        "word": word.word,
                        "start": word.start,
                        "end": word.end,
                    }
                    for word in segment.words
                ]

            all_text.append(segment.text)
            result["segments"].append(segment_info)

        elapsed_time = datetime.datetime.now() - now
        logger.info(f"Transcription time: {elapsed_time}")
        logger.info(f"Segments: {len(result['segments'])}")
        # Combine text from all segments
        result["text"] = " ".join(all_text)

        return result

    def __log_diagnostics__(self):
        """
        Log diagnostics information for debugging.
        This includes CUDA availability, library paths, installed packages,
        and environment variables.

        Very useful as the HF endpoint runtime is rather secretive about its environment.
        """
        logger.info("Logging environment diagnostics")

        logger.info("LD_LIBRARY_PATH:")
        if "LD_LIBRARY_PATH" in os.environ:
            logger.info(os.environ["LD_LIBRARY_PATH"])
        else:
            logger.info("LD_LIBRARY_PATH not set")

        # and the files
        logger.info("LD_LIBRARY_PATH files:")
        if "LD_LIBRARY_PATH" in os.environ:
            for ld_path in os.environ["LD_LIBRARY_PATH"].split(":"):
                if os.path.exists(ld_path):
                    logger.info(f"  {ld_path}:")
                    for file in os.listdir(ld_path):
                        logger.info(f"    {file}")
                else:
                    logger.info(f"  {ld_path} does not exist")

        logger.info(f"Installed Python packages:")
        import pkg_resources

        for package in pkg_resources.working_set:
            logger.info(f"  {package.key}=={package.version}, {package.location}")

        # dump environment variables
        logger.info("Environment variables:")
        for key, value in os.environ.items():
            logger.info(f"  {key}: {value}")

        logger.info("NVIDIA environment:")
        import subprocess

        try:
            result = subprocess.run(["nvidia-smi"], capture_output=True, text=True)
            logger.info(result.stdout)
        except Exception as e:
            logger.warning(f"Could not run nvidia-smi: {e}")