AImpower/StutteredSpeechASR

1"""2Custom Inference Handler for StutteredSpeechASR Model3Handles audio input and returns transcriptions for stuttered speech.4"""5 6import torch7import librosa8import numpy as np9import base6410import io11import logging12from copy import deepcopy13from typing import Dict, Any14from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor15 16# Configure logging17logging.basicConfig(level=logging.INFO)18logger = logging.getLogger(__name__)19 20 21class EndpointHandler:22    """23    Custom handler for StutteredSpeechASR inference endpoint.24    25    This handler processes audio inputs and returns transcriptions26    using the fine-tuned Whisper model for stuttered Mandarin speech.27    """28    29    def __init__(self, path: str = ""):30        """31        Initialize the handler by loading the model and processor.32        33        Args:34            path: Path to the model directory (provided by Inference Endpoints)35        """36        logger.info("Initializing StutteredSpeechASR handler...")37        38        # Determine device and dtype39        self.device = "cuda" if torch.cuda.is_available() else "cpu"40        self.torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float3241        42        logger.info(f"Using device: {self.device}")43        logger.info(f"Using dtype: {self.torch_dtype}")44        45        # Load model and processor46        try:47            self.model = AutoModelForSpeechSeq2Seq.from_pretrained(48                path,49                torch_dtype=self.torch_dtype50            )51            self.processor = AutoProcessor.from_pretrained(path)52            self.model.to(self.device)53            self.model.eval()  # Set to evaluation mode54            55            logger.info("Model and processor loaded successfully!")56        except Exception as e:57            logger.error(f"Error loading model: {e}")58            raise59    60    def _load_audio_from_bytes(self, audio_bytes: bytes) -> np.ndarray:61        """62        Load audio from bytes and resample to 16kHz.63        64        Args:65            audio_bytes: Raw audio bytes66            67        Returns:68            Audio waveform as numpy array69        """70        try:71            # Load audio from bytes using librosa72            audio_buffer = io.BytesIO(audio_bytes)73            waveform, _ = librosa.load(audio_buffer, sr=16000, mono=True)74            return waveform75        except Exception as e:76            logger.error(f"Error loading audio from bytes: {e}")77            raise ValueError(f"Failed to load audio: {e}")78    79    def _load_audio_from_base64(self, base64_string: str) -> np.ndarray:80        """81        Load audio from base64-encoded string.82        83        Args:84            base64_string: Base64-encoded audio data85            86        Returns:87            Audio waveform as numpy array88        """89        try:90            # Decode base64 string91            audio_bytes = base64.b64decode(base64_string)92            return self._load_audio_from_bytes(audio_bytes)93        except Exception as e:94            logger.error(f"Error decoding base64 audio: {e}")95            raise ValueError(f"Failed to decode base64 audio: {e}")96    97    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:98        """99        Process incoming requests and return transcriptions.100        101        Expected input formats:102        1. {"inputs": "base64_encoded_audio_string"}103        2. {"inputs": {"audio": "base64_encoded_audio_string"}}104        3. Binary audio data in request body105 106        Dictionary requests may include {"parameters": {"language": "en"}}.107        Omit language (or use null) to keep the model's default behavior.108        109        Args:110            data: Input data dictionary111            112        Returns:113            Dictionary containing transcription and metadata114        """115        try:116            logger.info("Processing inference request...")117 118            generate_kwargs = {}119            if isinstance(data, dict):120                parameters = data.get("parameters", {})121                if not isinstance(parameters, dict):122                    raise ValueError("'parameters' must be a dictionary")123                language = parameters.get("language")124                if language is not None:125                    if not isinstance(language, str) or not language.strip():126                        raise ValueError("'language' must be a non-empty string or null")127                    # Keep language overrides local to this request, including on128                    # Transformers versions that modify the generation config.129                    generate_kwargs = {130                        "generation_config": deepcopy(self.model.generation_config),131                        "language": language.strip(),132                        "task": "transcribe",133                    }134            135            # Extract audio data from various input formats136            waveform = None137            138            if isinstance(data, dict):139                # Format 1: {"inputs": "base64_string"}140                if "inputs" in data:141                    inputs = data["inputs"]142                    143                    if isinstance(inputs, str):144                        # Base64-encoded audio145                        waveform = self._load_audio_from_base64(inputs)146                    147                    elif isinstance(inputs, dict):148                        # Format 2: {"inputs": {"audio": "base64_string"}}149                        if "audio" in inputs:150                            waveform = self._load_audio_from_base64(inputs["audio"])151                        else:152                            raise ValueError("Missing 'audio' field in inputs dictionary")153                    154                    elif isinstance(inputs, bytes):155                        # Binary audio data156                        waveform = self._load_audio_from_bytes(inputs)157                    158                    else:159                        raise ValueError(f"Unsupported input type: {type(inputs)}")160                161                # Direct audio field162                elif "audio" in data:163                    audio_data = data["audio"]164                    if isinstance(audio_data, str):165                        waveform = self._load_audio_from_base64(audio_data)166                    elif isinstance(audio_data, bytes):167                        waveform = self._load_audio_from_bytes(audio_data)168                169                else:170                    raise ValueError("No valid audio data found in request. Expected 'inputs' or 'audio' field.")171            172            elif isinstance(data, (bytes, bytearray)):173                # Format 3: Direct binary data174                waveform = self._load_audio_from_bytes(bytes(data))175            176            else:177                raise ValueError(f"Unsupported data type: {type(data)}")178            179            if waveform is None:180                raise ValueError("Failed to extract audio from request")181            182            logger.info(f"Audio loaded: {len(waveform)} samples at 16kHz")183            184            # Process audio with the processor185            input_features = self.processor(186                waveform,187                sampling_rate=16000,188                return_tensors="pt"189            ).input_features190            191            # Move to device192            input_features = input_features.to(self.device, dtype=self.torch_dtype)193            194            # Use the requested language, or the model's defaults when omitted.195            with torch.no_grad():196                predicted_ids = self.model.generate(input_features, **generate_kwargs)197            198            # Decode transcription199            transcription = self.processor.batch_decode(200                predicted_ids,201                skip_special_tokens=True202            )[0]203            204            logger.info(f"Transcription complete: {transcription[:100]}...")205            206            # Return result207            return {208                "transcription": transcription.strip(),209                "audio_duration_seconds": float(len(waveform) / 16000),210                "model": "AImpower/StutteredSpeechASR"211            }212        213        except Exception as e:214            logger.error(f"Error during inference: {e}", exc_info=True)215            return {216                "error": str(e),217                "transcription": None218            }219