Update handler.py

handler.py

@@ -2,11 +2,12 @@ from typing import Dict
 from faster_whisper import WhisperModel
 
 class EndpointHandler:
-    def __init__(self, *args, **kwargs):
-        # Initialize WhisperModel; assume that "large-v2" model files are in /repository
-        self.model = WhisperModel("large-v2", device="cpu")
+    def __init__(self, model_dir=None):
+        # If a model_dir is provided, use it; otherwise, default to 'large-v2'.
+        model_size = "large-v2" if model_dir is None else model_dir
+        self.model = WhisperModel(model_size, device="cpu")
 
-    def __call__(self, data: Dict) -> Dict:
+    def __call__(self, data: Dict) -> Dict[str, str]:
         # Process the input data expected to be in 'inputs' key containing audio file bytes
         audio_bytes = data["inputs"]