Upload KotobaWhisperPipeline

kotoba_whisper.py

@@ -49,6 +49,7 @@ class KotobaWhisperPipeline(AutomaticSpeechRecognitionPipeline):
                  device: Union[int, "torch.device"] = None,
                  device_diarizarization: Union[int, "torch.device"] = None,
                  torch_dtype: Optional[Union[str, "torch.dtype"]] = None,
+                 return_unique_speaker: bool = False,
                  **kwargs):
         self.type = "seq2seq_whisper"
         if device is None:
@@ -58,6 +59,7 @@ class KotobaWhisperPipeline(AutomaticSpeechRecognitionPipeline):
         if type(device_diarizarization) is str:
             device_diarizarization = torch.device(device_diarizarization)
         self.model_speaker_diarization = SpeakerDiarization(model_diarizarization, device_diarizarization)
+        self.return_unique_speaker = return_unique_speaker
         super().__init__(
             model=model,
             feature_extractor=feature_extractor,
@@ -192,6 +194,7 @@ class KotobaWhisperPipeline(AutomaticSpeechRecognitionPipeline):
             )
 
         # custom processing for Whisper timestamps and word-level timestamps
+        generate_kwargs["return_timestamps"] = True
         if inputs.shape[-1] > self.feature_extractor.nb_max_frames:
             generate_kwargs["input_features"] = inputs
         else:
@@ -215,7 +218,7 @@ class KotobaWhisperPipeline(AutomaticSpeechRecognitionPipeline):
                     *args,
                     **kwargs):
         assert len(model_outputs) > 0
-        audio_array = list(model_outputs)[0].pop("audio_array")
+        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(
@@ -229,35 +232,48 @@ class KotobaWhisperPipeline(AutomaticSpeechRecognitionPipeline):
         new_chunks = []
         while True:
             if pointer_ts == len(timelines):
-                new_chunks += outputs["chunks"][pointer_chunk:]
+                ts = timelines[-1]
+                for chunk in outputs["chunks"][pointer_chunk:]:
+                    chunk["speaker"] = sd.get_labels(ts)
+                    new_chunks.append(chunk)
                 break
             if pointer_chunk == len(outputs["chunks"]):
                 break
             ts = timelines[pointer_ts]
+
             chunk = outputs["chunks"][pointer_chunk]
             if "speaker" not in chunk:
-                chunk["speaker"] = set()
+                chunk["speaker"] = []
+
             start, end = chunk["timestamp"]
             if ts.end <= start:
-                chunk["speaker"].update(sd.get_labels(ts))
                 pointer_ts += 1
             elif end <= ts.start:
+                if len(chunk["speaker"]) == 0:
+                    chunk["speaker"] += list(sd.get_labels(ts))
                 new_chunks.append(chunk)
                 pointer_chunk += 1
             else:
+                chunk["speaker"] += list(sd.get_labels(ts))
                 if ts.end >= end:
                     new_chunks.append(chunk)
                     pointer_chunk += 1
                 else:
-                    chunk["speaker"].update(sd.get_labels(ts))
                     pointer_ts += 1
         for i in new_chunks:
             if "speaker" in i:
-                i["speaker"] = list(i["speaker"])
+                if self.return_unique_speaker:
+                    i["speaker"] = [i["speaker"][0]]
+                else:
+                    i["speaker"] = list(set(i["speaker"]))
             else:
                 i["speaker"] = []
         outputs["chunks"] = new_chunks
         outputs["text"] = "".join([c["text"] for c in outputs["chunks"]])
         outputs["speakers"] = sd.labels()
+        outputs.pop("audio_array")
+        for s in outputs["speakers"]:
+            outputs[f"text/{s}"] = "".join([c["text"] for c in outputs["chunks"] if s in c["speaker"]])
+            outputs[f"chunks/{s}"] = [c for c in outputs["chunks"] if s in c["speaker"]]
         return outputs