added tests for Faster-Whisper and HF models

-added tests for Faster-Whisper and Hugging Face models; cap `tokenizers` to 0.20.3 for tests with python versions <3.9 due to: huggingface/tokenizers#1691

-added test for PyTorch 2.6.0

Author: jianfch

.github/workflows/test.yml

@@ -16,34 +16,54 @@ jobs:
           - python-version: '3.8'
             pytorch-version: 1.10.1
             numpy-requirement: "'numpy<2'"
+            tokenizers-requirement: "'tokenizers<=0.20.3'"
           - python-version: '3.8'
             pytorch-version: 1.13.1
             numpy-requirement: "'numpy<2'"
+            tokenizers-requirement: "'tokenizers<=0.20.3'"
           - python-version: '3.8'
             pytorch-version: 2.0.1
             numpy-requirement: "'numpy<2'"
+            tokenizers-requirement: "'tokenizers<=0.20.3'"
           - python-version: '3.9'
             pytorch-version: 2.1.2
             numpy-requirement: "'numpy<2'"
+            tokenizers-requirement: "'tokenizers'"
           - python-version: '3.10'
             pytorch-version: 2.2.2
             numpy-requirement: "'numpy<2'"
+            tokenizers-requirement: "'tokenizers'"
           - python-version: '3.11'
             pytorch-version: 2.3.1
             numpy-requirement: "'numpy'"
+            tokenizers-requirement: "'tokenizers'"
           - python-version: '3.12'
             pytorch-version: 2.4.1
             numpy-requirement: "'numpy'"
+            tokenizers-requirement: "'tokenizers'"
           - python-version: '3.12'
             pytorch-version: 2.5.0
             numpy-requirement: "'numpy'"
+            tokenizers-requirement: "'tokenizers'"
+          - python-version: '3.12'
+            pytorch-version: 2.6.0
+            numpy-requirement: "'numpy'"
+            tokenizers-requirement: "'tokenizers'"
     steps:
       - uses: conda-incubator/setup-miniconda@v3
       - run: conda install -n test ffmpeg python=${{ matrix.python-version }}
       - uses: actions/checkout@v4
       - run: echo "$CONDA/envs/test/bin" >> $GITHUB_PATH
-      - run: pip3 install .["dev"] ${{ matrix.numpy-requirement }} torch==${{ matrix.pytorch-version }}+cpu --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple
+      - run: pip3 install . ${{ matrix.numpy-requirement }} torch==${{ matrix.pytorch-version }}+cpu --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple
       - run: python test/test_transcribe.py
       - run: python test/test_align.py
       - run: python test/test_refine.py
       - run: python test/test_locate.py
+      - run: pip3 install .["fw"] ${{ matrix.tokenizers-requirement }}
+      - run: python test/test_transcribe.py load_faster_whisper
+      - run: python test/test_align.py load_faster_whisper
+      - run: python test/test_refine.py load_faster_whisper
+      - run: pip3 install .["hf"] 'transformers<=4.46.3'
+      - run: python test/test_transcribe.py load_hf_whisper
+      - run: python test/test_align.py load_hf_whisper
+      - run: python test/test_refine.py load_hf_whisper

test/test_align.py

@@ -1,51 +1,60 @@
 import os
+import sys
 import torch
 import stable_whisper
 
 
+def get_load_method():
+    if len(sys.argv) >= 2:
+        return getattr(stable_whisper, sys.argv[1]), sys.argv[1]
+    return stable_whisper.load_model, 'load_model'
+
+
 def check_result(result, expected_text: str, test_name: str):
     assert result.text == expected_text
 
     timing_checked = False
-    for segment in result:
-        for word in segment:
-            assert word.start < word.end, (word.start, word.end, test_name)
-            if word.word.strip(" ,") == "americans":
-                assert word.start <= 1.8, (word.start, test_name)
-                assert word.end >= 1.8, (word.end, test_name)
-                timing_checked = True
+    all_words = result.all_words()
+    fail_count = 0
+    for word in all_words:
+        if word.start >= word.end:
+            fail_count += 1
+        if word.word.strip(" ,") == "americans":
+            assert word.start <= 1.8, (word.start, test_name)
+            assert word.end >= 1.8, (word.end, test_name)
+            timing_checked = True
+    fail_rate = fail_count / len(all_words)
+    print(f'Fail Count: {fail_count} / {len(all_words)}  ({test_name})\n')
+    assert fail_rate < 0.1, (fail_rate, fail_count, test_name)
 
     assert timing_checked, test_name
 
 
 def test_align(model_names):
     device = "cuda" if torch.cuda.is_available() else "cpu"
     audio_path = os.path.join(os.path.dirname(__file__), "jfk.flac")
