add task, language keyword to transcribe

This commit is contained in:
Simon
2023-05-20 15:30:25 +02:00
parent d8a2b4ffc9
commit 1fc965bc1a

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)
@ -229,13 +229,12 @@ 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 language = language or self.tokenizer.language_code
if self.tokenizer is None: task = task or self.tokenizer.task
language = self.detect_language(audio) 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="transcribe", language=language) self.tokenizer = faster_whisper.tokenizer.Tokenizer(self.model.hf_tokenizer,
del_tokenizer = True self.model.model.is_multilingual, task=task,
else: language=language)
language = self.tokenizer.language_code
segments: List[SingleSegment] = [] segments: List[SingleSegment] = []
batch_size = batch_size or self._batch_size batch_size = batch_size or self._batch_size
@ -251,9 +250,6 @@ class FasterWhisperPipeline(Pipeline):
} }
) )
if del_tokenizer:
self.tokenizer = None
return {"segments": segments, "language": language} return {"segments": segments, "language": language}