init

README.md

@@ -9,7 +9,7 @@ tags:
 widget:
 - example_title: Sample 1
   src: >-
-    https://huggingface.co/datasets/japanese-asr/ja_asr.common_voice_8_0/resolve/main/sample.flac
+    https://huggingface.co/kotoba-tech/kotoba-whisper-v2.2/resolve/main/sample_audio/sample_diarization_japanese.mp3
 pipeline_tag: automatic-speech-recognition
 ---
 
@@ -58,8 +58,9 @@ install the latest version of Transformers.
 ```bash
 pip install --upgrade pip
 pip install --upgrade transformers accelerate torchaudio
-pip install stable-ts==2.16.0
-pip install punctuators==0.0.5
+pip install "punctuators==0.0.5"
+pip install "pyannote.audio"
+pip install git+https://github.com/huggingface/diarizers.git
 ```
 
 ### Transcription

pipeline/kotoba_whisper.py

@@ -13,6 +13,7 @@ from transformers.feature_extraction_sequence_utils import SequenceFeatureExtrac
 from pyannote.audio import Pipeline
 from pyannote.core.annotation import Annotation
 from punctuators.models import PunctCapSegModelONNX
+from diarizers import SegmentationModel
 
 
 class Punctuator:
@@ -45,25 +46,29 @@ class Punctuator:
 
 class SpeakerDiarization:
 
-    def __init__(self, model_id: str, device: torch.device):
+    def __init__(self,
+                 device: torch.device,
+                 model_id: str = "pyannote/speaker-diarization-3.1",
+                 model_id_diarizers: Optional[str] = None):
         self.device = device
         self.pipeline = Pipeline.from_pretrained(model_id)
         self.pipeline = self.pipeline.to(self.device)
-
-    def __call__(self,
-                 audio: Union[str, torch.Tensor, np.ndarray],
-                 sampling_rate: Optional[int] = None) -> Annotation:
-        if type(audio) is torch.Tensor or type(audio) is np.ndarray:
-            if sampling_rate is None:
-                raise ValueError("sampling_rate must be provided")
-            if type(audio) is np.ndarray:
-                audio = torch.as_tensor(audio)
-            audio = torch.as_tensor(audio, dtype=torch.float32)
-            if len(audio.shape) == 1:
-                audio = audio.unsqueeze(0)
-            elif len(audio.shape) > 3:
-                raise ValueError("audio shape must be (channel, time)")
-            audio = {"waveform": audio.to(self.device), "sample_rate": sampling_rate}
+        if model_id_diarizers:
+            self.pipeline._segmentation.model = SegmentationModel().from_pretrained(
+                model_id_diarizers
+            ).to_pyannote_model().to(self.device)
+
+    def __call__(self, audio: Union[torch.Tensor, np.ndarray], sampling_rate: int) -> Annotation:
+        if sampling_rate is None:
+            raise ValueError("sampling_rate must be provided")
+        if type(audio) is np.ndarray:
+            audio = torch.as_tensor(audio)
+        audio = torch.as_tensor(audio, dtype=torch.float32)
+        if len(audio.shape) == 1:
+            audio = audio.unsqueeze(0)
+        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)
         return output
 
@@ -72,23 +77,28 @@ class KotobaWhisperPipeline(AutomaticSpeechRecognitionPipeline):
 
     def __init__(self,
                  model: "PreTrainedModel",
-                 model_diarizarization: str="pyannote/speaker-diarization-3.1",
+                 model_pyannote: str = "pyannote/speaker-diarization-3.1",
+                 model_diarizers: Optional[str] = "diarizers-community/speaker-segmentation-fine-tuned-callhome-jpn",
                  feature_extractor: Union["SequenceFeatureExtractor", str] = None,
                  tokenizer: Optional[PreTrainedTokenizer] = None,
                  device: Union[int, "torch.device"] = None,
-                 device_diarizarization: Union[int, "torch.device"] = None,
+                 device_pyannote: Union[int, "torch.device"] = None,
                  torch_dtype: Optional[Union[str, "torch.dtype"]] = None,
-                 return_unique_speaker: bool = False,
+                 return_unique_speaker: bool = True,
                  punctuator: bool = False,
                  **kwargs):
         self.type = "seq2seq_whisper"
         if device is None:
             device = "cpu"
