Upload KotobaWhisperPipeline

kotoba_whisper.py

@@ -1,5 +1,6 @@
-from typing import Union, Optional, Dict, List, Any
 import requests
+from typing import Union, Optional, Dict, List, Any
+from collections import defaultdict
 
 import torch
 import numpy as np
@@ -38,12 +39,12 @@ class Punctuator:
         return [
             {
                 'timestamp': c['timestamp'],
+                'speaker': c['speaker'],
                 'text': validate_punctuation(c['text'], "".join(e))
             } for c, e in zip(pipeline_chunk, text_edit)
         ]
 
 
-
 class SpeakerDiarization:
 
     def __init__(self,
@@ -58,7 +59,12 @@ class SpeakerDiarization:
                 model_id_diarizers
             ).to_pyannote_model().to(self.device)
 
-    def __call__(self, audio: Union[torch.Tensor, np.ndarray], sampling_rate: int) -> Annotation:
+    def __call__(self,
+                 audio: Union[torch.Tensor, np.ndarray],
+                 sampling_rate: int,
+                 num_speakers: Optional[int] = None,
+                 min_speakers: Optional[int] = None,
+                 max_speakers: Optional[int] = None) -> Annotation:
         if sampling_rate is None:
             raise ValueError("sampling_rate must be provided")
         if type(audio) is np.ndarray:
@@ -69,7 +75,7 @@ class SpeakerDiarization:
         elif len(audio.shape) > 3:
             raise ValueError("audio shape must be (channel, time)")
         audio = {"waveform": audio.to(self.device), "sample_rate": sampling_rate}
-        output = self.pipeline(audio)
+        output = self.pipeline(audio, num_speakers=num_speakers, min_speakers=min_speakers, max_speakers=max_speakers)
         return output
 
 
@@ -84,8 +90,6 @@ class KotobaWhisperPipeline(AutomaticSpeechRecognitionPipeline):
                  device: Union[int, "torch.device"] = None,
                  device_pyannote: Union[int, "torch.device"] = None,
                  torch_dtype: Optional[Union[str, "torch.dtype"]] = None,
-                 return_unique_speaker: bool = True,
-                 punctuator: bool = False,
                  **kwargs):
         self.type = "seq2seq_whisper"
         if device is None:
@@ -99,11 +103,7 @@ class KotobaWhisperPipeline(AutomaticSpeechRecognitionPipeline):
             model_id=model_pyannote,
             model_id_diarizers=model_diarizers
         )
-        self.return_unique_speaker = return_unique_speaker
-        if punctuator:
-            self.punctuator = Punctuator()
-        else:
-            self.punctuator = None
+        self.punctuator = None
         super().__init__(
             model=model,
             feature_extractor=feature_extractor,
@@ -113,6 +113,71 @@ class KotobaWhisperPipeline(AutomaticSpeechRecognitionPipeline):
             **kwargs
         )
 
