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 suppress_numerals

Browse Source
This commit is contained in:
briguetjo
2023-07-29 19:34:51 +02:00
parent 864976af23
commit 225f6b4d69
1 changed files with 28 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
whisperx/asr.py
View File

@ -1,6 +1,6 @@
import os import os
import warnings import warnings
from typing import List, Union, Optional from typing import List, Union, Optional, NamedTuple
import ctranslate2 import ctranslate2
import faster_whisper import faster_whisper
@ -83,13 +83,7 @@ def load_model(whisper_arch,
if asr_options is not None: if asr_options is not None:
default_asr_options.update(asr_options) default_asr_options.update(asr_options)
if default_asr_options["suppress_numerals"]: suppress_numerals = default_asr_options["suppress_numerals"]
if tokenizer is None:
tokenizer = faster_whisper.tokenizer.Tokenizer(model.hf_tokenizer, model.model.is_multilingual, task=task, language="en")
numeral_symbol_tokens = find_numeral_symbol_tokens(tokenizer)
print(f"Suppressing numeral and symbol tokens: {numeral_symbol_tokens}")
default_asr_options["suppress_tokens"] += numeral_symbol_tokens
default_asr_options["suppress_tokens"] = list(set(default_asr_options["suppress_tokens"]))
del default_asr_options["suppress_numerals"] del default_asr_options["suppress_numerals"]
default_asr_options = faster_whisper.transcribe.TranscriptionOptions(**default_asr_options) default_asr_options = faster_whisper.transcribe.TranscriptionOptions(**default_asr_options)
@ -104,7 +98,14 @@ 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, language) return FasterWhisperPipeline(
model=model,
vad=vad_model,
options=default_asr_options,
tokenizer=tokenizer,
language=language,
suppress_numerals=suppress_numerals,
)
class WhisperModel(faster_whisper.WhisperModel): class WhisperModel(faster_whisper.WhisperModel):
''' '''
@ -179,17 +180,19 @@ class FasterWhisperPipeline(Pipeline):
self, self,
model, model,
vad, vad,
options, options : NamedTuple,
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",
language : Optional[str] = None,
suppress_numerals: bool = False,
**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.preset_language = language
self.suppress_numerals = suppress_numerals
self._batch_size = kwargs.pop("batch_size", None) self._batch_size = kwargs.pop("batch_size", None)
self._num_workers = 1 self._num_workers = 1
self._preprocess_params, self._forward_params, self._postprocess_params = self._sanitize_parameters(**kwargs) self._preprocess_params, self._forward_params, self._postprocess_params = self._sanitize_parameters(**kwargs)
@ -272,6 +275,14 @@ class FasterWhisperPipeline(Pipeline):
self.model.model.is_multilingual, task=task, self.model.model.is_multilingual, task=task,
language=language) language=language)
if self.suppress_numerals:
previous_suppress_tokens = self.options.suppress_tokens
numeral_symbol_tokens = find_numeral_symbol_tokens(self.tokenizer)
print(f"Suppressing numeral and symbol tokens: {numeral_symbol_tokens}")
new_suppressed_tokens = numeral_symbol_tokens + self.options.suppress_tokens
new_suppressed_tokens = list(set(new_suppressed_tokens))
self.options = self.options._replace(suppress_tokens=new_suppressed_tokens)
segments: List[SingleSegment] = [] segments: List[SingleSegment] = []
batch_size = batch_size or self._batch_size batch_size = batch_size or self._batch_size
for idx, out in enumerate(self.__call__(data(audio, vad_segments), batch_size=batch_size, num_workers=num_workers)): for idx, out in enumerate(self.__call__(data(audio, vad_segments), batch_size=batch_size, num_workers=num_workers)):
@ -286,9 +297,14 @@ class FasterWhisperPipeline(Pipeline):
} }
) )
# revert the tokenizer if multilingual inference is enabled
if self.preset_language is None: if self.preset_language is None:
self.tokenizer = None self.tokenizer = None
# revert suppressed tokens if suppress_numerals is enabled
if self.suppress_numerals:
self.options = self.options._replace(suppress_tokens=previous_suppress_tokens)
return {"segments": segments, "language": language} return {"segments": segments, "language": language}