-        if device_diarizarization is None:
-            device_diarizarization = device
-        if type(device_diarizarization) is str:
-            device_diarizarization = torch.device(device_diarizarization)
-        self.model_speaker_diarization = SpeakerDiarization(model_diarizarization, device_diarizarization)
+        if device_pyannote is None:
+            device_pyannote = device
+        if type(device_pyannote) is str:
+            device_pyannote = torch.device(device_pyannote)
+        self.model_speaker_diarization = SpeakerDiarization(
+            device=device_pyannote,
+            model_id=model_pyannote,
+            model_id_diarizers=model_diarizers
+        )
         self.return_unique_speaker = return_unique_speaker
         if punctuator:
             self.punctuator = Punctuator()
@@ -308,8 +318,13 @@ class KotobaWhisperPipeline(AutomaticSpeechRecognitionPipeline):
         outputs["text"] = "".join([c["text"] for c in outputs["chunks"]])
         outputs["speakers"] = sd.labels()
         outputs.pop("audio_array")
+        speakers = []
         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"]]
+            chunk_s = [c for c in outputs["chunks"] if s in c["speaker"]]
+            if len(chunk_s) != 0:
+                outputs[f"chunks/{s}"] = chunk_s
+                outputs[f"text/{s}"] = "".join([c["text"] for c in outputs["chunks"] if s in c["speaker"]])
+                speakers.append(s)
+        outputs["speakers"] = speakers
         return outputs
 

pipeline/push_pipeline.py

@@ -1,9 +1,11 @@
+"""wget https://huggingface.co/kotoba-tech/kotoba-whisper-v2.2/resolve/main/sample_audio/sample_diarization_japanese.mp3"""
 from pprint import pprint
 from kotoba_whisper import KotobaWhisperPipeline
 from transformers.pipelines import PIPELINE_REGISTRY, pipeline
 from transformers import WhisperForConditionalGeneration, TFWhisperForConditionalGeneration
 
 
