Fixing the missing sample_rate argument in mixin calling in Sortformer model file (#15228)

* Adding flexible input source for Diarization Mixin

Signed-off-by: taejinp <tango4j@gmail.com>

* Apply isort and black reformatting

Signed-off-by: tango4j <tango4j@users.noreply.github.com>

* Letting diarize() function to use lhotse dataloader

Signed-off-by: taejinp <tango4j@gmail.com>

* Apply isort and black reformatting

Signed-off-by: tango4j <tango4j@users.noreply.github.com>

* One for loop to handle everything

Signed-off-by: taejinp <tango4j@gmail.com>

* Fixing the missing sample_rate and fixed some outdated comments

Signed-off-by: taejinp <tango4j@gmail.com>

* Apply isort and black reformatting

Signed-off-by: tango4j <tango4j@users.noreply.github.com>

* Added sortformer_diar_models.py

Signed-off-by: taejinp <tango4j@gmail.com>

---------

Signed-off-by: taejinp <tango4j@gmail.com>
Signed-off-by: tango4j <tango4j@users.noreply.github.com>
Co-authored-by: tango4j <tango4j@users.noreply.github.com>

Author: tango4j

nemo/collections/asr/models/sortformer_diar_models.py

@@ -1142,6 +1142,7 @@ def diarize(
         """
         return super().diarize(
             audio=audio,
+            sample_rate=sample_rate,
             batch_size=batch_size,
             include_tensor_outputs=include_tensor_outputs,
             postprocessing_yaml=postprocessing_yaml,

nemo/collections/asr/parts/mixins/diarization.py

@@ -49,7 +49,6 @@ def resample_audio(samples: np.ndarray, orig_sr: int, target_sr: int) -> np.ndar
         return samples.astype(np.float32, copy=False)
 
     resampled_samples = samples.astype(np.float32, copy=False)
-    # User-requested API
     resampled_samples = librosa.core.resample(resampled_samples, orig_sr=orig_sr, target_sr=target_sr)
     return resampled_samples.astype(np.float32, copy=False)
 
@@ -132,7 +131,7 @@ class SpkDiarizationMixin(ABC):
     """
     An abstract class for diarize-able models.
 
-    Creates a template function `diarize()` that provides an interface to perform transcription of audio tensors or
+    Creates a template function `diarize()` that provides an interface to perform diarization of audio tensors or
     filepaths.
     """
 
@@ -409,7 +408,7 @@ def _diarize_on_begin(self, audio: Union[str, List[str]], diarcfg: DiarizeConfig
         # Model's mode and device
         diarcfg._internal.training_mode = self.training
 
-        # Switch model to evaluation mode
+        # Save preprocessor settings before switching to evaluation mode
         if hasattr(self, 'preprocessor'):
             if hasattr(self.preprocessor, 'featurizer') and hasattr(self.preprocessor.featurizer, 'dither'):
                 diarcfg._internal.dither_value = self.preprocessor.featurizer.dither
@@ -541,7 +540,8 @@ def _diarize_input_processing(self, audio, diarcfg: DiarizeConfig):
 
         else:
             raise ValueError(
-                f"Input `audio` is of type {type(audio[0])}. " "Only `str` (path to audio file) is supported as input."
+                f"Input `audio` is of type {type(audio[0])}. "
+                "Only `str` (path to audio file) or `np.ndarray` are supported as input."
             )
 
     def _diarize_input_manifest_processing(
@@ -632,7 +632,7 @@ def _diarize_output_processing(self, outputs, uniq_ids, diarcfg: DiarizeConfig)
 
     def _diarize_on_end(self, diarcfg: DiarizeConfig):
         """
-        Internal function to teardown the model after transcription. Perform all teardown and post-checks here.
+        Internal function to teardown the model after diarization. Perform all teardown and post-checks here.
 
         Args:
             diarcfg: The diarization config dataclass. Subclasses can change this to a different dataclass if needed.

tests/collections/speaker_tasks/mixins/test_diarization.py

@@ -200,8 +200,34 @@ def test_diarize_manifest_jsonl_path(self, dummy_model, audio_files, tmp_path: P
     def test_diarize_numpy_single_requires_sample_rate(self, dummy_model, audio_files):
         dummy_model = dummy_model.eval()
         audio1, _, _, _, _ = audio_files
+
+        # Check if it raises an error without sample rate when using a single numpy variable input
         with pytest.raises(ValueError):
-            _ = dummy_model.diarize(audio1, batch_size=1)
+            _ = dummy_model.diarize(audio=audio1, batch_size=1)
+
+        # Set sample rate and check if it works
+        sample_rate = 16000
+        outputs = dummy_model.diarize(audio1, batch_size=1, sample_rate=sample_rate)
+        assert isinstance(outputs, list)
+        assert len(outputs) == 1
+        assert outputs[0] > 0
+
+    @pytest.mark.unit
+    def test_diarize_numpy_list_requires_sample_rate(self, dummy_model, audio_files):
+        dummy_model = dummy_model.eval()
+        audio1, audio2, _, _, _ = audio_files
+        numpy_audio_list = [audio1, audio2]
+        # Check if it raises an error without sample rate when using numpy list input
+        with pytest.raises(ValueError):
+            _ = dummy_model.diarize(audio=numpy_audio_list, batch_size=2)
+
+        # Set sample rate and check if it works
+        sample_rate = 16000
+        outputs = dummy_model.diarize(audio=numpy_audio_list, batch_size=2, sample_rate=sample_rate)
+        assert isinstance(outputs, list)
+        assert len(outputs) == 2
+        assert outputs[0] > 0
+        assert outputs[1] > 0
 
     @pytest.mark.unit
     def test_diarize_numpy_single(self, dummy_model, audio_files):