Add Wespeaker embeddings model

Author: istupakov

src/onnx_asr/adapters.py

@@ -13,6 +13,7 @@
 
 from onnx_asr.asr import Asr, TimestampedResult
 from onnx_asr.preprocessors.resampler import Resampler
+from onnx_asr.se import SpeakerEmbedding
 from onnx_asr.utils import SampleRates, read_wav_files
 from onnx_asr.vad import SegmentResult, TimestampedSegmentResult, Vad
 
@@ -230,3 +231,50 @@ def _recognize_batch(
                 self.asr, waveforms, waveforms_len, self.asr._get_sample_rate(), {**kwargs}, **self._vadargs
             )
         )
+
+
+class SeAdapter:
+    """Speaker Embedding adapter class."""
+
+    se: SpeakerEmbedding
+    resampler: Resampler
+
+    def __init__(self, se: SpeakerEmbedding, resampler: Resampler):
+        """Create SE adapter."""
+        self.se = se
+        self.resampler = resampler
+
+    def embedding(
+        self,
+        waveform: str | Path | npt.NDArray[np.float32] | list[str | Path | npt.NDArray[np.float32]],
+        *,
+        sample_rate: SampleRates = 16_000,
+        channel: int | Literal["mean"] | None = None,
+    ) -> npt.NDArray[np.float32]:
+        """Compute speaker embedding (single or batch).
+
+        Args:
+            waveform: Path to wav file (only PCM_U8, PCM_16, PCM_24 and PCM_32 formats are supported)
+                      or Numpy array with PCM waveform.
+                      A list of file paths or numpy arrays for batch recognition are also supported.
+            sample_rate: Sample rate for Numpy arrays in waveform.
+            channel: Channel selector for multi-channel audio.
+
+        Returns:
+            speaker embedding results.
+
+        Raises:
+            utils.AudioLoadingError: Audio loading error (onnx-asr specific).
+            FileNotFoundError: File not found error.
+            wave.Error: WAV file reading error.
+            OSError: Other IO errors.
+
+        """
+        if isinstance(waveform, list) and not waveform:
+            return np.array(None, dtype=np.float32)
+
+        waveform_batch = waveform if isinstance(waveform, list) else [waveform]
+        result = self.se.embedding(*self.resampler(*read_wav_files(waveform_batch, sample_rate, channel)))
+        if isinstance(waveform, list):
+            return result
+        return result.squeeze(0)

src/onnx_asr/loader.py

@@ -7,14 +7,15 @@
 
 import onnxruntime as rt
 
-from onnx_asr.adapters import TextResultsAsrAdapter
+from onnx_asr.adapters import SeAdapter, TextResultsAsrAdapter
 from onnx_asr.asr import Asr, Preprocessor
 from onnx_asr.models.gigaam import GigaamV2Ctc, GigaamV2Rnnt, GigaamV3E2eCtc, GigaamV3E2eRnnt
 from onnx_asr.models.kaldi import KaldiTransducer
 from onnx_asr.models.nemo import NemoConformerAED, NemoConformerCtc, NemoConformerRnnt, NemoConformerTdt
 from onnx_asr.models.pyannote import PyAnnoteVad
 from onnx_asr.models.silero import SileroVad
 from onnx_asr.models.tone import TOneCtc
