[ASR] Add optimization util for linear sum assignment algorithm
State: closed | Pull request: yes
Labels: name: ASR

# What does this PR do ?

1. Add a optimization util file for linear sum assignment (LSA) solver algorithm for online-diarization/multi-speaker-ASR
LSA problem solver is needed for the following tasks in NeMo: 
        (1) Permutation Invariant Loss (PIL) for diarization model training
        (2) Label permutation matching for online speaker diarzation 
        (3) Concatenated minimum-permutation Word Error Rate (cp-WER) calculation

What is LSA solver algorithm? [Google OR-tools LSA Solver](https://developers.google.com/optimization/assignment/linear_assignment)

The NeMo linear_sum_assignment function is compared with `scipy.optimization.linear_sum_assingment`.
In the [unit-test ](https://github.com/NVIDIA/NeMo/blob/49d4e80ddf1138cea9593392a9a9e76be9ca58ec/tests/collections/asr/test_diar_utils.py#L792)for NeMo LSA solver, the result is compared with the `scipy` version of `linear_sum_assignment`. 


2. Removing `@torch.jit.script` decorator in speaker_utils.py since it creates type-errors when the code is not used for production purpose. 
     Instead, all `torch.jit.script` required classes and functions are tested in `test_diar_utils.py`. 
     Take a look at [these tests](https://github.com/NVIDIA/NeMo/blob/49d4e80ddf1138cea9593392a9a9e76be9ca58ec/tests/collections/asr/test_diar_utils.py#L606) for checking jit_script = [True/False] and cuda = [True/False] (testing total 4 combinations)

3. Also refactored some of the functions in online diarization 
    - replaced scipy LSA solver to NeMo LSA solver in `online_clustering.py`.

5. Added a couple of functions in `der.py` for online diarization DER calculation. 
    - replaced scipy LSA solver to NeMo LSA solver in `der.py`.

**Collection**: [ASR]

# Changelog 
- nemo/collections/asr/metrics/der.py
: replaced scipy LSA solver to NeMo LSA solver in `calculate_session_cpWER ` function. 
: Added two functions for online diarization evaluations: `get_partial_ref_labels` and `get_online_DER_stats`.

- nemo/collections/asr/models/online_diarizer.py
: Made `_perform_online_clustering` function simpler by moving `get_reduced_mat` and `match_labels` into online clustering function. 

- nemo/collections/asr/parts/utils/offline_clustering.py
: Added `laplacian = laplacian.float().to(torch.device('cpu'))` to avoid jit-scripted module uses GPU even when CPU is specified or vice-versa. This behavior is always tested/checked in `test_diar_utils.py`.

- nemo/collections/asr/parts/utils/online_clustering.py
 : replaced scipy LSA solver to NeMo LSA solver in `get_lsa_speaker_mapping` function.
 : Modified the docstrings of `update_speaker_history_buffer` to make the example easier.

- nemo/collections/asr/parts/utils/optimization_utils.py
  : Fully torch-jit-scriptable, linear sum assignment problem solver class and function were added.

- nemo/collections/asr/parts/utils/speaker_utils.py
 : Removed `@torch.jit.script` decorators since this creates unnecessary warning messages and type related errors when used without scripting.

- tests/collections/asr/test_diar_metrics.py
: Added unit-tests for the newly added function `get_partial_ref_labels` and `get_online_DER_stats`.

- tests/collections/asr/test_diar_utils.py
: Added tests for offline clustering and online clustering for many different cases including:
[jit-script=True, cuda=True], 
[jit-script=True, cuda=False], 
[jit-script=False, cuda=True], 
[jit-script=False, cuda=False] cases 
 which is using the torch-jit-scripted NeMo linear_sum_assignment function.

# Usage

```python
from nemo.collections.asr.parts.utils.optimization_utils import linear_sum_assignment
#An example cost matrix to be solved
cost_matrix = \
torch.tensor([[7, 6, 2], [6, 2, 1], [5, 6, 8]])
row_ind_nm, col_ind_nm = linear_sum_assignment(cost_matrix)
```

# Before your PR is "Ready for review"
**Pre checks**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [x] 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**:
- [x] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed.