update test

linto-ai · Jan 28, 2024 · e1fded5 · e1fded5
1 parent 801a13c
commit e1fded5
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tests/test_transcribe.py b/tests/test_transcribe.py
@@ -661,15 +661,17 @@ def test_hugging_face_model(self):
         )
 
         import tempfile
-        from transformers import WhisperForConditionalGeneration
+        from transformers import WhisperForConditionalGeneration, WhisperProcessor, GenerationConfig
         tempfolder = os.path.join(tempfile.gettempdir(), "tmp_whisper-tiny-french-cased")
 
         for safe_serialization in False, True,:
             for max_shard_size in "100MB", "10GB", :
                 shutil.rmtree(tempfolder, ignore_errors=True)
                 model = WhisperForConditionalGeneration.from_pretrained("qanastek/whisper-tiny-french-cased")
+                processor = WhisperProcessor.from_pretrained("qanastek/whisper-tiny-french-cased")
                 try:
                     model.save_pretrained(tempfolder, safe_serialization=safe_serialization, max_shard_size=max_shard_size)
+                    processor.save_pretrained(tempfolder)
                     self._test_cli_(
                         ["--model", tempfolder, "--verbose", "True"],
                         "verbose", files=["bonjour.wav"], extensions=None,