diff --git a/constraints.txt b/constraints.txt index fcb9f1e..0d031ad 100644 --- a/constraints.txt +++ b/constraints.txt @@ -1,2 +1,2 @@ -huggingface_hub<0.24 numpy<2 +indic-numtowords @ git+https://github.com/AI4Bharat/indic-numtowords.git \ No newline at end of file diff --git a/diarization/__init__.py b/diarization/__init__.py new file mode 100644 index 0000000..fc4d095 --- /dev/null +++ b/diarization/__init__.py @@ -0,0 +1,3 @@ +from .msdd.msdd import MSDDDiarizer + +__all__ = ["MSDDDiarizer"] diff --git a/nemo_msdd_configs/diar_infer_telephonic.yaml b/diarization/msdd/diar_infer_telephonic.yaml old mode 100755 new mode 100644 similarity index 72% rename from nemo_msdd_configs/diar_infer_telephonic.yaml rename to diarization/msdd/diar_infer_telephonic.yaml index 9062101..d7ca70f --- a/nemo_msdd_configs/diar_infer_telephonic.yaml +++ b/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. diff --git a/diarization/msdd/msdd.py b/diarization/msdd/msdd.py new file mode 100644 index 0000000..0e16f3c --- /dev/null +++ b/diarization/msdd/msdd.py @@ -0,0 +1,100 @@ +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]))) + + 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 diff --git a/diarize.py b/diarize.py index 16dab16..60b5b41 100644 --- a/diarize.py +++ b/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,8 @@ pid = os.getpid() temp_outputs_dir = f"temp_outputs_{pid}" +temp_path = os.path.join(os.getcwd(), "temp_outputs") +os.makedirs(temp_path, exist_ok=True) # Initialize parser parser = argparse.ArgumentParser() @@ -92,6 +91,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 +192,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, -) + diarizer_model = MSDDDiarizer(device=args.device) - -# Initialize NeMo MSDD diarization model -msdd_model = NeuralDiarizer(cfg=create_config(temp_path)).to(args.device) -msdd_model.diarize() - -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: diff --git a/diarize_parallel.py b/diarize_parallel.py index f75fcaa..15705bb 100644 --- a/diarize_parallel.py +++ b/diarize_parallel.py @@ -1,8 +1,8 @@ import argparse import logging +import multiprocessing as mp import os import re -import subprocess import faster_whisper import torch @@ -17,6 +17,7 @@ ) from deepmultilingualpunctuation import PunctuationModel +from diarization import MSDDDiarizer from helpers import ( cleanup, find_numeral_symbol_tokens, @@ -31,226 +32,238 @@ write_srt, ) -mtypes = {"cpu": "int8", "cuda": "float16"} -pid = os.getpid() -temp_outputs_dir = f"temp_outputs_{pid}" +def diarize_parallel(audio: torch.Tensor, device, queue: mp.Queue): + model = MSDDDiarizer(device=device) + result = model.diarize(audio) + queue.put(result) -# Initialize parser -parser = argparse.ArgumentParser() -parser.add_argument( - "-a", "--audio", help="name of the target audio file", required=True -) -parser.add_argument( - "--no-stem", - action="store_false", - dest="stemming", - default=True, - help="Disables source separation." - "This helps with long files that don't contain a lot of music.", -) -parser.add_argument( - "--suppress_numerals", - action="store_true", - dest="suppress_numerals", - default=False, - help="Suppresses Numerical Digits." - "This helps the diarization accuracy but converts all digits into written text.", -) +mp.set_start_method("spawn", force=True) -parser.add_argument( - "--whisper-model", - dest="model_name", - default="large-v2", - help="name of the Whisper model to use", -) +if __name__ == "__main__": + mtypes = {"cpu": "int8", "cuda": "float16"} + pid = os.getpid() + temp_outputs_dir = f"temp_outputs_{pid}" + temp_path = os.path.join(os.getcwd(), temp_outputs_dir) + os.makedirs(temp_path, exist_ok=True) -parser.add_argument( - "--batch-size", - type=int, - dest="batch_size", - default=4, - help="Batch size for batched inference, reduce if you run out of memory, " - "set to 0 for original whisper longform inference", -) + # Initialize parser + parser = argparse.ArgumentParser() + parser.add_argument( + "-a", "--audio", help="name of the target audio file", required=True + ) + parser.add_argument( + "--no-stem", + action="store_false", + dest="stemming", + default=True, + help="Disables source separation." + "This helps with long files that don't contain a lot of music.", + ) -parser.add_argument( - "--language", - type=str, - default=None, - choices=whisper_langs, - help="Language spoken in the audio, specify None to perform language detection", -) + parser.add_argument( + "--suppress_numerals", + action="store_true", + dest="suppress_numerals", + default=False, + help="Suppresses Numerical Digits." + "This helps the diarization accuracy but converts all digits into written text.", + ) -parser.add_argument( - "--device", - dest="device", - default="cuda" if torch.cuda.is_available() else "cpu", - help="if you have a GPU use 'cuda', otherwise 'cpu'", -) + parser.add_argument( + "--whisper-model", + dest="model_name", + default="large-v2", + help="name of the Whisper model to use", + ) -args = parser.parse_args() -language = process_language_arg(args.language, args.model_name) + parser.add_argument( + "--batch-size", + type=int, + dest="batch_size", + default=4, + help="Batch size for batched inference, reduce if you run out of memory, " + "set to 0 for original whisper longform inference", + ) -if args.stemming: - # Isolate vocals from the rest of the audio + parser.add_argument( + "--language", + type=str, + default=None, + choices=whisper_langs, + help="Language spoken in the audio, specify None to perform language detection", + ) - return_code = os.system( - f'python -m demucs.separate -n htdemucs --two-stems=vocals "{args.audio}" -o "{temp_outputs_dir}" --device "{args.device}"' + parser.add_argument( + "--device", + dest="device", + default="cuda" if torch.cuda.is_available() else "cpu", + help="if you have a GPU use 'cuda', otherwise 'cpu'", ) - if return_code != 0: - logging.warning( - "Source splitting failed, using original audio file. " - "Use --no-stem argument to disable it." + 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) + + if args.stemming: + # Isolate vocals from the rest of the audio + + return_code = os.system( + f'python -m demucs.separate -n htdemucs --two-stems=vocals "{args.audio}" -o "{temp_outputs_dir}" --device "{args.device}"' ) + + if return_code != 0: + logging.warning( + "Source splitting failed, using original audio file. " + "Use --no-stem argument to disable it." + ) + vocal_target = args.audio + else: + vocal_target = os.path.join( + temp_outputs_dir, + "htdemucs", + os.path.splitext(os.path.basename(args.audio))[0], + "vocals.wav", + ) + else: vocal_target = args.audio + + audio_waveform = faster_whisper.decode_audio(vocal_target) + + logging.info("Starting Nemo process with vocal_target: ", vocal_target) + results_queue = mp.Queue() + nemo_process = mp.Process( + target=diarize_parallel, + args=( + torch.from_numpy(audio_waveform).unsqueeze(0), + args.device, + results_queue, + ), + ) + nemo_process.start() + # Transcribe the audio file + + whisper_model = faster_whisper.WhisperModel( + args.model_name, device=args.device, compute_type=mtypes[args.device] + ) + whisper_pipeline = faster_whisper.BatchedInferencePipeline(whisper_model) + + suppress_tokens = ( + find_numeral_symbol_tokens(whisper_model.hf_tokenizer) + if args.suppress_numerals + else [-1] + ) + + if args.batch_size > 0: + transcript_segments, info = whisper_pipeline.transcribe( + audio_waveform, + language, + suppress_tokens=suppress_tokens, + batch_size=args.batch_size, + ) else: - vocal_target = os.path.join( - temp_outputs_dir, - "htdemucs", - os.path.splitext(os.path.basename(args.audio))[0], - "vocals.wav", + transcript_segments, info = whisper_model.transcribe( + audio_waveform, + language, + suppress_tokens=suppress_tokens, + vad_filter=True, ) -else: - vocal_target = args.audio -logging.info("Starting Nemo process with vocal_target: ", vocal_target) -nemo_process = subprocess.Popen( - ["python", "nemo_process.py", "-a", vocal_target, "--device", args.device], - stderr=subprocess.PIPE, -) -# Transcribe the audio file + full_transcript = "".join(segment.text for segment in transcript_segments) -whisper_model = faster_whisper.WhisperModel( - args.model_name, device=args.device, compute_type=mtypes[args.device] -) -whisper_pipeline = faster_whisper.BatchedInferencePipeline(whisper_model) -audio_waveform = faster_whisper.decode_audio(vocal_target) -suppress_tokens = ( - find_numeral_symbol_tokens(whisper_model.hf_tokenizer) - if args.suppress_numerals - else [-1] -) + # clear gpu vram + del whisper_model, whisper_pipeline + torch.cuda.empty_cache() -if args.batch_size > 0: - transcript_segments, info = whisper_pipeline.transcribe( - audio_waveform, - language, - suppress_tokens=suppress_tokens, + # Forced Alignment + alignment_model, alignment_tokenizer = load_alignment_model( + args.device, + dtype=torch.float16 if args.device == "cuda" else torch.float32, + ) + + emissions, stride = generate_emissions( + alignment_model, + torch.from_numpy(audio_waveform) + .to(alignment_model.dtype) + .to(alignment_model.device), batch_size=args.batch_size, ) -else: - transcript_segments, info = whisper_model.transcribe( - audio_waveform, - language, - suppress_tokens=suppress_tokens, - vad_filter=True, + + del alignment_model + torch.cuda.empty_cache() + + tokens_starred, text_starred = preprocess_text( + full_transcript, + romanize=True, + language=langs_to_iso[info.language], ) -full_transcript = "".join(segment.text for segment in transcript_segments) + segments, scores, blank_token = get_alignments( + emissions, + tokens_starred, + alignment_tokenizer, + ) -# clear gpu vram -del whisper_model, whisper_pipeline -torch.cuda.empty_cache() + spans = get_spans(tokens_starred, segments, blank_token) + word_timestamps = postprocess_results(text_starred, spans, stride, scores) -# Forced Alignment -alignment_model, alignment_tokenizer = load_alignment_model( - args.device, - dtype=torch.float16 if args.device == "cuda" else torch.float32, -) + nemo_process.join() + if results_queue.empty(): + raise RuntimeError("Diarization process did not return any results.") -emissions, stride = generate_emissions( - alignment_model, - torch.from_numpy(audio_waveform) - .to(alignment_model.dtype) - .to(alignment_model.device), - batch_size=args.batch_size, -) + speaker_ts = results_queue.get_nowait() -del alignment_model -torch.cuda.empty_cache() + wsm = get_words_speaker_mapping(word_timestamps, speaker_ts, "start") -tokens_starred, text_starred = preprocess_text( - full_transcript, - romanize=True, - language=langs_to_iso[info.language], -) + if info.language in punct_model_langs: + # restoring punctuation in the transcript to help realign the sentences + punct_model = PunctuationModel(model="kredor/punctuate-all") -segments, scores, blank_token = get_alignments( - emissions, - tokens_starred, - alignment_tokenizer, -) + words_list = list(map(lambda x: x["word"], wsm)) -spans = get_spans(tokens_starred, segments, blank_token) + labled_words = punct_model.predict(words_list, chunk_size=230) -word_timestamps = postprocess_results(text_starred, spans, stride, scores) + ending_puncts = ".?!" + model_puncts = ".,;:!?" -# Reading timestamps <> Speaker Labels mapping + # We don't want to punctuate U.S.A. with a period. Right? + is_acronym = lambda x: re.fullmatch(r"\b(?:[a-zA-Z]\.){2,}", x) -nemo_return_code = nemo_process.wait() -nemo_error_trace = nemo_process.stderr.read() -assert nemo_return_code == 0, ( - "Diarization failed with the following error:" - f"\n{nemo_error_trace.decode('utf-8')}" -) + for word_dict, labeled_tuple in zip(wsm, labled_words): + word = word_dict["word"] + if ( + word + and labeled_tuple[1] in ending_puncts + and (word[-1] not in model_puncts or is_acronym(word)) + ): + word += labeled_tuple[1] + if word.endswith(".."): + word = word.rstrip(".") + word_dict["word"] = word -ROOT = os.getcwd() -temp_path = os.path.join(ROOT, temp_outputs_dir) - -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: - # restoring punctuation in the transcript to help realign the sentences - punct_model = PunctuationModel(model="kredor/punctuate-all") - - words_list = list(map(lambda x: x["word"], wsm)) - - labled_words = punct_model.predict(words_list, chunk_size=230) - - ending_puncts = ".?!" - model_puncts = ".,;:!?" - - # We don't want to punctuate U.S.A. with a period. Right? - is_acronym = lambda x: re.fullmatch(r"\b(?:[a-zA-Z]\.){2,}", x) - - for word_dict, labeled_tuple in zip(wsm, labled_words): - word = word_dict["word"] - if ( - word - and labeled_tuple[1] in ending_puncts - and (word[-1] not in model_puncts or is_acronym(word)) - ): - word += labeled_tuple[1] - if word.endswith(".."): - word = word.rstrip(".") - word_dict["word"] = word - -else: - logging.warning( - f"Punctuation restoration is not available for {info.language} language." - " Using the original punctuation." - ) + else: + logging.warning( + f"Punctuation restoration is not available for {info.language} language." + " Using the original punctuation." + ) -wsm = get_realigned_ws_mapping_with_punctuation(wsm) -ssm = get_sentences_speaker_mapping(wsm, speaker_ts) + wsm = get_realigned_ws_mapping_with_punctuation(wsm) + ssm = get_sentences_speaker_mapping(wsm, speaker_ts) -with open(f"{os.path.splitext(args.audio)[0]}.txt", "w", encoding="utf-8-sig") as f: - get_speaker_aware_transcript(ssm, f) + with open(f"{os.path.splitext(args.audio)[0]}.txt", "w", encoding="utf-8-sig") as f: + get_speaker_aware_transcript(ssm, f) -with open(f"{os.path.splitext(args.audio)[0]}.srt", "w", encoding="utf-8-sig") as srt: - write_srt(ssm, srt) + with open( + f"{os.path.splitext(args.audio)[0]}.srt", "w", encoding="utf-8-sig" + ) as srt: + write_srt(ssm, srt) -cleanup(temp_path) + cleanup(temp_path) diff --git a/helpers.py b/helpers.py index df36fb2..b1d0f7f 100644 --- a/helpers.py +++ b/helpers.py @@ -1,11 +1,7 @@ -import json import os import shutil import nltk -import wget - -from omegaconf import OmegaConf punct_model_langs = [ "en", @@ -249,60 +245,6 @@ } -def create_config(output_dir): - DOMAIN_TYPE = "telephonic" - CONFIG_LOCAL_DIRECTORY = "nemo_msdd_configs" - CONFIG_FILE_NAME = f"diar_infer_{DOMAIN_TYPE}.yaml" - MODEL_CONFIG_PATH = os.path.join(CONFIG_LOCAL_DIRECTORY, CONFIG_FILE_NAME) - if not os.path.exists(MODEL_CONFIG_PATH): - os.makedirs(CONFIG_LOCAL_DIRECTORY, exist_ok=True) - CONFIG_URL = f"https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/speaker_tasks/diarization/conf/inference/{CONFIG_FILE_NAME}" - MODEL_CONFIG_PATH = wget.download(CONFIG_URL, MODEL_CONFIG_PATH) - - config = OmegaConf.load(MODEL_CONFIG_PATH) - - data_dir = os.path.join(output_dir, "data") - os.makedirs(data_dir, exist_ok=True) - - meta = { - "audio_filepath": os.path.join(output_dir, "mono_file.wav"), - "offset": 0, - "duration": None, - "label": "infer", - "text": "-", - "rttm_filepath": None, - "uem_filepath": None, - } - with open(os.path.join(data_dir, "input_manifest.json"), "w") as fp: - json.dump(meta, fp) - fp.write("\n") - - pretrained_vad = "vad_multilingual_marblenet" - pretrained_speaker_model = "titanet_large" - config.num_workers = 0 - config.diarizer.manifest_filepath = os.path.join(data_dir, "input_manifest.json") - config.diarizer.out_dir = ( - output_dir # Directory to store intermediate files and prediction outputs - ) - - 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 - - def get_word_ts_anchor(s, e, option="start"): if option == "end": return e diff --git a/nemo_msdd_configs/diar_infer_general.yaml b/nemo_msdd_configs/diar_infer_general.yaml deleted file mode 100755 index 860c77f..0000000 --- a/nemo_msdd_configs/diar_infer_general.yaml +++ /dev/null @@ -1,95 +0,0 @@ -# This YAML file is created for all types of offline speaker diarization inference tasks in `/example/speaker_tasks/diarization` folder. -# The inference parameters for VAD, speaker embedding extractor, clustering module, MSDD module, ASR decoder are all included in this YAML file. -# All the keys under `diarizer` key (`vad`, `speaker_embeddings`, `clustering`, `msdd_model`, `asr`) can be selectively used for its own purpose and also can be ignored if the module is not used. -# The configurations in this YAML file is optimized to show balanced performances on various types of domain. VAD is optimized on multilingual ASR datasets and diarizer is optimized on DIHARD3 development set. -# An example line in an input manifest file (`.json` format): -# {"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 -sample_rate: 16000 -batch_size: 64 -device: null # can specify a specific device, i.e: cuda:1 (default cuda if cuda available, else cpu) -verbose: True # enable additional logging - -diarizer: - manifest_filepath: ??? - out_dir: ??? - oracle_vad: False # If True, uses RTTM files provided in the manifest file to get speech activity (VAD) timestamps - collar: 0.25 # Collar value for scoring - ignore_overlap: True # Consider or ignore overlap segments while scoring - - vad: - model_path: vad_multilingual_marblenet # .nemo local model path or pretrained VAD model name - external_vad_manifest: null # This option is provided to use external vad and provide its speech activity labels for speaker embeddings extraction. Only one of model_path or external_vad_manifest should be set - - parameters: # Tuned by detection error rate (false alarm + miss) on multilingual ASR evaluation datasets - window_length_in_sec: 0.63 # Window length in sec for VAD context input - shift_length_in_sec: 0.08 # Shift length in sec for generate frame level VAD prediction - smoothing: False # False or type of smoothing method (eg: median) - overlap: 0.5 # Overlap ratio for overlapped mean/median smoothing filter - onset: 0.5 # Onset threshold for detecting the beginning and end of a speech - offset: 0.3 # Offset threshold for detecting the end of a speech - pad_onset: 0.2 # Adding durations before each speech segment - pad_offset: 0.2 # Adding durations after each speech segment - min_duration_on: 0.5 # Threshold for small non_speech deletion - min_duration_off: 0.5 # Threshold for short speech segment deletion - filter_speech_first: True - - speaker_embeddings: - model_path: titanet_large # .nemo local model path or pretrained model name (titanet_large, ecapa_tdnn or speakerverification_speakernet) - parameters: - window_length_in_sec: [1.9,1.2,0.5] # Window length(s) in sec (floating-point number). either a number or a list. ex) 1.5 or [1.5,1.0,0.5] - shift_length_in_sec: [0.95,0.6,0.25] # Shift length(s) in sec (floating-point number). either a number or a list. ex) 0.75 or [0.75,0.5,0.25] - multiscale_weights: [1,1,1] # Weight for each scale. should be null (for single scale) or a list matched with window/shift scale count. ex) [0.33,0.33,0.33] - save_embeddings: True # If True, save speaker embeddings in pickle format. This should be True if clustering result is used for other models, such as `msdd_model`. - - clustering: - parameters: - oracle_num_speakers: False # If True, use num of speakers value provided in manifest file. - max_num_speakers: 8 # Max number of speakers for each recording. If an oracle number of speakers is passed, this value is ignored. - enhanced_count_thres: 80 # If the number of segments is lower than this number, enhanced speaker counting is activated. - max_rp_threshold: 0.25 # Determines the range of p-value search: 0 < p <= max_rp_threshold. - sparse_search_volume: 10 # The higher the number, the more values will be examined with more time. - maj_vote_spk_count: False # If True, take a majority vote on multiple p-values to estimate the number of speakers. - chunk_cluster_count: 50 # Number of forced clusters (overclustering) per unit chunk in long-form audio clustering. - embeddings_per_chunk: 10000 # Number of embeddings in each chunk for long-form audio clustering. Adjust based on GPU memory capacity. (default: 10000, approximately 40 mins of audio) - - - msdd_model: - model_path: null # .nemo local model path or pretrained model name for multiscale diarization decoder (MSDD) - parameters: - use_speaker_model_from_ckpt: True # If True, use speaker embedding model in checkpoint. If False, the provided speaker embedding model in config will be used. - infer_batch_size: 25 # Batch size for MSDD inference. - sigmoid_threshold: [0.7] # Sigmoid threshold for generating binarized speaker labels. The smaller the more generous on detecting overlaps. - seq_eval_mode: False # If True, use oracle number of speaker and evaluate F1 score for the given speaker sequences. Default is False. - 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: null # 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. - diff --git a/nemo_msdd_configs/diar_infer_meeting.yaml b/nemo_msdd_configs/diar_infer_meeting.yaml deleted file mode 100755 index 0011a85..0000000 --- a/nemo_msdd_configs/diar_infer_meeting.yaml +++ /dev/null @@ -1,94 +0,0 @@ -# This YAML file is created for all types of offline speaker diarization inference tasks in `/example/speaker_tasks/diarization` folder. -# The inference parameters for VAD, speaker embedding extractor, clustering module, MSDD module, ASR decoder are all included in this YAML file. -# All the keys under `diarizer` key (`vad`, `speaker_embeddings`, `clustering`, `msdd_model`, `asr`) can be selectively used for its own purpose and also can be ignored if the module is not used. -# The configurations in this YAML file is suitable for 3~5 speakers participating in a meeting and may not show the best performance on other types of dialogues. -# An example line in an input manifest file (`.json` format): -# {"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 -sample_rate: 16000 -batch_size: 64 -device: null # can specify a specific device, i.e: cuda:1 (default cuda if cuda available, else cpu) -verbose: True # enable additional logging - -diarizer: - manifest_filepath: ??? - out_dir: ??? - oracle_vad: False # If True, uses RTTM files provided in the manifest file to get speech activity (VAD) timestamps - collar: 0.25 # Collar value for scoring - ignore_overlap: True # Consider or ignore overlap segments while scoring - - vad: - model_path: vad_multilingual_marblenet # .nemo local model path or pretrained VAD model name - external_vad_manifest: null # This option is provided to use external vad and provide its speech activity labels for speaker embeddings extraction. Only one of model_path or external_vad_manifest should be set - - parameters: # Tuned parameters for CH109 (using the 11 multi-speaker sessions as dev set) - window_length_in_sec: 0.63 # Window length in sec for VAD context input - shift_length_in_sec: 0.01 # Shift length in sec for generate frame level VAD prediction - smoothing: False # False or type of smoothing method (eg: median) - overlap: 0.5 # Overlap ratio for overlapped mean/median smoothing filter - onset: 0.9 # Onset threshold for detecting the beginning and end of a speech - offset: 0.5 # Offset threshold for detecting the end of a speech - pad_onset: 0 # Adding durations before each speech segment - pad_offset: 0 # Adding durations after each speech segment - min_duration_on: 0 # Threshold for small non_speech deletion - min_duration_off: 0.6 # Threshold for short speech segment deletion - filter_speech_first: True - - speaker_embeddings: - model_path: titanet_large # .nemo local model path or pretrained model name (titanet_large, ecapa_tdnn or speakerverification_speakernet) - parameters: - window_length_in_sec: [3.0,2.5,2.0,1.5,1.0,0.5] # Window length(s) in sec (floating-point number). either a number or a list. ex) 1.5 or [1.5,1.0,0.5] - shift_length_in_sec: [1.5,1.25,1.0,0.75,0.5,0.25] # Shift length(s) in sec (floating-point number). either a number or a list. ex) 0.75 or [0.75,0.5,0.25] - multiscale_weights: [1,1,1,1,1,1] # Weight for each scale. should be null (for single scale) or a list matched with window/shift scale count. ex) [0.33,0.33,0.33] - save_embeddings: True # If True, save speaker embeddings in pickle format. This should be True if clustering result is used for other models, such as `msdd_model`. - - clustering: - parameters: - oracle_num_speakers: False # If True, use num of speakers value provided in manifest file. - max_num_speakers: 8 # Max number of speakers for each recording. If an oracle number of speakers is passed, this value is ignored. - enhanced_count_thres: 80 # If the number of segments is lower than this number, enhanced speaker counting is activated. - max_rp_threshold: 0.25 # Determines the range of p-value search: 0 < p <= max_rp_threshold. - sparse_search_volume: 30 # The higher the number, the more values will be examined with more time. - maj_vote_spk_count: False # If True, take a majority vote on multiple p-values to estimate the number of speakers. - chunk_cluster_count: 50 # Number of forced clusters (overclustering) per unit chunk in long-form audio clustering. - embeddings_per_chunk: 10000 # Number of embeddings in each chunk for long-form audio clustering. Adjust based on GPU memory capacity. (default: 10000, approximately 40 mins of audio) - - msdd_model: - model_path: null # .nemo local model path or pretrained model name for multiscale diarization decoder (MSDD) - parameters: - use_speaker_model_from_ckpt: True # If True, use speaker embedding model in checkpoint. If False, the provided speaker embedding model in config will be used. - infer_batch_size: 25 # Batch size for MSDD inference. - sigmoid_threshold: [0.7] # Sigmoid threshold for generating binarized speaker labels. The smaller the more generous on detecting overlaps. - seq_eval_mode: False # If True, use oracle number of speaker and evaluate F1 score for the given speaker sequences. Default is False. - 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. - diff --git a/nemo_process.py b/nemo_process.py deleted file mode 100644 index 55b9837..0000000 --- a/nemo_process.py +++ /dev/null @@ -1,35 +0,0 @@ -import argparse -import os - -import torch - -from nemo.collections.asr.models.msdd_models import NeuralDiarizer -from pydub import AudioSegment - -from helpers import create_config - -pid = os.getppid() -temp_outputs_dir = f"temp_outputs_{pid}" - -parser = argparse.ArgumentParser() -parser.add_argument( - "-a", "--audio", help="name of the target audio file", required=True -) -parser.add_argument( - "--device", - dest="device", - default="cuda" if torch.cuda.is_available() else "cpu", - help="if you have a GPU use 'cuda', otherwise 'cpu'", -) -args = parser.parse_args() - -# convert audio to mono for NeMo combatibility -sound = AudioSegment.from_file(args.audio).set_channels(1) -ROOT = os.getcwd() -temp_path = os.path.join(ROOT, temp_outputs_dir) -os.makedirs(temp_path, exist_ok=True) -sound.export(os.path.join(temp_path, "mono_file.wav"), format="wav") - -# Initialize NeMo MSDD diarization model -msdd_model = NeuralDiarizer(cfg=create_config(temp_path)).to(args.device) -msdd_model.diarize() diff --git a/requirements.txt b/requirements.txt index 928fcb4..f67ff5a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,6 @@ -wget -nemo_toolkit[asr]==2.0.0rc0 +nemo_toolkit[asr]>=2.3.0 nltk faster-whisper>=1.1.0 git+https://github.com/MahmoudAshraf97/demucs.git git+https://github.com/oliverguhr/deepmultilingualpunctuation.git -git+https://github.com/MahmoudAshraf97/ctc-forced-aligner.git - +git+https://github.com/MahmoudAshraf97/ctc-forced-aligner.git \ No newline at end of file