+    def _sanitize_parameters(self,
+                             chunk_length_s=None,
+                             stride_length_s=None,
+                             ignore_warning=None,
+                             decoder_kwargs=None,
+                             return_timestamps=None,
+                             return_language=None,
+                             generate_kwargs=None,
+                             max_new_tokens=None,
+                             add_punctuation: bool =False,
+                             return_unique_speaker: bool =True,
+                             num_speakers: Optional[int] = None,
+                             min_speakers: Optional[int] = None,
+                             max_speakers: Optional[int] = None):
+        # No parameters on this pipeline right now
+        preprocess_params = {}
+        if chunk_length_s is not None:
+            preprocess_params["chunk_length_s"] = chunk_length_s
+        if stride_length_s is not None:
+            preprocess_params["stride_length_s"] = stride_length_s
+
+        forward_params = defaultdict(dict)
+        if max_new_tokens is not None:
+            forward_params["max_new_tokens"] = max_new_tokens
+        if generate_kwargs is not None:
+            if max_new_tokens is not None and "max_new_tokens" in generate_kwargs:
+                raise ValueError(
+                    "`max_new_tokens` is defined both as an argument and inside `generate_kwargs` argument, please use"
+                    " only 1 version"
+                )
+            forward_params.update(generate_kwargs)
+
+        postprocess_params = {}
+        if decoder_kwargs is not None:
+            postprocess_params["decoder_kwargs"] = decoder_kwargs
+        if return_timestamps is not None:
+            # Check whether we have a valid setting for return_timestamps and throw an error before we perform a forward pass
+            if self.type == "seq2seq" and return_timestamps:
+                raise ValueError("We cannot return_timestamps yet on non-CTC models apart from Whisper!")
+            if self.type == "ctc_with_lm" and return_timestamps != "word":
+                raise ValueError("CTC with LM can only predict word level timestamps, set `return_timestamps='word'`")
+            if self.type == "ctc" and return_timestamps not in ["char", "word"]:
+                raise ValueError(
+                    "CTC can either predict character level timestamps, or word level timestamps. "
+                    "Set `return_timestamps='char'` or `return_timestamps='word'` as required."
+                )
+            if self.type == "seq2seq_whisper" and return_timestamps == "char":
+                raise ValueError(
+                    "Whisper cannot return `char` timestamps, only word level or segment level timestamps. "
+                    "Use `return_timestamps='word'` or `return_timestamps=True` respectively."
+                )
+            forward_params["return_timestamps"] = return_timestamps
+            postprocess_params["return_timestamps"] = return_timestamps
+        if return_language is not None:
+            if self.type != "seq2seq_whisper":
+                raise ValueError("Only Whisper can return language for now.")
+            postprocess_params["return_language"] = return_language
+        postprocess_params["return_language"] = return_language
+        postprocess_params["add_punctuation"] = add_punctuation
+        postprocess_params["return_unique_speaker"] = return_unique_speaker
+        postprocess_params["num_speakers"] = num_speakers
+        postprocess_params["min_speakers"] = min_speakers
+        postprocess_params["max_speakers"] = max_speakers
+        return preprocess_params, forward_params, postprocess_params
+
     def preprocess(self, inputs, chunk_length_s=0, stride_length_s=None):
         if isinstance(inputs, str):
             if inputs.startswith("http://") or inputs.startswith("https://"):
@@ -259,18 +324,31 @@ class KotobaWhisperPipeline(AutomaticSpeechRecognitionPipeline):
                     model_outputs,
                     decoder_kwargs: Optional[Dict] = None,
                     return_language=None,
+                    add_punctuation: bool = False,
+                    return_unique_speaker: bool = True,
+                    num_speakers: Optional[int] = None,
+                    min_speakers: Optional[int] = None,
+                    max_speakers: Optional[int] = None,
                     *args,
                     **kwargs):
         assert len(model_outputs) > 0
-        audio_array = list(model_outputs)[0]["audio_array"]
-        sd = self.model_speaker_diarization(audio_array, sampling_rate=self.feature_extractor.sampling_rate)
-        timelines = sd.get_timeline()
         outputs = super().postprocess(
             model_outputs=model_outputs,
             decoder_kwargs=decoder_kwargs,
             return_timestamps=True,
             return_language=return_language
         )
+        audio_array = outputs.pop("audio_array")[0]
+        sd = self.model_speaker_diarization(
+            audio_array,
+            num_speakers=num_speakers,
+            min_speakers=min_speakers,
+            max_speakers=max_speakers,
+            sampling_rate=self.feature_extractor.sampling_rate
+        )
+        diarization_result = {s: [[i.start, i.end] for i in sd.label_timeline(s)] for s in sd.labels()}
+        timelines = sd.get_timeline()
+
         pointer_ts = 0
         pointer_chunk = 0
         new_chunks = []
@@ -306,18 +384,19 @@ class KotobaWhisperPipeline(AutomaticSpeechRecognitionPipeline):
                     pointer_ts += 1
         for i in new_chunks:
             if "speaker" in i:
-                if self.return_unique_speaker:
+                if return_unique_speaker:
                     i["speaker"] = [i["speaker"][0]]
                 else:
                     i["speaker"] = list(set(i["speaker"]))
             else:
                 i["speaker"] = []
         outputs["chunks"] = new_chunks
-        if self.punctuator:
+        if add_punctuation:
+            if self.punctuator is None:
+                self.punctuator = Punctuator()
             outputs["chunks"] = self.punctuator.punctuate(outputs["chunks"])
         outputs["text"] = "".join([c["text"] for c in outputs["chunks"]])
         outputs["speakers"] = sd.labels()
-        outputs.pop("audio_array")
         speakers = []
         for s in outputs["speakers"]:
             chunk_s = [c for c in outputs["chunks"] if s in c["speaker"]]
@@ -326,5 +405,5 @@ class KotobaWhisperPipeline(AutomaticSpeechRecognitionPipeline):
                 outputs[f"text/{s}"] = "".join([c["text"] for c in outputs["chunks"] if s in c["speaker"]])
                 speakers.append(s)
         outputs["speakers"] = speakers
+        outputs["diarization_result"] = diarization_result
         return outputs
-