Add file class based inference API for diarization

[SeanNaren]

What does this PR do?

Supports:

from nemo.collections.asr.models import NeuralDiarizer

diarizer = NeuralDiarizer.from_pretrained(model_name='diar_msdd_telephonic')

annotation = diarizer("an4_diarize_test.wav")

print(annotation.to_rttm())

User can change the device:

import torch
from nemo.collections.asr.models import NeuralDiarizer

diarizer = NeuralDiarizer.from_pretrained(model_name='diar_msdd_telephonic')
diarizer.to(torch.device('cuda'))

Supports passing the maximum number of speakers, or the actual number of speakers in a file (oracle):

annotation = diarizer("an4_diarize_test.wav", num_speakers=2)

annotation = diarizer("an4_diarize_test.wav", max_speakers=4)

Relies on a temporary directory, which can be specified if preferred:

annotation = diarizer("an4_diarize_test.wav", out_dir="my_temp/")

List of available models:

from nemo.collections.asr.models.msdd_models import NeuralDiarizer

print(NeuralDiarizer.list_available_models())
"""
PretrainedModelInfo(
	pretrained_model_name=diar_msdd_telephonic,
	description=For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:diar_msdd_telephonic,
location=https://api.ngc.nvidia.com/v2/models/nvidia/nemo/diar_msdd_telephonic/versions/1.0.0/files/diar_msdd_telephonic.nemo
)
"""

TODO

• Add tests, ensure that config options are passed correctly and updated correctly. The config logic is quite messy due to mutable objects, but this should be saved for a larger refactor.
• Hopefully in the future, NeuralDiarizer will be a NeMo PT model file too, save for a later refactor

Collection: ASR, Speaker Tasks

Changelog

• Add file class based inference API for Diarization.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

[tango4j]

I resolved the conflict in speaker_utils.py and tested with the script you provided.

@SeanNaren I think it is required to add a simple CI-test for this too? I guess it won't be that difficult.

[SeanNaren]

Thanks guys. I have made modifications as per requested to the code.

I think it is required to add a simple CI-test for this too? I guess it won't be that difficult.

Yep, once we're happy with the API and stuff, can start to add tests for functionality.

The current code seems to be a one run setup; the code isn't setup for diarize to be called multiple times. It instantiates all the models again in the pipeline.

I'll be modifying this behaviour.

[SeanNaren]

High level updates:

• Tried to reduce logging verbosity as much as I can. This is controlled by the verbose flag. This will reduce info logging to a minimum. Much more readable, however I'd like to reduce the amount of spam even further
• NeuralDiarizer is now based on a LightningModule allowing us to easily swap between devices
• It supports multiple files now. Not the cleanest as we're updating a config data-class to do so, but works for now

[tango4j]

I have sent you the branch feat/inference_diar_add with the necessary changes to make sure these things:

• we only use MSDD for telephonic and the rest is clustering diarizer
• download yaml file from raw.githubusercontent.com for automate the yaml file
  update for API.

https://github.com/NVIDIA/NeMo/blob/feat/inference_diar_add/nemo/collections/asr/models/configs/msdd_config.py
https://github.com/NVIDIA/NeMo/blob/feat/inference_diar_add/nemo/collections/asr/models/msdd_models.py

Few points that I strongly insist to change:

• MSDD is currently writing to pred_ovl_rttm folder. We should change it to pred_rttm to unify the output folder Line in msdd_models.py
• Config file should be automatically downloaded from NeMo repo. I implemented all to get it work.
• I temporarily changed __call__ function to diarize_audio and this could probably changed to use __call__ but I should probably spend more time on this. Also there are some ugly config assignment to get clustering diarizer work in current format.

Please ping me if you need to clarify something.

[SeanNaren]

After offline discussion the high level changes are as below:

• Move telephonic configuration to the data-class. This will be our new “default” for the neural diarizer (especially considering the pre-trained model is labelled telephonic). Remove the telephonic yaml file
• The general and meeting yaml files will inherit parameters from the data-class config, just like ASR does. We remove all the parameters that are not changed as they will be inherited from the data-class when loaded

When changes to the defaults are made, they are made in the data-class config file (the python data-class). When changes regarding meeting or general are made, they are made in the yaml files. This means parameters are not duplicated, but there is two different places to update (which I think is fine as maintainers, as the config file is mainly to try different configurations out).

One note to make, is that our neural diarizer model is telephonic. If we release further adapted models to different domains (which is not certain yet) we will need to come up with a mechanism via the API to enable this.

[tango4j]

After offline discussion the high level changes are as below:

• Move telephonic configuration to the data-class. This will be our new “default” for the neural diarizer (especially considering the pre-trained model is labelled telephonic). Remove the telephonic yaml file
• The general and meeting yaml files will inherit parameters from the data-class config, just like ASR does. We remove all the parameters that are not changed as they will be inherited from the data-class when loaded

When changes to the defaults are made, they are made in the data-class config file (the python data-class). When changes regarding meeting or general are made, they are made in the yaml files. This means parameters are not duplicated, but there is two different places to update (which I think is fine as maintainers, as the config file is mainly to try different configurations out).

One note to make, is that our neural diarizer model is telephonic. If we release further adapted models to different domains (which is not certain yet) we will need to come up with a mechanism via the API to enable this.

I recognize that making data class is necessary but editing diarization yaml file this way is too confusing. Even I am having a hard time grasping what is happening with the configuration. Also, setting telephonic as a default configuration for the whole diarization pipeline is also confusing.

Can we just duplicate the configuration (for now, it is telephonic) and just revert all yaml files? This is too much of change on the original diarization pipeline, which is only focusing on API purpose.

[SeanNaren]

Thanks @tango4j I have updated the PR based on your suggestions.

Would we be able to review and merge this now?

[tango4j]

Couple of points to be addressed.
Also it would be great if we can remove the scale-specific tqdm in clustering diarizer.
It is overkill to show it.

[SeanNaren]

Thanks @tango4j everything addressed.

I've reduced the TQDM logging by introducing verbosity to those functions. Everything left is just general NeMo logging, which also needs to definitely be reduced (but in a different PR I believe).

[tango4j]

I will do PRs when I need to update the configs later.
Thank you for your effort on building this.

[titu1994]

Very cool

[nithinraok]

Great work Sean.