+from onnx_asr.models.wespeaker import WespeakerEmbeddings
 from onnx_asr.models.whisper import WhisperHf, WhisperOrt
 from onnx_asr.onnx import OnnxSessionOptions, get_onnx_providers, update_onnx_providers
 from onnx_asr.preprocessors.numpy_preprocessor import (
@@ -26,6 +27,7 @@
 from onnx_asr.preprocessors.preprocessor import ConcurrentPreprocessor, IdentityPreprocessor, OnnxPreprocessor
 from onnx_asr.preprocessors.resampler import Resampler
 from onnx_asr.resolver import Resolver
+from onnx_asr.se import SpeakerEmbedding
 from onnx_asr.utils import (
     ModelNotSupportedError,
 )
@@ -82,7 +84,7 @@
 
 
 def create_asr_resolver(
-    model: str, local_dir: str | Path | None = None, *, offline: bool | None = None
+    model: str | None = None, local_dir: str | Path | None = None, *, offline: bool | None = None
 ) -> Resolver[AsrTypes]:
     """Create resolver for ASR models."""
     model_types: dict[str, type[AsrTypes]] = {
@@ -120,13 +122,20 @@ def create_asr_resolver(
 
 
 def create_vad_resolver(
-    model: str, local_dir: str | Path | None = None, *, offline: bool | None = None
+    model: str | None = None, local_dir: str | Path | None = None, *, offline: bool | None = None
 ) -> Resolver[VadTypes]:
     """Create resolver for VAD models."""
     model_types: dict[str, type[VadTypes]] = {"silero": SileroVad, "pyannote": PyAnnoteVad}
     return Resolver(model_types, model, local_dir, offline=offline)
 
 
+def create_se_resolver(
+    model: str | None = None, local_dir: str | Path | None = None, *, offline: bool | None = None
+) -> Resolver[WespeakerEmbeddings]:
+    """Create resolver for SE models."""
+    return Resolver(WespeakerEmbeddings, model, local_dir, offline=offline)
+
+
 class PreprocessorRuntimeConfig(OnnxSessionOptions, total=False):
     """Preprocessor runtime config."""
 
@@ -206,30 +215,34 @@ def _create_preprocessor(self, name: str) -> Preprocessor:
     def _create_resampler(self, sample_rate: Literal[8000, 16000]) -> Resampler:
         return Resampler(sample_rate, self.resampler_config)
 
+    def _create_asr_adapter(self, asr: Asr) -> TextResultsAsrAdapter:
+        return TextResultsAsrAdapter(asr, self._create_resampler(asr._get_sample_rate()))
+
+    def _create_se_adapter(self, se: SpeakerEmbedding) -> SeAdapter:
+        return SeAdapter(se, self._create_resampler(se._get_sample_rate()))
+
     def create_asr(
         self,
-        model: str,
+        model: str | ModelNames | ModelTypes | None = None,
         local_dir: str | Path | None = None,
         *,
         quantization: str | None = None,
         offline: bool | None = None,
         config: OnnxSessionOptions | None = None,
-    ) -> Asr:
+    ) -> TextResultsAsrAdapter:
         """Create ASR model."""
         resolver = create_asr_resolver(model, local_dir, offline=offline)
         if config is None:
             config = update_onnx_providers(
                 self.default_onnx_config, excluded_providers=resolver.model_type._get_excluded_providers()
             )
-        return resolver.model_type(resolver.resolve_model(quantization=quantization), self._create_preprocessor, config)
-
-    def create_adapter(self, asr: Asr) -> TextResultsAsrAdapter:
-        """Create ASR adapter."""
-        return TextResultsAsrAdapter(asr, self._create_resampler(asr._get_sample_rate()))
+        return self._create_asr_adapter(
+            resolver.model_type(resolver.resolve_model(quantization=quantization), self._create_preprocessor, config)
+        )
 
     def create_vad(
         self,
-        model: str,
+        model: str | VadNames | None = None,
         local_dir: str | Path | None = None,
         *,
         quantization: str | None = None,
@@ -244,6 +257,25 @@ def create_vad(
             )
         return resolver.model_type(resolver.resolve_model(quantization=quantization), config)
 
+    def create_se(
+        self,
+        model: str | None = None,
+        local_dir: str | Path | None = None,
+        *,
+        quantization: str | None = None,
+        offline: bool | None = None,
+        config: OnnxSessionOptions | None = None,
+    ) -> SeAdapter:
+        """Create SE model."""
+        resolver = create_se_resolver(model, local_dir, offline=offline)
+        if config is None:
+            config = update_onnx_providers(
+                self.default_onnx_config, excluded_providers=resolver.model_type._get_excluded_providers()
+            )
+        return self._create_se_adapter(
+            resolver.model_type(resolver.resolve_model(quantization=quantization), self._create_preprocessor, config)
+        )
+
 
 def load_model(
     model: str | ModelNames | ModelTypes,
@@ -304,7 +336,7 @@ def load_model(
         )
 
     manager = Manager(sess_options, providers, provider_options, preprocessor_config, resampler_config)
-    return manager.create_adapter(manager.create_asr(model, path, quantization=quantization, config=asr_config))
+    return manager.create_asr(model, path, quantization=quantization, config=asr_config)
 
 
 def load_vad(

src/onnx_asr/models/wespeaker.py

@@ -0,0 +1,53 @@
+"""Wespeaker SE implementation."""
+
+from collections.abc import Callable
+from pathlib import Path
+
+import numpy as np
+import numpy.typing as npt
+import onnxruntime as rt
+
+from onnx_asr.asr import Preprocessor
+from onnx_asr.onnx import OnnxSessionOptions
+from onnx_asr.se import SpeakerEmbedding
+from onnx_asr.utils import is_float32_array
+
+
+class WespeakerEmbeddings(SpeakerEmbedding):
+    """Wespeaker embeddings model."""
+
+    def __init__(
+        self,
+        model_files: dict[str, Path],
+        preprocessor_factory: Callable[[str], Preprocessor],
+        onnx_options: OnnxSessionOptions,
+    ):
+        """Create model.
+
+        Args:
+            model_files: Dict with paths to model files.
+            preprocessor_factory: Factory for preprocessor creation.
+            onnx_options: Options for onnxruntime InferenceSession.
+
+        """
+        self._model = rt.InferenceSession(model_files["model"], **onnx_options)
+        self._preprocessor = preprocessor_factory("wespeaker")
+
+    @staticmethod
+    def _get_excluded_providers() -> list[str]:
+        return []
+
+    @staticmethod
+    def _get_model_files(quantization: str | None = None) -> dict[str, str]:
+        suffix = "?" + quantization if quantization else ""
+        return {"config": "config.yaml", "model": f"*{suffix}.onnx"}
+
+    def embedding(
+        self, waveforms: npt.NDArray[np.float32], waveforms_len: npt.NDArray[np.int64]
+    ) -> npt.NDArray[np.float32]:
+        """Compute speaker embedding."""
+        features, _ = self._preprocessor(waveforms, waveforms_len)
+        features -= features.mean(axis=1, keepdims=True)
+        (embs,) = self._model.run(["embs"], {"feats": features})
+        assert is_float32_array(embs)
+        return embs

src/onnx_asr/se.py

@@ -0,0 +1,20 @@
+"""Base Speaker Embedding classes."""
+
+from typing import Literal, Protocol
+
+import numpy as np
+import numpy.typing as npt
+
+
+class SpeakerEmbedding(Protocol):
+    """Speaker Embedding protocol."""
+
+    @staticmethod
+    def _get_sample_rate() -> Literal[8_000, 16_000]:
+        return 16_000
+
+    def embedding(
+        self, waveforms: npt.NDArray[np.float32], waveforms_len: npt.NDArray[np.int64]
+    ) -> npt.NDArray[np.float32]:
+        """Compute speaker embedding."""
+        ...

tests/onnx_asr/test_embedding.py

@@ -0,0 +1,40 @@
+import numpy as np
+import pytest
+
+from onnx_asr.adapters import SeAdapter
+from onnx_asr.loader import Manager
+
+
+@pytest.fixture(scope="module", params=["wespeaker/wespeaker-voxceleb-resnet34"])
+def model(request: pytest.FixtureRequest) -> SeAdapter:
+    manager = Manager()
+    return manager.create_se(request.param)
+
+
+def test_embedding(model: SeAdapter) -> None:
+    rng = np.random.default_rng(0)
+    waveform = rng.random((1 * 16_000), dtype=np.float32)
+
+    result = model.embedding(waveform)
+    assert isinstance(result, np.ndarray)
+    assert result.dtype == np.float32
+    assert result.ndim == 1
+
+
+def test_empty_embedding(model: SeAdapter) -> None:
+    result = model.embedding([])
+    assert isinstance(result, np.ndarray)
+    assert result.dtype == np.float32
+    assert result.ndim == 0
+
+
+def test_embedding_batch(model: SeAdapter) -> None:
+    rng = np.random.default_rng(0)
+    waveform1 = rng.random((2 * 16_000), dtype=np.float32)
+    waveform2 = rng.random((1 * 16_000), dtype=np.float32)
+
+    result = model.embedding([waveform1, waveform2])
+    assert isinstance(result, np.ndarray)
+    assert result.dtype == np.float32
+    assert result.ndim == 2
+    assert result.shape[0] == 2

tests/onnx_asr/test_resolver.py

@@ -10,13 +10,15 @@
     ModelTypes,
     VadNames,
     create_asr_resolver,
+    create_se_resolver,
     create_vad_resolver,
 )
 from onnx_asr.models.kaldi import KaldiTransducer
 from onnx_asr.models.nemo import NemoConformerAED
 from onnx_asr.models.pyannote import PyAnnoteVad
 from onnx_asr.models.silero import SileroVad
 from onnx_asr.models.tone import TOneCtc
+from onnx_asr.models.wespeaker import WespeakerEmbeddings
 from onnx_asr.models.whisper import WhisperHf
 from onnx_asr.resolver import Resolver
 from onnx_asr.utils import (
@@ -208,3 +210,20 @@ def test_resolve_vad_file_not_found_error() -> None:
     loader = create_vad_resolver("silero")
     with pytest.raises(ModelFileNotFoundError):
         loader.resolve_model(quantization="xxx")
+
+
+@pytest.mark.parametrize("model", ["wespeaker/wespeaker-voxceleb-resnet34"])
+def test_se(model: str) -> None:
+    loader = create_se_resolver(model)
+    assert issubclass(loader.model_type, WespeakerEmbeddings)
+    assert not loader.offline
+    assert loader.local_dir is None
+    assert isinstance(loader.repo_id, str)
+
+
+def test_se_with_path(tmp_path: Path) -> None:
+    loader = create_se_resolver(local_dir=tmp_path)
+    assert issubclass(loader.model_type, WespeakerEmbeddings)
+    assert loader.offline
+    assert loader.local_dir == tmp_path
+    assert loader.repo_id is None