better model selection cli

src/transcribble/cli.py

@@ -10,6 +10,18 @@ class IgnoreHFWarning(logging.Filter):
 logger = logging.getLogger("huggingface_hub.utils._http")
 logger.addFilter(IgnoreHFWarning())
 
+class WhisperModel(argparse.Action):
+    def __call__(self, parser, namespace, value, option_string=None):
+        import whisper
+        models = whisper.available_models()
+
+        if value == "?":
+            parser.exit(message="\n".join(models) + "\n")
+        if value not in models:
+            parser.error(f"invalid mode '{value}'\nchoose from:\n    {'\n    '.join(models)}")
+
+        setattr(namespace, self.dest, value)
+
 def main():
     parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument("input")
@@ -17,28 +29,17 @@ def main():
                         help="Output subtitle (SRT) file path")
     parser.add_argument("-s", "--speakers", type=int, default=2,
                         help="Number of speakers in the audio")
-    parser.add_argument("-m", "--model", default='turbo',
+    parser.add_argument("-m", "--model", default='turbo', action=WhisperModel,
                         help="Whisper model to use ('?' to see available options)")
     parser.add_argument("--gpu", action="store_true", help="run on GPU if available")
     args = parser.parse_args()
 
     import torch
-    import whisper
 
     from .transcribe import transcribe
     from .diarize import diarize
     from .output import write_srt
 
-    models = whisper.available_models()
-
-    if args.model == "?":
-        for model in models:
-            print(model)
-        return
-
-    if args.model not in models:
-        parser.error(f"'{args.model}' is not a valid whisper model")
-
     has_gpu = torch.cuda.is_available()
     if args.gpu and not has_gpu:
         print("GPU not available. Falling back to CPU")