Add HuggingFace WavLM backend to speaker similarity metric

docs/supported_metrics.md

@@ -112,7 +112,7 @@ We include x mark if the metric is auto-installed in versa.
 | 7 |   | NVIDIA Conformer-Transducer X-Large Speech Recognition-based Error Rate | nemo_wer | nemo_wer |[NeMo](https://github.com/NVIDIA/NeMo) | [paper](https://arxiv.org/abs/2005.08100) |
 | 8 | x | Facebook Hubert-Large-Finetuned Speech Recognition-based Error Rate | hubert_wer | hubert_wer |[HuBERT](https://github.com/facebookresearch/fairseq/tree/main/examples/hubert) | [paper](https://arxiv.org/abs/2106.07447) |
 | 9 |   | Emotion2vec similarity (emo2vec) | emo2vec_similarity | emotion_similarity | [emo2vec](https://github.com/ftshijt/emotion2vec/tree/main) | [paper](https://arxiv.org/abs/2312.15185) |
-| 10 | x | Speaker Embedding Similarity  | speaker | spk_similarity | [espnet](https://github.com/espnet/espnet) | [paper](https://arxiv.org/abs/2401.17230) |
+| 10 | x | Speaker Embedding Similarity  | speaker | spk_similarity | [espnet](https://github.com/espnet/espnet), [transformers x-vector (e.g. WavLM)](https://huggingface.co/microsoft/wavlm-base-sv) | [paper](https://arxiv.org/abs/2401.17230) |
 | 11 |   | NOMAD: Unsupervised Learning of Perceptual Embeddings For Speech Enhancement and Non-Matching Reference Audio Quality Assessment |  nomad | nomad |[Nomad](https://github.com/shimhz/nomad/tree/main) | [paper](https://arxiv.org/abs/2309.16284) |
 | 12 |   | Contrastive Language-Audio Pretraining Score (CLAP Score) | clap_score | clap_score | [frechet-audio-distance](https://github.com/gudgud96/frechet-audio-distance) | [paper](https://arxiv.org/abs/2301.12661) |
 | 13 |   | Accompaniment Prompt Adherence (APA) | apa | apa | [Sony-audio-metrics](https://github.com/SonyCSLParis/audio-metrics) | [paper](https://arxiv.org/abs/2404.00775) |

egs/separate_metrics/spk_similarity.yaml

@@ -1,6 +1,10 @@
-# speaker related metrics
-# -- spk_similarity: speaker cosine similarity
-#                    model tag can be any ESPnet-SPK huggingface repo at 
-#                    https://huggingface.co/espnet
-- name: speaker
-  model_tag: default
+# speaker related metrics
+# -- spk_similarity: speaker cosine similarity
+#                    model tag can be:
+#                    - any ESPnet-SPK huggingface repo at
+#                      https://huggingface.co/espnet (default backend)
+#                    - any HuggingFace x-vector speaker model such as
+#                      microsoft/wavlm-base-sv (auto-detected huggingface
+#                      backend, see spk_similarity_wavlm.yaml)
+- name: speaker
+  model_tag: default

egs/separate_metrics/spk_similarity_wavlm.yaml

@@ -0,0 +1,9 @@
+# speaker related metrics (HuggingFace x-vector backend)
+# -- spk_similarity: speaker cosine similarity with microsoft/wavlm-base-sv
+#                    any AutoModelForAudioXVector checkpoint works, e.g.
+#                    microsoft/wavlm-base-sv, microsoft/wavlm-base-plus-sv,
+#                    microsoft/unispeech-sat-base-plus-sv
+#                    the backend is auto-detected from the model tag;
+#                    add `backend: huggingface` to force it explicitly
+- name: speaker
+  model_tag: microsoft/wavlm-base-sv

test/test_metrics/test_speaker.py

@@ -0,0 +1,242 @@
+import os
+
+import numpy as np
+import pytest
+
+from versa.utterance_metrics.speaker import (
+    is_transformers_available,
+    resolve_speaker_backend,
+)
+
+RUN_REAL_MODEL_TESTS = os.environ.get("VERSA_RUN_REAL_MODEL_TESTS") == "1"
+
+
+def test_resolve_backend_default_tag_is_espnet():
+    assert resolve_speaker_backend(model_tag="default") == "espnet"
+
+
+def test_resolve_backend_espnet_prefix_is_espnet():
+    assert resolve_speaker_backend(model_tag="espnet/voxcelebs12_rawnet3") == "espnet"
+
+
+def test_resolve_backend_hf_tag_is_huggingface():
+    assert resolve_speaker_backend(model_tag="microsoft/wavlm-base-sv") == "huggingface"
+
+
+def test_resolve_backend_local_espnet_files_win():
+    assert (
+        resolve_speaker_backend(
+            model_tag="microsoft/wavlm-base-sv",
+            model_path="/path/model.pth",
+            model_config="/path/config.yaml",
+        )
+        == "espnet"
+    )
+
+
+def test_resolve_backend_explicit_override():
+    assert (
+        resolve_speaker_backend(model_tag="default", backend="huggingface")
+        == "huggingface"
+    )
+    assert resolve_speaker_backend(model_tag="some/tag", backend="espnet") == "espnet"
+
+
+def test_resolve_backend_invalid_backend_raises():
+    with pytest.raises(ValueError, match="backend"):
+        resolve_speaker_backend(model_tag="default", backend="wavlm2000")
+
+
+def _fixed_audio(freq, duration=1.0, sample_rate=16000):
+    t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False)
+    envelope = 0.5 + 0.5 * np.sin(2 * np.pi * 0.5 * t)
+    return (envelope * np.sin(2 * np.pi * freq * t)).astype(np.float32)
+
+
+@pytest.mark.real_model
+@pytest.mark.skipif(
+    not RUN_REAL_MODEL_TESTS,
+    reason="Set VERSA_RUN_REAL_MODEL_TESTS=1 to run real model-backed checks",
+)
+@pytest.mark.skipif(
+    not is_transformers_available(), reason="Transformers not available"
+)
+def test_hf_speaker_model_embedding_shape():
+    from versa.utterance_metrics.speaker import hf_speaker_model_setup
+
+    model = hf_speaker_model_setup(model_tag="microsoft/wavlm-base-sv", use_gpu=False)
+    embedding = model(_fixed_audio(150))
+    assert embedding.dim() == 2
+    assert embedding.shape[0] == 1
+    assert embedding.shape[1] > 0
+
+
+@pytest.mark.real_model
+@pytest.mark.skipif(
+    not RUN_REAL_MODEL_TESTS,
+    reason="Set VERSA_RUN_REAL_MODEL_TESTS=1 to run real model-backed checks",
+)
+@pytest.mark.skipif(
+    not is_transformers_available(), reason="Transformers not available"
+)
+def test_hf_speaker_metric_identical_signals():
+    from versa.utterance_metrics.speaker import hf_speaker_model_setup, speaker_metric
+
+    model = hf_speaker_model_setup(model_tag="microsoft/wavlm-base-sv", use_gpu=False)
+    audio = _fixed_audio(150)
+    result = speaker_metric(model, audio, audio, 16000)
+    assert "spk_similarity" in result
+    assert result["spk_similarity"] == pytest.approx(1.0, abs=1e-4)
+
+
+@pytest.mark.real_model
+@pytest.mark.skipif(
+    not RUN_REAL_MODEL_TESTS,
+    reason="Set VERSA_RUN_REAL_MODEL_TESTS=1 to run real model-backed checks",
+)
+@pytest.mark.skipif(
+    not is_transformers_available(), reason="Transformers not available"
+)
+def test_hf_speaker_metric_different_signals():
+    from versa.utterance_metrics.speaker import hf_speaker_model_setup, speaker_metric
+
+    model = hf_speaker_model_setup(model_tag="microsoft/wavlm-base-sv", use_gpu=False)
+    same = speaker_metric(model, _fixed_audio(150), _fixed_audio(150), 16000)
+    diff = speaker_metric(model, _fixed_audio(150), _fixed_audio(420), 16000)
+    assert diff["spk_similarity"] < same["spk_similarity"]
+    assert -1.0 <= diff["spk_similarity"] <= 1.0
+
+
+@pytest.mark.real_model
+@pytest.mark.skipif(
+    not RUN_REAL_MODEL_TESTS,
+    reason="Set VERSA_RUN_REAL_MODEL_TESTS=1 to run real model-backed checks",
+)
+@pytest.mark.skipif(
+    not is_transformers_available(), reason="Transformers not available"
+)
+def test_speaker_metric_class_wavlm_backend():
+    from versa.utterance_metrics.speaker import SpeakerMetric
+
+    metric = SpeakerMetric({"model_tag": "microsoft/wavlm-base-sv", "use_gpu": False})
+    assert metric.backend == "huggingface"
+
+    audio = _fixed_audio(150)
+    result = metric.compute(audio, audio, metadata={"sample_rate": 16000})
+    assert result["spk_similarity"] == pytest.approx(1.0, abs=1e-4)
+
+
+@pytest.mark.real_model
+@pytest.mark.skipif(
+    not RUN_REAL_MODEL_TESTS,
+    reason="Set VERSA_RUN_REAL_MODEL_TESTS=1 to run real model-backed checks",
+)
+@pytest.mark.skipif(
+    not is_transformers_available(), reason="Transformers not available"
+)
+def test_speaker_metric_class_requires_both_signals():
+    from versa.utterance_metrics.speaker import SpeakerMetric
+
+    metric = SpeakerMetric({"model_tag": "microsoft/wavlm-base-sv", "use_gpu": False})
+    with pytest.raises(ValueError, match="Predicted signal"):
+        metric.compute(None, _fixed_audio(150), metadata={"sample_rate": 16000})
+    with pytest.raises(ValueError, match="Reference signal"):
+        metric.compute(_fixed_audio(150), None, metadata={"sample_rate": 16000})
+
+
+def test_speaker_metadata_mentions_both_backends():
+    from versa.utterance_metrics.speaker import _speaker_metadata
+
+    metadata = _speaker_metadata()
+    assert "transformers" in metadata.dependencies
+    assert "espnet2" in metadata.dependencies
+
+
+def test_cache_namespace_speaker_espnet_default():
+    from versa.scorer_shared import configure_metric_cache_dirs
+
+    configs = configure_metric_cache_dirs(
+        [{"name": "speaker", "model_tag": "default"}], cache_folder="/tmp/vc"
+    )
+    assert configs[0]["cache_dir"].endswith("espnet_model_zoo")
+
+
+def test_cache_namespace_speaker_hf_tag():
+    from versa.scorer_shared import configure_metric_cache_dirs
+
+    configs = configure_metric_cache_dirs(
+        [{"name": "speaker", "model_tag": "microsoft/wavlm-base-sv"}],
+        cache_folder="/tmp/vc",
+    )
+    assert configs[0]["cache_dir"].endswith("huggingface")
+
+
+def test_cache_namespace_speaker_alias_hf_tag():
+    from versa.scorer_shared import configure_metric_cache_dirs
+
+    configs = configure_metric_cache_dirs(
+        [{"name": "spk_similarity", "model_tag": "microsoft/wavlm-base-sv"}],
+        cache_folder="/tmp/vc",
+    )
+    assert configs[0]["cache_dir"].endswith("huggingface")
+
+
+def test_cache_namespace_explicit_cache_dir_wins():
+    """Explicit cache_dir wins at the config-plumbing level.
+
+    Note: at model-load time, get_hf_cache_dir gives the VERSA_HF_CACHE_DIR
+    environment variable precedence over this per-metric cache_dir.
+    """
+    from versa.scorer_shared import configure_metric_cache_dirs
+
+    configs = configure_metric_cache_dirs(
+        [
+            {
+                "name": "speaker",
+                "model_tag": "microsoft/wavlm-base-sv",
+                "cache_dir": "/custom/cache",
+            }
+        ],
+        cache_folder="/tmp/vc",
+    )
+    assert configs[0]["cache_dir"] == "/custom/cache"
+
+
+def test_cache_namespace_backend_override():
+    from versa.scorer_shared import configure_metric_cache_dirs
+
+    configs = configure_metric_cache_dirs(
+        [{"name": "speaker", "model_tag": "default", "backend": "huggingface"}],
+        cache_folder="/tmp/vc",
+    )
+    assert configs[0]["cache_dir"].endswith("huggingface")
+
+
+def test_resolve_backend_none_tag_is_espnet():
+    assert resolve_speaker_backend(model_tag=None) == "espnet"
+
+
+def test_resolve_backend_non_string_tag_does_not_crash():
+    assert resolve_speaker_backend(model_tag=123) == "huggingface"
+
+
+def test_effective_dependencies_wavlm_excludes_espnet():
+    from versa.config_validation import _effective_dependencies
+    from versa.utterance_metrics.speaker import _speaker_metadata
+
+    deps = _effective_dependencies(
+        "speaker", _speaker_metadata(), {"model_tag": "microsoft/wavlm-base-sv"}
+    )
+    assert "espnet2" not in deps
+    assert "transformers" in deps
+
+
+def test_effective_dependencies_default_excludes_transformers():
+    from versa.config_validation import _effective_dependencies
+    from versa.utterance_metrics.speaker import _speaker_metadata
+
+    deps = _effective_dependencies(
+        "speaker", _speaker_metadata(), {"model_tag": "default"}
+    )
+    assert "transformers" not in deps
+    assert "espnet2" in deps

test/test_pipeline/test_speaker_wavlm.py

@@ -0,0 +1,71 @@
+import math
+import os
+
+import pytest
+import yaml
+
+from versa.definition import MetricRegistry
+from versa.scorer_shared import VersaScorer, compute_summary
+from versa.utils_shared import find_files
+from versa.utterance_metrics.speaker import (
+    is_transformers_available,
+    register_speaker_metric,
+)
+
+RUN_REAL_MODEL_TESTS = os.environ.get("VERSA_RUN_REAL_MODEL_TESTS") == "1"
+
+
+def _load_wavlm_speaker_config():
+    with open(
+        "egs/separate_metrics/spk_similarity_wavlm.yaml", "r", encoding="utf-8"
+    ) as f:
+        return yaml.safe_load(f)
+
+
+def _sample_files():
+    gen_path = "test/test_samples/test2"
+    gt_path = "test/test_samples/test1"
+    if not os.path.isdir(gen_path) or not os.path.isdir(gt_path):
+        pytest.skip("Required test sample directories are not available")
+    return find_files(gen_path), find_files(gt_path)
+
+
+@pytest.mark.real_model
+@pytest.mark.skipif(
+    not RUN_REAL_MODEL_TESTS,
+    reason="Set VERSA_RUN_REAL_MODEL_TESTS=1 to run real model-backed checks",
+)
+@pytest.mark.skipif(
+    not is_transformers_available(), reason="Transformers not available"
+)
+def test_speaker_wavlm_pipeline_with_real_model():
+    """Run the WavLM speaker backend through the registry/scorer path."""
+    gen_files, gt_files = _sample_files()
+    score_config = _load_wavlm_speaker_config()
+
+    registry = MetricRegistry()
+    register_speaker_metric(registry)
+    scorer = VersaScorer(registry)
+    metric_suite = scorer.load_metrics(score_config, use_gt=True, use_gpu=False)
+
+    assert len(score_config) > 0, "no scoring function is provided"
+    assert set(metric_suite.metrics) == {"speaker"}
+
+    score_info = scorer.score_utterances(
+        gen_files,
+        metric_suite,
+        gt_files=gt_files,
+        output_file=None,
+        io="soundfile",
+    )
+    summary = compute_summary(score_info)
+
+    assert "spk_similarity" in summary
+    assert math.isfinite(summary["spk_similarity"])
+    assert -1.0 <= summary["spk_similarity"] <= 1.0
+
+
+if __name__ == "__main__":
+    if not RUN_REAL_MODEL_TESTS:
+        raise SystemExit("Set VERSA_RUN_REAL_MODEL_TESTS=1 to run this check")
+    pytest.main([__file__, "-q", "-s"])

versa/config_validation.py

@@ -20,6 +20,30 @@
 _GPU_REQUIRED_METRICS = {"audiobox_aesthetics"}
 
 
+def _effective_dependencies(metric_name, metadata, config):
+    """Return the dependencies actually required by this metric config.
+
+    The speaker metric bundles both backends in its metadata; only the
+    backend selected by the config is required at runtime.
+    """
+    name = str(metric_name).lower()
+    if name in {"speaker", "spk_similarity", "speaker_similarity"}:
+        from versa.utterance_metrics.speaker import resolve_speaker_backend
+
+        try:
+            backend = resolve_speaker_backend(
+                model_tag=config.get("model_tag", "default"),
+                backend=config.get("backend"),
+                model_path=config.get("model_path"),
+                model_config=config.get("model_config"),
+            )
+        except ValueError:
+            return list(metadata.dependencies)
+        excluded = "espnet2" if backend == "huggingface" else "transformers"
+        return [dep for dep in metadata.dependencies if dep != excluded]
+    return list(metadata.dependencies)
+
+
 @dataclass(frozen=True)
 class ScoreConfigValidationError:
     """One score configuration validation issue."""
@@ -142,7 +166,7 @@ def validate_score_config(
 
         missing_dependencies = [
             dependency
-            for dependency in metadata.dependencies
+            for dependency in _effective_dependencies(metric_name, metadata, config)
             if not _dependency_available(dependency)
         ]
         if missing_dependencies: