Update tools/extract/asr.py

tools/extract/asr.py

@@ -1,3 +1,4 @@
+import inspect
 from pathlib import Path
 
 import torch
@@ -22,7 +23,6 @@ class ASRProcessor:
     def __init__(self, model_name="large-v2", compute_type="float16"):
         self.model_name = model_name
         # Check if whisperx.load_model accepts compute_type parameter
-        import inspect
 
         if "compute_type" in inspect.signature(whisperx.load_model).parameters:
             self.model = whisperx.load_model(
@@ -34,7 +34,10 @@ class ASRProcessor:
     def get_asr(self, audio_file, return_duration=True):
         assert Path(audio_file).exists(), f"File {audio_file} does not exist"
         audio = whisperx.load_audio(audio_file)
-        result = self.model.transcribe(audio, batch_size=1)
+        if "batch_size" in inspect.signature(self.model.transcribe).parameters:
+            result = self.model.transcribe(audio, batch_size=1)
+        else:
+            result = self.model.transcribe(audio)
         language = result["language"]
         duration = audio.shape[0] / SAMPLE_RATE