Merge pull request #269 from sorgfresser/transcribe_keywords

Add transcribe keywords
This commit is contained in:
Max Bain
2023-05-21 12:08:44 +01:00
committed by GitHub

View File

@ -214,7 +214,7 @@ class FasterWhisperPipeline(Pipeline):
return final_iterator return final_iterator
def transcribe( def transcribe(
self, audio: Union[str, np.ndarray], batch_size=None, num_workers=0 self, audio: Union[str, np.ndarray], batch_size=None, num_workers=0, language=None, task=None
) -> TranscriptionResult: ) -> TranscriptionResult:
if isinstance(audio, str): if isinstance(audio, str):
audio = load_audio(audio) audio = load_audio(audio)
@ -228,14 +228,19 @@ 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)
del_tokenizer = False
if self.tokenizer is None: if self.tokenizer is None:
language = self.detect_language(audio) language = language or self.detect_language(audio)
self.tokenizer = faster_whisper.tokenizer.Tokenizer(self.model.hf_tokenizer, self.model.model.is_multilingual, task="transcribe", language=language) task = task or "transcribe"
del_tokenizer = True self.tokenizer = faster_whisper.tokenizer.Tokenizer(self.model.hf_tokenizer,
self.model.model.is_multilingual, task=task,
language=language)
else: else:
language = self.tokenizer.language_code language = language or self.tokenizer.language_code
task = task or self.tokenizer.task
if task != self.tokenizer.task or language != self.tokenizer.language_code:
self.tokenizer = faster_whisper.tokenizer.Tokenizer(self.model.hf_tokenizer,
self.model.model.is_multilingual, task=task,
language=language)
segments: List[SingleSegment] = [] segments: List[SingleSegment] = []
batch_size = batch_size or self._batch_size batch_size = batch_size or self._batch_size
@ -251,9 +256,6 @@ class FasterWhisperPipeline(Pipeline):
} }
) )
if del_tokenizer:
self.tokenizer = None
return {"segments": segments, "language": language} return {"segments": segments, "language": language}