diff --git a/faster_whisper/__init__.py b/faster_whisper/__init__.py index e2fe00d..9b56a39 100644 --- a/faster_whisper/__init__.py +++ b/faster_whisper/__init__.py @@ -1,9 +1,10 @@ from faster_whisper.audio import decode_audio from faster_whisper.transcribe import WhisperModel -from faster_whisper.utils import download_model, format_timestamp +from faster_whisper.utils import available_models, download_model, format_timestamp from faster_whisper.version import __version__ __all__ = [ + "available_models", "decode_audio", "WhisperModel", "download_model", diff --git a/faster_whisper/utils.py b/faster_whisper/utils.py index 5987aee..f020bc2 100644 --- a/faster_whisper/utils.py +++ b/faster_whisper/utils.py @@ -2,7 +2,7 @@ import logging import os import re -from typing import Optional +from typing import List, Optional import huggingface_hub import requests @@ -24,6 +24,11 @@ _MODELS = { } +def available_models() -> List[str]: + """Returns the names of available models.""" + return list(_MODELS.keys()) + + def get_assets_path(): """Returns the path to the assets directory.""" return os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets") diff --git a/tests/test_utils.py b/tests/test_utils.py index ee404bf..bb488fe 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,6 +1,12 @@ import os -from faster_whisper import download_model +from faster_whisper import available_models, download_model + + +def test_available_models(): + models = available_models() + assert isinstance(models, list) + assert "tiny" in models def test_download_model(tmpdir):