add model name check

AaltoRSE · Feb 2, 2024 · c3cd8b8 · c3cd8b8
1 parent d19cd7e
commit c3cd8b8
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/speech2text.py b/src/speech2text.py
@@ -15,7 +15,6 @@
 import whisperx
 from numba.core.errors import (NumbaDeprecationWarning,
                                NumbaPendingDeprecationWarning)
-from pydub import AudioSegment
 from whisperx.types import TranscriptionResult
 
 import settings
@@ -229,6 +228,10 @@ def load_whisperx_model(
     if device is None:
         device = "cuda" if torch.cuda.is_available() else "cpu"
 
+    if name not in settings.available_whisper_models:
+        logger.warning(f"Specified model '{name}' not among available models: {settings.available_whisper_models}. Opting to use the default model '{settings.default_whisper_model}' instead")
+        name = settings.default_whisper_model
+
     compute_type = "float16" if device == "cuda" else "int8"
     try:
         model = whisperx.load_model(