+test_audio = "sample_diarization_japanese.mp3"
 model_alias = "kotoba-tech/kotoba-whisper-v2.2"
 PIPELINE_REGISTRY.register_pipeline(
     "kotoba-whisper",
@@ -11,13 +13,9 @@ PIPELINE_REGISTRY.register_pipeline(
     pt_model=WhisperForConditionalGeneration,
     tf_model=TFWhisperForConditionalGeneration
 )
-test_audio = "/Users/asahiu/Desktop/speaker_diariazation_sample_1.wav"
-pipe = pipeline(task="kotoba-whisper", model="kotoba-tech/kotoba-whisper-v2.0", chunk_length_s=15, batch_size=16, return_unique_speaker=True)
-output = pipe(test_audio)
-pprint(output)
 pipe = pipeline(task="kotoba-whisper", model="kotoba-tech/kotoba-whisper-v2.0", chunk_length_s=15, batch_size=16)
-output = pipe(test_audio)
-pprint(output)
+# output = pipe(test_audio)
+# pprint(output)
 pipe.push_to_hub(model_alias)
 
 

pipeline/test_pipeline.py

@@ -1,7 +1,11 @@
+"""wget https://huggingface.co/kotoba-tech/kotoba-whisper-v2.2/resolve/main/sample_audio/sample_diarization_japanese.mp3"""
 from pprint import pprint
 from transformers.pipelines import pipeline
 
-test_audio = "/Users/asahiu/Desktop/speaker_diariazation_sample_1.wav"
-pipe = pipeline(model="kotoba-tech/kotoba-whisper-v2.2", chunk_length_s=15, batch_size=16, trust_remote_code=True)
-output = pipe(test_audio)
+pipe = pipeline(model="kotoba-tech/kotoba-whisper-v2.2", chunk_length_s=None, batch_size=16, trust_remote_code=True)
+output = pipe("sample_diarization_japanese.mp3")
+pprint(output)
+
+pipe = pipeline(model="kotoba-tech/kotoba-whisper-v2.2", chunk_length_s=None, batch_size=16, trust_remote_code=True, return_unique_speaker=False)
+output = pipe("sample_diarization_japanese.mp3")
 pprint(output)

pipeline/test_speaker_diarization.py

@@ -3,46 +3,42 @@
 # Requirement: Sumit access request for the following models.
 #     https://huggingface.co/pyannote/speaker-diarization-3.1
 #     https://huggingface.co/pyannote/segmentation-3.0
+# wget https://huggingface.co/kotoba-tech/kotoba-whisper-v2.2/resolve/main/sample_audio/sample_diarization_japanese.mp3
 import soundfile as sf
 import numpy as np
-from typing import Union, Optional, Dict, List
+from typing import Union, Dict, List
 
 import torch
 from pyannote.audio import Pipeline
+from diarizers import SegmentationModel
 
 
 class SpeakerDiarization:
 
-    def __init__(self, model_id: str):
-        self.pipeline = Pipeline.from_pretrained(model_id)
+    def __init__(self):
+        self.pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1")
+        self.pipeline._segmentation.model = SegmentationModel().from_pretrained(
+            'diarizers-community/speaker-segmentation-fine-tuned-callhome-jpn'
+        ).to_pyannote_model()
 
     def __call__(self,
-                 audio: Union[str, torch.Tensor, np.ndarray],
-                 sampling_rate: Optional[int] = None) -> Dict[str, List[List[float]]]:
-        if type(audio) is torch.Tensor or type(audio) is np.ndarray:
-            if sampling_rate is None:
-                raise ValueError("sampling_rate must be provided")
-            if type(audio) is np.ndarray:
-                audio = torch.as_tensor(audio)
-            audio = torch.as_tensor(audio, dtype=torch.float32)
-            if len(audio.shape) == 1:
-                audio = audio.unsqueeze(0)
-            elif len(audio.shape) > 3:
-                raise ValueError("audio shape must be (channel, time)")
-            audio = {"waveform": audio, "sample_rate": sampling_rate}
+                 audio: Union[torch.Tensor, np.ndarray],
+                 sampling_rate: int) -> Dict[str, List[List[float]]]:
+        if sampling_rate is None:
+            raise ValueError("sampling_rate must be provided")
+        if type(audio) is np.ndarray:
+            audio = torch.as_tensor(audio)
+        audio = torch.as_tensor(audio, dtype=torch.float32)
+        if len(audio.shape) == 1:
+            audio = audio.unsqueeze(0)
+        elif len(audio.shape) > 3:
+            raise ValueError("audio shape must be (channel, time)")
+        audio = {"waveform": audio, "sample_rate": sampling_rate}
         output = self.pipeline(audio)
-        # dictionary: {speaker_id: [[start, end],...]}
         return {s: [[i.start, i.end] for i in output.label_timeline(s)] for s in output.labels()}
 
 
-pipeline = SpeakerDiarization("pyannote/speaker-diarization-3.1")
-root_dir = "/Users/asahiu/Desktop"
-sample_audio_files = ["speaker_diariazation_sample_1.wav", "speaker_diariazation_sample_2.wav"]
+pipeline = SpeakerDiarization()
+a, sr = sf.read("sample_diarization_japanese.mp3")
+print(pipeline(a.T, sampling_rate=sr))
 
-print(sample_audio_file)
-a, sr = sf.read(f"{root_dir}/{sample_audio_file}")
-output = pipeline(a, sampling_rate=sr)
-print(output)
-output = pipeline(f"{root_dir}/{sample_audio_file}")
-print(output)
-print()