Create Handler.py

handler.py

@@ -0,0 +1,176 @@
+from typing import Dict, List, Any
+import torch
+from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor
+import logging
+import base64
+import tempfile
+import os
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        """
+        Initialize the handler with the Wolof Whisper model and fix the forced_decoder_ids issue
+        """
+        logger.info(f"Loading Wolof Whisper model from {path}")
+
+        try:
+            # Load the model and processor
+            self.model = AutoModelForSpeechSeq2Seq.from_pretrained(
+                path,
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                low_cpu_mem_usage=True,
+                use_safetensors=True
+            )
+            self.processor = AutoProcessor.from_pretrained(path)
+
+            # Fix the deprecated forced_decoder_ids parameter
+            if hasattr(self.model, 'generation_config'):
+                logger.info("Fixing deprecated forced_decoder_ids parameter for Wolof model...")
+                
+                # Remove deprecated parameters that cause 400 errors
+                self.model.generation_config.forced_decoder_ids = None
+                
+                # Clear suppress tokens that might cause issues
+                if hasattr(self.model.generation_config, 'suppress_tokens'):
+                    self.model.generation_config.suppress_tokens = []
+                
+                # Set correct parameters for Wolof transcription
+                self.model.generation_config.language = "wo"  # Wolof language code
+                self.model.generation_config.task = "transcribe"
+                
+                # Ensure we don't have conflicting parameters
+                if hasattr(self.model.generation_config, 'decoder_input_ids'):
+                    self.model.generation_config.decoder_input_ids = None
+                if hasattr(self.model.generation_config, 'input_ids'):
+                    self.model.generation_config.input_ids = None
+                    
+                logger.info("Successfully fixed model configuration for Wolof transcription")
+
+            # Create pipeline with fixed model
+            self.pipe = pipeline(
+                "automatic-speech-recognition",
+                model=self.model,
+                tokenizer=self.processor.tokenizer,
+                feature_extractor=self.processor.feature_extractor,
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                device=0 if torch.cuda.is_available() else -1
+            )
+            
+            logger.info("Wolof Whisper model loaded successfully with fixed configuration")
+            
+        except Exception as e:
+            logger.error(f"Error loading Wolof model: {e}")
+            raise e
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Process the audio input and return Wolof transcription
+        Args:
+            data: Input data containing audio (binary or base64)
+        Returns:
+            Transcription result in the expected format
+        """
+        try:
+            logger.info("Processing Wolof audio transcription request")
+            
+            # Get the audio input
+            inputs = data.get("inputs", data)
+            
+            # Handle different input types
+            if isinstance(inputs, str):
+                logger.info("Processing base64 encoded audio")
+                # Base64 encoded audio - decode and save to temp file
+                try:
+                    audio_bytes = base64.b64decode(inputs)
+                except Exception as e:
+                    logger.error(f"Failed to decode base64 audio: {e}")
+                    return [{"error": f"Invalid base64 audio data: {str(e)}"}]
+                    
+                # Save to temporary file for processing
+                with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as temp_file:
+                    temp_file.write(audio_bytes)
+                    temp_path = temp_file.name
+                    
+                try:
+                    result = self._transcribe_audio(temp_path)
+                finally:
+                    # Clean up temp file
+                    if os.path.exists(temp_path):
+                        os.unlink(temp_path)
+                        
+            elif isinstance(inputs, bytes):
+                logger.info("Processing binary audio data")
+                # Direct binary audio data
+                with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as temp_file:
+                    temp_file.write(inputs)
+                    temp_path = temp_file.name
+                    
+                try:
+                    result = self._transcribe_audio(temp_path)
+                finally:
+                    # Clean up temp file
+                    if os.path.exists(temp_path):
+                        os.unlink(temp_path)
+                        
+            else:
+                logger.info("Processing direct audio path/data")
+                # Direct audio path or numpy array
+                result = self._transcribe_audio(inputs)
+
+            logger.info(f"Wolof transcription completed successfully")
+            return [result] if not isinstance(result, list) else result
+
+        except Exception as e:
+            logger.error(f"Error during Wolof transcription: {e}")
+            return [{"error": f"Wolof transcription failed: {str(e)}"}]
+
+    def _transcribe_audio(self, audio_input):
+        """
+        Internal method to transcribe audio using the fixed pipeline
+        """
+        try:
+            # Use the pipeline with explicit parameters to avoid forced_decoder_ids
+            result = self.pipe(
+                audio_input,
+                generate_kwargs={
+                    "language": "wo",  # Wolof language code
+                    "task": "transcribe",
+                    # Explicitly avoid deprecated parameters
+                    "forced_decoder_ids": None,
+                    "suppress_tokens": [],
+                    # Use modern parameters
+                    "max_length": 448,
+                    "num_beams": 1,
+                    "do_sample": False,
+                }
+            )
+            
+            # Extract text from result
+            if isinstance(result, dict):
+                text = result.get("text", "")
+            elif isinstance(result, list) and len(result) > 0:
+                text = result[0].get("text", "") if isinstance(result[0], dict) else str(result[0])
+            else:
+                text = str(result)
+            
+            # Return in expected format
+            return {
+                "text": text.strip(),
+                "language": "wo",
+                "model": "Alwaly/whisper-medium-wolof"
+            }
+            
+        except Exception as e:
+            logger.error(f"Pipeline transcription error: {e}")
+            # If we get the forced_decoder_ids error, provide helpful message
+            if "forced_decoder_ids" in str(e):
+                raise Exception(
+                    "forced_decoder_ids parameter is deprecated. "
+                    "This handler.py file should fix this issue. "
+                    "Please redeploy the endpoint."
+                )
+            else:
+                raise e