chrissewell/whisperX
1
0
Fork 0
mirror of https://github.com/m-bain/whisperX.git synced 2025-07-01 18:17:27 -04:00
Code Issues Packages Projects Releases Wiki Activity

fix preset language

Browse Source
This commit is contained in:
briguetjo
2023-07-29 18:13:36 +02:00
parent d80b98601b
commit d87f6268d0
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
whisperx/asr.py
View File

@ -1,6 +1,6 @@
import os import os
import warnings import warnings
from typing import List, Union from typing import List, Union, Optional
import ctranslate2 import ctranslate2
import faster_whisper import faster_whisper
@ -27,7 +27,7 @@ def load_model(whisper_arch,
device_index=0, device_index=0,
compute_type="float16", compute_type="float16",
asr_options=None, asr_options=None,
language=None, language : Optional[str] = None,
vad_options=None, vad_options=None,
model=None, model=None,
task="transcribe", task="transcribe",
@ -104,9 +104,7 @@ def load_model(whisper_arch,
vad_model = load_vad_model(torch.device(device), use_auth_token=None, **default_vad_options) vad_model = load_vad_model(torch.device(device), use_auth_token=None, **default_vad_options)
return FasterWhisperPipeline(model, vad_model, default_asr_options, tokenizer) return FasterWhisperPipeline(model, vad_model, default_asr_options, tokenizer, language)
class WhisperModel(faster_whisper.WhisperModel): class WhisperModel(faster_whisper.WhisperModel):
''' '''
@ -183,11 +181,13 @@ class FasterWhisperPipeline(Pipeline):
vad, vad,
options, options,
tokenizer=None, tokenizer=None,
language : Optional[str] = None,
device: Union[int, str, "torch.device"] = -1, device: Union[int, str, "torch.device"] = -1,
framework = "pt", framework = "pt",
**kwargs **kwargs
): ):
self.model = model self.model = model
self.preset_language = language
self.tokenizer = tokenizer self.tokenizer = tokenizer
self.options = options self.options = options
self._batch_size = kwargs.pop("batch_size", None) self._batch_size = kwargs.pop("batch_size", None)
@ -258,7 +258,7 @@ class FasterWhisperPipeline(Pipeline):
vad_segments = self.vad_model({"waveform": torch.from_numpy(audio).unsqueeze(0), "sample_rate": SAMPLE_RATE}) 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, 30)
if self.tokenizer is None: if self.preset_language is None:
language = language or self.detect_language(audio) language = language or self.detect_language(audio)
task = task or "transcribe" task = task or "transcribe"
self.tokenizer = faster_whisper.tokenizer.Tokenizer(self.model.hf_tokenizer, self.tokenizer = faster_whisper.tokenizer.Tokenizer(self.model.hf_tokenizer,