Refactor diarization

diarization/__init__.py

@@ -0,0 +1,4 @@
+from .msdd.msdd import MSDDDiarizer
+from .sortformer.sortformer import SortformerDiarizer
+
+__all__ = ["MSDDDiarizer", "SortformerDiarizer"]

nemo_msdd_configs/diar_infer_telephonic.yaml → diarization/msdd/diar_infer_telephonic.yaml

@@ -6,7 +6,7 @@
 # {"audio_filepath": "/path/to/audio_file", "offset": 0, "duration": null, "label": "infer", "text": "-", "num_speakers": null, "rttm_filepath": "/path/to/rttm/file", "uem_filepath": "/path/to/uem/file"}
 name: &name "ClusterDiarizer"
 
-num_workers: 1
+num_workers: 0
 sample_rate: 16000
 batch_size: 64
 device: null # can specify a specific device, i.e: cuda:1 (default cuda if cuda available, else cpu)
@@ -65,30 +65,5 @@ diarizer:
       split_infer: True # If True, break the input audio clip to short sequences and calculate cluster average embeddings for inference.
       diar_window_length: 50 # The length of split short sequence when split_infer is True.
       overlap_infer_spk_limit: 5 # If the estimated number of speakers are larger than this number, overlap speech is not estimated.
-
-  asr:
-    model_path: stt_en_conformer_ctc_large # Provide NGC cloud ASR model name. stt_en_conformer_ctc_* models are recommended for diarization purposes.
-    parameters:
-      asr_based_vad: False # if True, speech segmentation for diarization is based on word-timestamps from ASR inference.
-      asr_based_vad_threshold: 1.0 # Threshold (in sec) that caps the gap between two words when generating VAD timestamps using ASR based VAD.
-      asr_batch_size: null # Batch size can be dependent on each ASR model. Default batch sizes are applied if set to null.
-      decoder_delay_in_sec: null # Native decoder delay. null is recommended to use the default values for each ASR model.
-      word_ts_anchor_offset: null # Offset to set a reference point from the start of the word. Recommended range of values is [-0.05  0.2]. 
-      word_ts_anchor_pos: "start" # Select which part of the word timestamp we want to use. The options are: 'start', 'end', 'mid'.
-      fix_word_ts_with_VAD: False # Fix the word timestamp using VAD output. You must provide a VAD model to use this feature.
-      colored_text: False # If True, use colored text to distinguish speakers in the output transcript.
-      print_time: True # If True, the start and end time of each speaker turn is printed in the output transcript.
-      break_lines: False # If True, the output transcript breaks the line to fix the line width (default is 90 chars)
-
-    ctc_decoder_parameters: # Optional beam search decoder (pyctcdecode)
-      pretrained_language_model: null # KenLM model file: .arpa model file or .bin binary file.
-      beam_width: 32
-      alpha: 0.5
-      beta: 2.5
 
-    realigning_lm_parameters: # Experimental feature
-      arpa_language_model: null # Provide a KenLM language model in .arpa format.
-      min_number_of_words: 3 # Min number of words for the left context.
-      max_number_of_words: 10 # Max number of words for the right context.
-      logprob_diff_threshold: 1.2  # The threshold for the difference between two log probability values from two hypotheses.
 

diarization/msdd/msdd.py

