Skip to content

Commit f67d27b

Browse files
satraclaude
satra
and
claude
authored
Expand speech model coverage: S3PRL, NeMo ASR, Pyannote VAD, model registry (#491)
* Expand speech model coverage: S3PRL, NeMo ASR, Pyannote VAD, model registry New backends: - S3PRL subprocess venv for SSL embeddings (APC, TERA, CPC, MockingJay, DeCoAR2 + more) — pins torch<2.5 for torchaudio compatibility - SpeechBrain routing in ssl_embeddings (ECAPA-TDNN, x-vector, ResNet via existing speaker_embeddings infrastructure) - NeMo ASR subprocess venv (Conformer CTC/RNNT via existing NeMo venv) - Pyannote dedicated VAD (distinct from diarization-based approach) Unified SSL API: - extract_ssl_embeddings_from_audios() routes by model type: HFModel → HuggingFace, SpeechBrainModel → SpeechBrain, str → S3PRL Model registry: - YAML data source with 30+ models across all tasks - Generated Markdown table organized by task Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Fix formatting (ruff format + end of files) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Fix review: squeeze with dim, S3PRL sampling rate check - SpeechBrain SSL: squeeze(dim=1) and squeeze(dim=0) to avoid collapsing batch dimension on single-audio input - S3PRL: validate 16kHz sampling rate before subprocess call Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 48c19b5 commit f67d27b

19 files changed

Lines changed: 1816 additions & 57 deletions

File tree

‎scripts/generate_model_registry.py‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
#!/usr/bin/env python3
2+
"""Generate model_registry.md from model_registry.yaml."""
3+
4+
from pathlib import Path
5+
6+
import yaml
7+
8+
9+
def main() -> None:
10+
"""Read YAML and output Markdown table."""
11+
registry_path = Path(__file__).parent.parent / "docs" / "model_registry.yaml"
12+
with open(registry_path) as f:
13+
models = yaml.safe_load(f)
14+
15+
# Group by task
16+
tasks: dict = {}
17+
for model in models:
18+
task = model["task"]
19+
if task not in tasks:
20+
tasks[task] = []
21+
tasks[task].append(model)
22+
23+
print("# Senselab Model Registry\n")
24+
print("All models supported by senselab, organized by task.\n")
25+
26+
for task, task_models in tasks.items():
27+
title = task.replace("_", " ").title()
28+
print(f"## {title}\n")
29+
print("| Model | Source | Model ID | Embedding Dim | Parameters | Recommended For |")
30+
print("|-------|--------|----------|---------------|------------|-----------------|")
31+
for m in task_models:
32+
name = m["name"]
33+
source = m["source"]
34+
model_id = f'`{m["model_id"]}`'
35+
emb = m.get("embedding_dim", "—")
36+
params = m.get("parameters", "—")
37+
rec = m.get("recommended_for", "—")
38+
print(f"| {name} | {source} | {model_id} | {emb} | {params} | {rec} |")
39+
print()
40+
41+
42+
if __name__ == "__main__":
43+
main()

‎specs/20260428-101838-expand-speech-models/checklists/requirements.md‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
# Specification Quality Checklist: Expand Speech Representation Model Coverage
2+
3+
**Purpose**: Validate specification completeness and quality before proceeding to planning
4+
**Created**: 2026-04-28
5+
**Feature**: [spec.md](../spec.md)
6+
7+
## Content Quality
8+
9+
- [x] No implementation details (languages, frameworks, APIs)
10+
- [x] Focused on user value and business needs
11+
- [x] Written for non-technical stakeholders
12+
- [x] All mandatory sections completed
13+
14+
## Requirement Completeness
15+
16+
- [x] No [NEEDS CLARIFICATION] markers remain
17+
- [x] Requirements are testable and unambiguous
18+
- [x] Success criteria are measurable
19+
- [x] Success criteria are technology-agnostic (no implementation details)
20+
- [x] All acceptance scenarios are defined
21+
- [x] Edge cases are identified
22+
- [x] Scope is clearly bounded
23+
- [x] Dependencies and assumptions identified
24+
25+
## Feature Readiness
26+
27+
- [x] All functional requirements have clear acceptance criteria
28+
- [x] User scenarios cover primary flows
29+
- [x] Feature meets measurable outcomes defined in Success Criteria
30+
- [x] No implementation details leak into specification
31+
32+
## Notes
33+
34+
- All items pass validation. Spec is ready for `/speckit.plan`.
35+
- Speaker identity coding paper cloned at /tmp/speaker_identity_coding_paper for reference
36+
- Paper uses 17 models across S3PRL (7), HuggingFace (5), SpeechBrain (3), SERAB BYOL-S (1), cochlear (1)
37+
- S3PRL models will need subprocess venv (dependency conflicts with main env)

‎specs/20260428-101838-expand-speech-models/data-model.md‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
# Data Model: Expand Speech Representation Model Coverage
2+
3+
**Date**: 2026-04-28
4+
5+
## New Entities
6+
7+
### S3PRLModel (new model type)
8+
- Extends SenselabModel pattern
9+
- `path_or_uri`: S3PRL model name (e.g., "apc", "tera", "cpc")
10+
- Runs in isolated subprocess venv
11+
- Returns embeddings as List[torch.Tensor]
12+
13+
### ModelRegistryEntry (documentation entity)
14+
- `name`: Human-readable model name
15+
- `task`: Which senselab task it applies to
16+
- `source`: Backend (huggingface, speechbrain, s3prl, nemo, pyannote)
17+
- `model_id`: Identifier used to load the model
18+
- `embedding_dim`: Output dimension
19+
- `parameters`: Parameter count
20+
- `training_data`: What it was trained on
21+
- `recommended_for`: Use case guidance
22+
23+
## Modified Entities
24+
25+
### ssl_embeddings module
26+
- Currently: HuggingFace-only
27+
- After: HuggingFace + S3PRL + SpeechBrain backends
28+
- API signature stays the same; backend selected by model type
29+
30+
### speech_to_text module
31+
- Currently: HuggingFace-only
32+
- After: HuggingFace + NeMo (subprocess venv)
33+
34+
### voice_activity_detection module
35+
- Currently: Reuses diarization output
36+
- After: Also supports dedicated Pyannote VAD pipeline

‎specs/20260428-101838-expand-speech-models/plan.md‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
# Implementation Plan: Expand Speech Representation Model Coverage
2+
3+
**Branch**: `20260428-101838-expand-speech-models` | **Date**: 2026-04-28 | **Spec**: [spec.md](spec.md)
4+
5+
## Summary
6+
7+
Expand senselab's speech representation model coverage by integrating S3PRL (30+ SSL models), broadening SpeechBrain access, adding NeMo ASR, enabling dedicated Pyannote VAD, creating a model registry, and providing a tutorial reproducing the speaker identity coding paper's benchmarking pipeline.
8+
9+
## Technical Context
10+
11+
**Language/Version**: Python 3.11-3.12
12+
**Primary Dependencies**: s3prl (subprocess venv), speechbrain, pyannote-audio, nemo_toolkit (subprocess venv), transformers
13+
**Storage**: N/A
14+
**Testing**: pytest for unit tests, papermill for tutorials
15+
**Target Platform**: Linux (CI), macOS (dev), Google Colab (tutorials)
16+
**Project Type**: Library extension + tutorials + documentation
17+
18+
## Constitution Check
19+
20+
| Principle | Status | Notes |
21+
|-----------|--------|-------|
22+
| I. UV-Managed Python | PASS | All execution via uv; subprocess venvs for S3PRL/NeMo |
23+
| II. Encapsulated Testing | PASS | Tests in uv-managed venv |
24+
| III. Commit Early and Often | PASS | Incremental per-phase commits |
25+
| IV. CI Must Stay Green | PASS | New tests + existing tests must pass |
26+
| V. Memory-Driven Anti-Pattern Avoidance | PASS | Follows subprocess venv pattern from prior work |
27+
| VI. No Unnecessary API Calls | PASS | Models cached after first download |
28+
| VII. Simplicity First | PASS | Uses existing patterns (subprocess venv, model classes) |
29+
| VIII. No Hardcoded Parameters | PASS | Model names, venv paths configurable |
30+
31+
## Project Structure
32+
33+
```text
34+
src/senselab/audio/tasks/
35+
├── ssl_embeddings/
36+
│ ├── self_supervised_features.py # UPDATED (add S3PRL + SpeechBrain backends)
37+
│ ├── s3prl.py # NEW (S3PRL subprocess venv worker)
38+
│ └── doc.md # UPDATED (comprehensive model list)
39+
├── speech_to_text/
40+
│ ├── nemo.py # NEW (NeMo ASR subprocess venv)
41+
│ └── api.py # UPDATED (NeMo backend option)
42+
├── voice_activity_detection/
43+
│ ├── pyannote_vad.py # NEW (dedicated Pyannote VAD)
44+
│ └── api.py # UPDATED (VAD model option)
45+
46+
tutorials/audio/
47+
├── ssl_embeddings_comparison.ipynb # NEW (multi-model embedding comparison)
48+
├── speaker_identity_benchmark.ipynb # NEW (reproduce paper's 17-model benchmark)
49+
50+
docs/
51+
└── model_registry.md # NEW (or generated page)
52+
```
53+
54+
## Implementation Phases
55+
56+
### Phase 1: S3PRL Subprocess Venv Integration (US1 — P1)
57+
58+
**Goal**: Add S3PRL models as a backend for SSL embedding extraction via subprocess venv.
59+
60+
1. Create `s3prl.py` with subprocess worker script that:
61+
- Loads model via `s3prl.hub`
62+
- Accepts audio as FLAC files
63+
- Extracts hidden states / embeddings
64+
- Returns as numpy arrays via JSON
65+
2. Create `S3PRLModel` or use string-based model selection
66+
3. Update `ssl_embeddings` API to route S3PRL models to subprocess backend
67+
4. Write tests for at least 3 S3PRL models (APC, TERA, CPC)
68+
5. Update `ssl_embeddings/doc.md` with S3PRL model table
69+
70+
### Phase 2: SpeechBrain Embedding Unification (US1 — P1)
71+
72+
**Goal**: Make SpeechBrain speaker encoders accessible through ssl_embeddings in addition to speaker_embeddings.
73+
74+
1. Update `ssl_embeddings` API to accept SpeechBrainModel and route to existing speaker encoder infrastructure
75+
2. This is primarily an API routing change — the actual SpeechBrain encode_batch() already works
76+
3. Document in ssl_embeddings/doc.md
77+
78+
### Phase 3: NeMo ASR (US3 — P2)
79+
80+
**Goal**: Add NeMo Conformer ASR as a subprocess venv option for speech-to-text.
81+
82+
1. Create `nemo.py` in speech_to_text with subprocess worker script
83+
2. Reuse existing NeMo subprocess venv (add nemo ASR dependencies if missing)
84+
3. Update `transcribe_audios` API to accept NeMo models
85+
4. Test with `nvidia/stt_en_conformer_ctc_large`
86+
87+
### Phase 4: Pyannote Dedicated VAD (US3 — P2)
88+
89+
**Goal**: Expose Pyannote's VAD pipeline as a distinct option.
90+
91+
1. Create `pyannote_vad.py` with dedicated VAD pipeline (not diarization-based)
92+
2. Update VAD API to support model selection
93+
3. Test with `pyannote/voice-activity-detection`
94+
95+
### Phase 5: Model Registry + Documentation (US4 — P2)
96+
97+
**Goal**: Create a comprehensive model registry and update all docs.
98+
99+
1. Create YAML data source with all supported models
100+
2. Generate model_registry.md from YAML
101+
3. Update ssl_embeddings/doc.md with complete model table
102+
4. Update README.md features section if needed
103+
104+
### Phase 6: Expand Speaker Embeddings Tutorial (US2 — P2)
105+
106+
**Goal**: Expand existing `extract_speaker_embeddings.ipynb` with multi-backend comparison exploring how different representations capture speaker identity.
107+
108+
1. Add S3PRL, HuggingFace SSL sections to existing tutorial
109+
2. Add comparison visualization (t-SNE, cosine similarity)
110+
3. Pedagogical explanations of supervised vs self-supervised vs generative approaches
111+
112+
### Phase 7: Polish
113+
114+
1. Run all tests locally
115+
2. Run pre-commit
116+
3. Push and verify CI
117+
4. Merge
118+
119+
## Risk Mitigation
120+
121+
| Risk | Mitigation |
122+
|------|-----------|
123+
| S3PRL venv creation slow | Cache venv; only create on first use |
124+
| S3PRL model download large | Tests use smallest models (APC < 5MB) |
125+
| NeMo ASR model large (~500MB) | CPU timeout generous; test with small model if available |
126+
| Some paper models unavailable | Document which are accessible (15/17 target) |

‎specs/20260428-101838-expand-speech-models/quickstart.md‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
# Quickstart: Expand Speech Representation Model Coverage
2+
3+
## Test S3PRL embeddings
4+
```bash
5+
uv run python -c "
6+
from senselab.audio.data_structures import Audio
7+
from senselab.audio.tasks.ssl_embeddings import extract_ssl_embeddings
8+
audio = Audio(filepath='tutorial_audio_files/audio_48khz_mono_16bits.wav')
9+
embeddings = extract_ssl_embeddings([audio], model='apc') # S3PRL model
10+
print(f'Shape: {embeddings[0].shape}')
11+
"
12+
```
13+
14+
## Test NeMo ASR
15+
```bash
16+
uv run python -c "
17+
from senselab.audio.data_structures import Audio
18+
from senselab.audio.tasks.speech_to_text.api import transcribe_audios
19+
from senselab.utils.data_structures import NeMoModel
20+
audio = Audio(filepath='tutorial_audio_files/audio_48khz_mono_16bits.wav')
21+
result = transcribe_audios([audio], model=NeMoModel(path_or_uri='nvidia/stt_en_conformer_ctc_large'))
22+
print(result[0].text)
23+
"
24+
```
25+
26+
## Build model registry
27+
```bash
28+
uv run python scripts/generate_model_registry.py > docs/model_registry.md
29+
```

‎specs/20260428-101838-expand-speech-models/research.md‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
# Research: Expand Speech Representation Model Coverage
2+
3+
**Date**: 2026-04-28
4+
5+
## R1: S3PRL Integration Approach
6+
7+
**Decision**: Add S3PRL as a subprocess venv backend for SSL embedding extraction. The S3PRL toolkit pins specific torch/torchaudio versions and has a large dependency tree that conflicts with the main environment.
8+
9+
**Rationale**: S3PRL provides 30+ pre-trained SSL models through a uniform `hub` API (`s3prl.hub`). The paper's encoder.py shows the pattern: `model = getattr(hub, model_name)()` → `model(wavs)` returns hidden states. This is a clean API that maps well to a subprocess worker script.
10+
11+
**S3PRL models to support** (from the paper + s3prl hub):
12+
- APC, TERA, MockingJay, DeCoAR2, CPC (from paper)
13+
- NPC, VQ-APC, Audio ALBERT (additional popular models)
14+
- Filterbanks, MFCCs, mel (feature extractors, not SSL but useful baselines)
15+
16+
**Venv requirements**: `s3prl`, `torch>=2.0`, `torchaudio`, `numpy`, `soundfile`
17+
18+
**Alternatives considered**:
19+
- Direct dependency: Too many conflicts (pins torch versions)
20+
- Docker container: Overkill for model inference
21+
- Skip S3PRL, use only HF equivalents: Many S3PRL models (APC, TERA, CPC) have no HF equivalent
22+
23+
## R2: SpeechBrain Speaker Encoder Access
24+
25+
**Decision**: Expose SpeechBrain's `EncoderClassifier.encode_batch()` for embedding extraction (not just verification). This is already partially done — the speaker_embeddings module uses it. The fix is to also expose it in ssl_embeddings with a clear "SpeechBrain speaker encoders" option.
26+
27+
**Rationale**: The paper uses 3 SpeechBrain models (ECAPA-TDNN, ResNet, x-vector) via `EncoderClassifier.from_hparams()` → `encode_batch()`. Senselab already has this in `speaker_embeddings/` but not unified with the SSL embeddings module.
28+
29+
**Action**: Unify the embedding extraction API so users can request embeddings from any backend (HF, S3PRL, SpeechBrain) through a single entry point, or keep them as separate but documented paths.
30+
31+
## R3: NeMo ASR via Subprocess Venv
32+
33+
**Decision**: Add NeMo ASR as a subprocess venv option in the speech_to_text module. NeMo has Conformer-CTC and Conformer-RNNT models that are state-of-the-art for English ASR.
34+
35+
**Rationale**: NeMo already has a subprocess venv pattern established for diarization. ASR can reuse the same venv with minimal additions.
36+
37+
**NeMo ASR models**: `nvidia/stt_en_conformer_ctc_large`, `nvidia/stt_en_conformer_transducer_large`
38+
39+
## R4: Pyannote Dedicated VAD
40+
41+
**Decision**: Expose `pyannote/voice-activity-detection` as a first-class option separate from the diarization-based VAD. The current VAD is just diarization with all segments merged.
42+
43+
**Rationale**: Pyannote's dedicated VAD model is faster and more accurate for voice detection than using full diarization. It's also a prerequisite for other tasks (trimming silence, detecting speech regions for processing).
44+
45+
## R5: Model Registry Documentation
46+
47+
**Decision**: Create a `docs/model_registry.md` (or `MODEL_REGISTRY.md` in repo root) that is a structured table of all supported models, generated from a YAML data source. This becomes part of the pdoc-generated docs.
48+
49+
**Format per model**:
50+
```yaml
51+
- name: ECAPA-TDNN
52+
task: speaker_embeddings
53+
source: speechbrain
54+
model_id: speechbrain/spkrec-ecapa-voxceleb
55+
embedding_dim: 192
56+
parameters: 7.3M
57+
training_data: VoxCeleb
58+
recommended_for: Speaker verification, speaker identification
59+
```
60+
61+
## R6: Speaker Identity Coding Paper Reproducibility
62+
63+
**Decision**: Create a tutorial notebook that demonstrates extracting embeddings from the paper's 17 models using senselab, then performs basic speaker verification benchmarking. The paper's custom train/benchmark scripts are NOT reproduced — only the embedding extraction and evaluation steps.
64+
65+
**Models accessible through senselab** (15 of 17):
66+
- HuggingFace (5/5): wav2vec2, W2V-BERT, HuBERT, WavLM, data2vec — all work via existing ssl_embeddings
67+
- SpeechBrain (3/3): ECAPA-TDNN, ResNet, x-vector — work via speaker_embeddings
68+
- S3PRL (5/7): APC, TERA, MockingJay, DeCoAR2, CPC — need new subprocess venv; filterbanks/MFCCs/mel are baseline features
69+
- Not accessible (2): Hybrid BYOL-S (custom SERAB package), cochlear model (custom implementation)

0 commit comments

Comments
 (0)