feat: Add merge chunks chunk_size as arguments.

Suggest from https://github.com/m-bain/whisperX/issues/200#issuecomment-1666507780
This commit is contained in:
陳鈞
2023-08-29 23:09:02 +08:00
parent ef965a03ed
commit eb771cf56d
2 changed files with 6 additions and 3 deletions

View File

@ -247,7 +247,7 @@ class FasterWhisperPipeline(Pipeline):
return final_iterator
def transcribe(
self, audio: Union[str, np.ndarray], batch_size=None, num_workers=0, language=None, task=None
self, audio: Union[str, np.ndarray], batch_size=None, num_workers=0, language=None, task=None, chunk_size=30
) -> TranscriptionResult:
if isinstance(audio, str):
audio = load_audio(audio)
@ -260,7 +260,7 @@ class FasterWhisperPipeline(Pipeline):
yield {'inputs': audio[f1:f2]}
vad_segments = self.vad_model({"waveform": torch.from_numpy(audio).unsqueeze(0), "sample_rate": SAMPLE_RATE})
vad_segments = merge_chunks(vad_segments, 30)
vad_segments = merge_chunks(vad_segments, chunk_size)
if self.tokenizer is None:
language = language or self.detect_language(audio)
task = task or "transcribe"