@@ -0,0 +1,105 @@
+import json
+import os
+import tempfile
+
+from typing import Union
+
+import torch
+import torchaudio
+
+from nemo.collections.asr.models.msdd_models import NeuralDiarizer
+from nemo.collections.asr.parts.utils.speaker_utils import rttm_to_labels
+from omegaconf import OmegaConf
+
+
+class MSDDDiarizer:
+    def __init__(self, device: Union[str, torch.device]):
+        self.model: NeuralDiarizer = NeuralDiarizer(cfg=create_config()).to(device)
+
+    def diarize(self, audio: torch.Tensor):
+        with tempfile.TemporaryDirectory() as temp_path:
+            torchaudio.save(
+                os.path.join(temp_path, "mono_file.wav"),
+                audio,
+                16000,
+                channels_first=True,
+            )
+
+            manifest_path = os.path.join(temp_path, "manifest.json")
+            meta = {
+                "audio_filepath": os.path.join(temp_path, "mono_file.wav"),
+                "offset": 0,
+                "duration": None,
+                "label": "infer",
+                "text": "-",
+                "rttm_filepath": None,
+                "uem_filepath": None,
+            }
+
+            with open(manifest_path, "w") as f:
+                json.dump(meta, f)
+
+            self.model._initialize_configs(
+                manifest_path=manifest_path,
+                max_speakers=8,
+                num_speakers=None,
+                tmpdir=temp_path,
+                batch_size=24,
+                num_workers=0,
+                verbose=True,
+            )
+            self.model.clustering_embedding.clus_diar_model._diarizer_params.out_dir = (
+                temp_path
+            )
+            self.model.clustering_embedding.clus_diar_model._diarizer_params.manifest_filepath = (
+                manifest_path
+            )
+            self.model.msdd_model.cfg.test_ds.manifest_filepath = manifest_path
+            self.model.diarize()
+
+            pred_labels_clus = rttm_to_labels(
+                os.path.join(temp_path, "pred_rttms", "mono_file.rttm")
+            )
+
+            labels = []
+            for label in pred_labels_clus:
+                start, end, speaker = label.split()
+                start, end = float(start), float(end)
+                start, end = int(start * 1000), int(end * 1000)
+                labels.append((start, end, int(speaker.split("_")[1])))
+
+            # pred_labels_clus = [label.split() for label in pred_labels_clus]
+            # labels = [
+            #     (int(start * 1000), int(end*1000), int(speaker.split("_")[1]))
+            #     for start, end, speaker in pred_labels_clus
+            # ]
+            labels = sorted(labels, key=lambda x: x[0])
+
+        return labels
+
+
+def create_config():
+    config = OmegaConf.load(
+        os.path.join(os.path.dirname(__file__), "diar_infer_telephonic.yaml")
+    )
+    pretrained_vad = "vad_multilingual_marblenet"
+    pretrained_speaker_model = "titanet_large"
+
+    config.diarizer.out_dir = None
+    config.diarizer.manifest_filepath = None
+    config.diarizer.speaker_embeddings.model_path = pretrained_speaker_model
+    config.diarizer.oracle_vad = (
+        False  # compute VAD provided with model_path to vad config
+    )
+    config.diarizer.clustering.parameters.oracle_num_speakers = False
+
+    # Here, we use our in-house pretrained NeMo VAD model
+    config.diarizer.vad.model_path = pretrained_vad
+    config.diarizer.vad.parameters.onset = 0.8
+    config.diarizer.vad.parameters.offset = 0.6
+    config.diarizer.vad.parameters.pad_offset = -0.05
+    config.diarizer.msdd_model.model_path = (
+        "diar_msdd_telephonic"  # Telephonic speaker diarization model
+    )
+
+    return config

diarize.py

@@ -5,7 +5,6 @@
 
 import faster_whisper
 import torch
-import torchaudio
 
 from ctc_forced_aligner import (
     generate_emissions,
@@ -16,11 +15,9 @@
     preprocess_text,
 )
 from deepmultilingualpunctuation import PunctuationModel
-from nemo.collections.asr.models.msdd_models import NeuralDiarizer
 
 from helpers import (
     cleanup,
-    create_config,
     find_numeral_symbol_tokens,
     get_realigned_ws_mapping_with_punctuation,
     get_sentences_speaker_mapping,
@@ -37,6 +34,9 @@
 
 pid = os.getpid()
 temp_outputs_dir = f"temp_outputs_{pid}"
+ROOT = os.getcwd()
+temp_path = os.path.join(ROOT, "temp_outputs")
+os.makedirs(temp_path, exist_ok=True)
 
 # Initialize parser
 parser = argparse.ArgumentParser()
@@ -92,6 +92,13 @@
     help="if you have a GPU use 'cuda', otherwise 'cpu'",
 )
 
+parser.add_argument(
+    "--diarizer",
+    default="msdd",
+    choices=["msdd"],
+    help="Choose the diarization model to use",
+)
+
 args = parser.parse_args()
 language = process_language_arg(args.language, args.model_name)
 
@@ -186,38 +193,15 @@
 
 word_timestamps = postprocess_results(text_starred, spans, stride, scores)
 
+if args.diarizer == "msdd":
+    from diarization import MSDDDiarizer
 
-# convert audio to mono for NeMo combatibility
-ROOT = os.getcwd()
-temp_path = os.path.join(ROOT, temp_outputs_dir)
-os.makedirs(temp_path, exist_ok=True)
-torchaudio.save(
-    os.path.join(temp_path, "mono_file.wav"),
-    torch.from_numpy(audio_waveform).unsqueeze(0).float(),
-    16000,
-    channels_first=True,
-)
-
-
-# Initialize NeMo MSDD diarization model
-msdd_model = NeuralDiarizer(cfg=create_config(temp_path)).to(args.device)
-msdd_model.diarize()
+    diarizer_model = MSDDDiarizer(device=args.device)
 
-del msdd_model
+speaker_ts = diarizer_model.diarize(torch.from_numpy(audio_waveform).unsqueeze(0))
+del diarizer_model
 torch.cuda.empty_cache()
 
-# Reading timestamps <> Speaker Labels mapping
-
-
-speaker_ts = []
-with open(os.path.join(temp_path, "pred_rttms", "mono_file.rttm"), "r") as f:
-    lines = f.readlines()
-    for line in lines:
-        line_list = line.split(" ")
-        s = int(float(line_list[5]) * 1000)
-        e = s + int(float(line_list[8]) * 1000)
-        speaker_ts.append([s, e, int(line_list[11].split("_")[-1])])
-
 wsm = get_words_speaker_mapping(word_timestamps, speaker_ts, "start")
 
 if info.language in punct_model_langs: