Update README.md

README.md

@@ -144,43 +144,46 @@ The model is available for use in the NeMo Framework[5], and can be used as a pr
 from nemo.collections.asr.models import SortformerEncLabelModel
 
 # load model
-diar_model = SortformerEncLabelModel.restore_from(restore_path="diar_sortformer_4spk-v1", map_location=torch.device('cuda'), strict=False)
+diar_model = SortformerEncLabelModel.restore_from(restore_path="/path/to/diar_sortformer_4spk-v1.nemo", map_location=torch.device('cuda'), strict=False)
 ```
 
 ### Input Format
-Input to Sortformer can be either a list of paths to audio files or a jsonl manifest file.
-
+Input to Sortformer can be an individual audio file:
 ```python
-pred_outputs = diar_model.diarize(audio=["/path/to/multispeaker_audio1.wav", "/path/to/multispeaker_audio2.wav"], batch_size=1)
+audio_input="/path/to/multispeaker_audio1.wav"
 ```
-
-Individual audio file can be fed into Sortformer model as follows:
+or a list of paths to audio files:
 ```python
-pred_output1 = diar_model.diarize(audio="/path/to/multispeaker_audio1.wav", batch_size=1)
+audio_input=["/path/to/multispeaker_audio1.wav", "/path/to/multispeaker_audio2.wav"]
 ```
-
-
-To use Sortformer for performing diarization on a multi-speaker audio recording, specify the input as jsonl manifest file, where each line in the file is a dictionary containing the following fields: 
-
+or a jsonl manifest file:
+```python
+audio_input="/path/to/multispeaker_manifest.json"
+```
+where each line is a dictionary containing the following fields:
 ```yaml
 # Example of a line in `multispeaker_manifest.json`
 {
     "audio_filepath": "/path/to/multispeaker_audio1.wav",  # path to the input audio file 
-    "offset": 0 # offset (start) time of the input audio
+    "offset": 0, # offset (start) time of the input audio
     "duration": 600,  # duration of the audio, can be set to `null` if using NeMo main branch
 }
 {
     "audio_filepath": "/path/to/multispeaker_audio2.wav",  
-    "offset": 0,
+    "offset": 900,
     "duration": 580,  
 }
 ```
 
-and then use:
+### Getting Diarization Results
+To perform speaker diarization and get a list of speaker-marked speech segments in the format 'begin_seconds, end_seconds, speaker_index', simply use:
 ```python
-pred_outputs = diar_model.diarize(audio="/path/to/multispeaker_manifest.json", batch_size=1)
+predicted_segments = diar_model.diarize(audio=audio_input, batch_size=1)
+```
+To also obtain tensors of speaker activity probabilities, use:
+```python
+predicted_segments, predicted_probs = diar_model.diarize(audio=audio_input, batch_size=1, include_tensor_outputs=True)
 ```
-
 
 ### Input
 
@@ -190,7 +193,7 @@ This model accepts single-channel (mono) audio sampled at 16,000 Hz.
 
 ### Output
 
-The output of the model is an T x S matrix, where:  
+The output of the model is a T x S matrix, where:  
 - S is the maximum number of speakers (in this model, S = 4).  
 - T is the total number of frames, including zero-padding. Each frame corresponds to a segment of 0.08 seconds of audio.  
 Each element of the T x S matrix represents the speaker activity probability in the [0, 1] range.  For example, a matrix element a(150, 2) = 0.95 indicates a 95% probability of activity for the second speaker during the time range [12.00, 12.08] seconds.
@@ -202,9 +205,27 @@ Each element of the T x S matrix represents the speaker activity probability in
 Sortformer diarizer models are trained on 8 nodes of 8×NVIDIA Tesla V100 GPUs. We use 90 second long training samples and batch size of 4.
 The model can be trained using this [example script](https://github.com/NVIDIA/NeMo/blob/main/examples/speaker_tasks/diarization/neural_diarizer/sortformer_diar_train.py) and [base config](https://github.com/NVIDIA/NeMo/blob/main/examples/speaker_tasks/diarization/conf/neural_diarizer/sortformer_diarizer_hybrid_loss_4spk-v1.yaml).
 
-### Inference
+### Evaluation
 
-Sortformer diarizer models can be performed with post-processing algorithms using inference [example script](https://github.com/NVIDIA/NeMo/blob/main/examples/speaker_tasks/diarization/neural_diarizer/e2e_diarize_speech.py). If you provide the post-processing YAML configs in [`post_processing` folder](https://github.com/NVIDIA/NeMo/tree/main/examples/speaker_tasks/diarization/conf/post_processing) to reproduce the optimized post-processing algorithm for each development dataset.
+To evaluate Sortformer diarizer and save diarization results in RTTM format, use the inference [example script](https://github.com/NVIDIA/NeMo/blob/main/examples/speaker_tasks/diarization/neural_diarizer/e2e_diarize_speech.py):
+```shell
+python [NEMO_GIT_FOLDER]/examples/speaker_tasks/diarization/neural_diarizer/e2e_diarize_speech.py 
+ model_path="/path/to/diar_sortformer_4spk-v1.nemo" 
+ manifest_filepath="/path/to/multispeaker_manifest_with_reference_rttms.json"
+ collar=COLLAR
+ out_rttm_dir="/path/to/output_rttms"
+```
+
+You can provide the post-processing YAML configs from [`post_processing` folder](https://github.com/NVIDIA/NeMo/tree/main/examples/speaker_tasks/diarization/conf/post_processing) to reproduce the optimized post-processing algorithm for each development dataset:
+```shell
+python [NEMO_GIT_FOLDER]/examples/speaker_tasks/diarization/neural_diarizer/e2e_diarize_speech.py 
+ model_path="/path/to/diar_sortformer_4spk-v1.nemo" 
+ manifest_filepath="/path/to/multispeaker_manifest_with_reference_rttms.json"
+ collar=COLLAR
+ bypass_postprocessing=False
+ postprocessing_yaml="/path/to/postprocessing_config.yaml"
+ out_rttm_dir="/path/to/output_rttms"
+```
 
 ### Technical Limitations