-    models = [stable_whisper.load_model(name, device=device) for name in model_names]
-    orig_result = models[0].transcribe(
+    load, load_name = get_load_method()
+    models = [(load(name, device=device), f'{load_name}->{name}') for name in model_names]
+    orig_result = models[0][0].transcribe(
         audio_path, language='en', temperature=0.0, word_timestamps=True
     )
     for word in orig_result.all_words():
         word.word = word.word.replace('Americans', 'americans')
 
     def single_test(m, meth: str, prep, extra_check, **kwargs):
-        model_type = 'multilingual-model' if m.is_multilingual else 'en-model'
+        m, model_type = m
         meth = getattr(m, meth)
-        test_name = f'{model_type} {meth.__name__}(WhisperResult)'
-        try:
-            result = meth(audio_path, orig_result, **kwargs)
-            check_same_segment_text(orig_result, result)
-        except Exception as e:
-            raise Exception(f'failed test {test_name} -> {e.__class__.__name__}: {e}')
+        test_name = f'{model_type} + {meth.__name__}(WhisperResult)'
+        print(f'Start Test: {test_name}')
+        result = meth(audio_path, orig_result, **kwargs)
+        check_same_segment_text(orig_result, result)
         check_result(result, orig_result.text, test_name)
 
-        test_name = f'{model_type} {meth.__name__}(plain-text)'
-        try:
-            result = meth(audio_path, prep(orig_result), language=orig_result.language)
-            if extra_check:
-                extra_check(orig_result, result)
-        except Exception as e:
-            raise Exception(f'failed test {test_name} -> {e.__class__.__name__}: {e}')
+        test_name = f'{model_type} + {meth.__name__}(plain-text)'
+        print(f'Start Test: {test_name}')
+        result = meth(audio_path, prep(orig_result), language=orig_result.language)
+        if extra_check:
+            extra_check(orig_result, result)
         check_result(result, orig_result.text, test_name)
 
     def get_text(res):

test/test_refine.py

@@ -1,17 +1,24 @@
 import os
+import sys
 import torch
 import stable_whisper
 
 
-def check_result(result, orig_result, expect_change: bool = True):
+def get_load_method():
+    if len(sys.argv) >= 2:
+        return getattr(stable_whisper, sys.argv[1]), sys.argv[1]
+    return stable_whisper.load_model, 'load_model'
+
+
+def check_result(result, orig_result, test_name: str, expect_change: bool = True):
 
     timing_checked = False
     for segment in result:
         for word in segment:
             assert word.start < word.end
             if word.word.strip(" ,").lower() == "americans":
-                assert word.start <= 1.8, word.start
-                assert word.end >= 1.8, word.end
+                assert word.start <= 1.8, (word.start, test_name)
+                assert word.end >= 1.8, (word.end, test_name)
                 timing_checked = True
 
     if expect_change:
@@ -25,7 +32,8 @@ def check_result(result, orig_result, expect_change: bool = True):
 
 def test_refine(model0_name: str, model1_name: str):
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model0 = stable_whisper.load_model(model0_name, device=device)
+    load, load_name = get_load_method()
+    model0 = load(model0_name, device=device)
     audio_path = os.path.join(os.path.dirname(__file__), "jfk.flac")
 
     language = "en" if model0_name.endswith(".en") else None
@@ -36,7 +44,7 @@ def test_refine(model0_name: str, model1_name: str):
     model1 = stable_whisper.load_model(model1_name, device=device)
 
     result = model1.refine(audio_path, orig_result, inplace=False)
-    check_result(result, orig_result, True)
+    check_result(result, orig_result, load_name, True)
 
 
 def test():

test/test_transcribe.py

@@ -1,42 +1,52 @@
 import os
+import sys
 import torch
 import stable_whisper
 
 
-def check_result(result):
-    assert result.language == "en"
+def get_load_method():
+    if len(sys.argv) >= 2:
+        return getattr(stable_whisper, sys.argv[1]), sys.argv[1]
+    return stable_whisper.load_model, 'load_model'
+
+
+def check_result(result, test_name: str):
+    assert result.language in ('en', 'english'), result.language
 
     transcription = result.text.lower()
-    assert "my fellow americans" in transcription
-    assert "your country" in transcription
-    assert "do for you" in transcription
+    assert "my fellow americans" in transcription, test_name
+    assert "your country" in transcription, test_name
+    assert "do for you" in transcription, test_name
 
     timing_checked = False
     for segment in result:
         for word in segment:
             assert word.start < word.end
             if word.word.strip(" ,").lower() == "americans":
-                assert word.start <= 1.8, word.start
-                assert word.end >= 1.8, word.end
+                assert word.start <= 1.8, (word.start, test_name)
+                assert word.end >= 1.8, (word.end, test_name)
                 timing_checked = True
 
     assert timing_checked
 
 
 def test_transcribe(model_name: str):
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = stable_whisper.load_model(model_name, device=device)
+    load, load_name = get_load_method()
+    model = load(model_name, device=device)
     audio_path = os.path.join(os.path.dirname(__file__), "jfk.flac")
 
     language = "en" if model_name.endswith(".en") else None
     result = model.transcribe(
         audio_path, language=language, temperature=0.0, word_timestamps=True
     )
-    check_result(result)
+    check_result(result, f'{load_name}->transcribe')
+    if not hasattr(model, 'transcribe_minimal'):
+        return
     result = model.transcribe_minimal(
         audio_path, language=language, temperature=0.0, word_timestamps=True
     )
-    check_result(result)
+    check_result(result, f'{load_name}->transcribe_minimal')
 
 
 def test():