Update handler.py

handler.py

@@ -1,26 +1,33 @@
+from typing import Dict
 from faster_whisper import WhisperModel
-import torch
 
 class EndpointHandler:
     def __init__(self):
-        model_size = "large-v2"  # Update model size if different
-        device = "cpu"  # Use CPU for Azure deployment
-        self.model = WhisperModel(model_size, device=device)
+        # Initialize WhisperModel
+        self.model = WhisperModel("large-v2")
 
-    def __call__(self, data):
-        # Extract audio bytes from the request data
-        audio_bytes = data.get("inputs")
+    def __call__(self, data: Dict) -> Dict:
+        # Get the audio file bytes from the request data
+        audio_bytes = data["inputs"]
         
-        # Convert audio bytes to audio samples
-        # Note: Additional conversion might be needed depending on the format of the incoming audio bytes
-        
-        # Transcribe the audio using the Whisper model without writing to disk
+        # Perform transcription
+        results = []
         segments, info = self.model.transcribe(audio_bytes)
+        for segment in segments:
+            result = {
+                "start": segment.start,
+                "end": segment.end,
+                "text": segment.text
+            }
+            results.append(result)
         
-        # Combine the text from all segments
-        text = " ".join(segment.text for segment in segments)
-        
-        # Return the transcribed text and the detected language
-        return {"text": text, "language": info.language_code}
-
-# If applicable, write additional conversion code to get samples from bytes
+        # Return the transcribed text along with language data
+        language_code, language_prob = info.language, info.language_probability
+        response = {
+            "transcription": results,
+            "language": {
+                "code": language_code,
+                "probability": language_prob
+            }
+        }
+        return response