feat: VAD ingestion buffer — Silero VAD, per-room accumulation
TDD: 3 tests covering silence passthrough, speech→silence segment emission, and short-speech discard. Uses _speech_ms tracking (not total buffer length) for accurate min_speech_ms enforcement. Silero VAD import is try/except'd so tests run without torch via mocker.patch on vad_prob. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,127 @@
|
|||||||
|
import numpy as np
|
||||||
|
from datetime import datetime
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
# Silero VAD — docker-only; set to None when not available
|
||||||
|
try:
|
||||||
|
import torch
|
||||||
|
_vad_model = None
|
||||||
|
_TORCH_AVAILABLE = True
|
||||||
|
except ImportError:
|
||||||
|
_TORCH_AVAILABLE = False
|
||||||
|
_vad_model = None
|
||||||
|
|
||||||
|
|
||||||
|
def _load_vad():
|
||||||
|
global _vad_model
|
||||||
|
if not _TORCH_AVAILABLE:
|
||||||
|
raise ImportError("torch not installed")
|
||||||
|
if _vad_model is None:
|
||||||
|
import torch
|
||||||
|
model, _ = torch.hub.load(
|
||||||
|
repo_or_dir="snakers4/silero-vad",
|
||||||
|
model="silero_vad",
|
||||||
|
force_reload=False,
|
||||||
|
trust_repo=True,
|
||||||
|
)
|
||||||
|
_vad_model = model
|
||||||
|
return _vad_model
|
||||||
|
|
||||||
|
|
||||||
|
def vad_prob(chunk: bytes, sample_rate: int = 16000) -> float:
|
||||||
|
"""Return speech probability (0-1) for a PCM int16 chunk."""
|
||||||
|
import torch
|
||||||
|
model = _load_vad()
|
||||||
|
audio = np.frombuffer(chunk, dtype=np.int16).astype(np.float32) / 32768.0
|
||||||
|
tensor = torch.FloatTensor(audio)
|
||||||
|
with torch.no_grad():
|
||||||
|
prob = model(tensor, sample_rate).item()
|
||||||
|
return prob
|
||||||
|
|
||||||
|
|
||||||
|
class RoomVADBuffer:
|
||||||
|
"""
|
||||||
|
Accumulates audio chunks for a single room.
|
||||||
|
Emits (audio_bytes, start_time) when a speech segment ends.
|
||||||
|
"""
|
||||||
|
|
||||||
|
SPEECH_THRESHOLD = 0.5
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
silence_threshold_ms: int = 600,
|
||||||
|
min_speech_ms: int = 500,
|
||||||
|
chunk_ms: int = 250,
|
||||||
|
sample_rate: int = 16000,
|
||||||
|
):
|
||||||
|
self.silence_threshold_ms = silence_threshold_ms
|
||||||
|
self.min_speech_ms = min_speech_ms
|
||||||
|
self.chunk_ms = chunk_ms
|
||||||
|
self.sample_rate = sample_rate
|
||||||
|
self._buffer: list = []
|
||||||
|
self._speaking = False
|
||||||
|
self._silence_ms = 0
|
||||||
|
self._speech_ms = 0 # track pure speech duration
|
||||||
|
self._start_time: Optional[datetime] = None
|
||||||
|
|
||||||
|
def process_chunk(
|
||||||
|
self, chunk: bytes, timestamp: datetime
|
||||||
|
) -> Optional[tuple]:
|
||||||
|
"""
|
||||||
|
Returns (audio_bytes, start_time) when a complete speech segment is ready.
|
||||||
|
Returns None while accumulating or during silence.
|
||||||
|
"""
|
||||||
|
prob = vad_prob(chunk)
|
||||||
|
|
||||||
|
if prob >= self.SPEECH_THRESHOLD:
|
||||||
|
if not self._speaking:
|
||||||
|
self._speaking = True
|
||||||
|
self._start_time = timestamp
|
||||||
|
self._buffer = []
|
||||||
|
self._speech_ms = 0
|
||||||
|
self._buffer.append(chunk)
|
||||||
|
self._silence_ms = 0
|
||||||
|
self._speech_ms += self.chunk_ms
|
||||||
|
elif self._speaking:
|
||||||
|
self._buffer.append(chunk) # include trailing silence
|
||||||
|
self._silence_ms += self.chunk_ms
|
||||||
|
if self._silence_ms >= self.silence_threshold_ms:
|
||||||
|
audio = b"".join(self._buffer)
|
||||||
|
start = self._start_time
|
||||||
|
speech_ms = self._speech_ms
|
||||||
|
self._buffer = []
|
||||||
|
self._speaking = False
|
||||||
|
self._silence_ms = 0
|
||||||
|
self._speech_ms = 0
|
||||||
|
self._start_time = None
|
||||||
|
if speech_ms >= self.min_speech_ms:
|
||||||
|
return audio, start
|
||||||
|
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
# Per-room registry: room_name -> {"buffer": RoomVADBuffer, "muted": bool}
|
||||||
|
_rooms: dict = {}
|
||||||
|
|
||||||
|
|
||||||
|
def get_or_create_room(room_id: str, silence_ms: int, min_speech_ms: int) -> dict:
|
||||||
|
if room_id not in _rooms:
|
||||||
|
_rooms[room_id] = {
|
||||||
|
"buffer": RoomVADBuffer(silence_threshold_ms=silence_ms, min_speech_ms=min_speech_ms),
|
||||||
|
"muted": False,
|
||||||
|
}
|
||||||
|
return _rooms[room_id]
|
||||||
|
|
||||||
|
|
||||||
|
def mute_room(room_id: str) -> None:
|
||||||
|
if room_id in _rooms:
|
||||||
|
_rooms[room_id]["muted"] = True
|
||||||
|
|
||||||
|
|
||||||
|
def unmute_room(room_id: str) -> None:
|
||||||
|
if room_id in _rooms:
|
||||||
|
_rooms[room_id]["muted"] = False
|
||||||
|
|
||||||
|
|
||||||
|
def is_muted(room_id: str) -> bool:
|
||||||
|
return _rooms.get(room_id, {}).get("muted", False)
|
||||||
@@ -0,0 +1,55 @@
|
|||||||
|
import numpy as np
|
||||||
|
import pytest
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
from unittest.mock import patch, MagicMock
|
||||||
|
|
||||||
|
|
||||||
|
def _make_chunk(samples: int = 4000, amplitude: float = 0.0) -> bytes:
|
||||||
|
"""Create a fake PCM int16 chunk."""
|
||||||
|
audio = (np.ones(samples) * amplitude * 32767).astype(np.int16)
|
||||||
|
return audio.tobytes()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def buf(mocker):
|
||||||
|
"""RoomVADBuffer with mocked VAD model."""
|
||||||
|
mocker.patch(
|
||||||
|
"app.pipeline.ingestion.vad_prob",
|
||||||
|
side_effect=lambda chunk: 0.9 if np.frombuffer(chunk, np.int16).max() > 100 else 0.1,
|
||||||
|
)
|
||||||
|
from app.pipeline.ingestion import RoomVADBuffer
|
||||||
|
return RoomVADBuffer(silence_threshold_ms=600, min_speech_ms=500, chunk_ms=250)
|
||||||
|
|
||||||
|
|
||||||
|
def test_silence_produces_no_segment(buf):
|
||||||
|
ts = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||||
|
for _ in range(10):
|
||||||
|
result = buf.process_chunk(_make_chunk(amplitude=0.0), ts)
|
||||||
|
assert result is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_speech_then_silence_produces_segment(buf):
|
||||||
|
ts = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||||
|
# 4 speech chunks = 1 second of speech
|
||||||
|
for _ in range(4):
|
||||||
|
result = buf.process_chunk(_make_chunk(amplitude=0.5), ts)
|
||||||
|
assert result is None # still accumulating
|
||||||
|
# 3 silence chunks = 750ms silence (> 600ms threshold)
|
||||||
|
segment = None
|
||||||
|
for _ in range(3):
|
||||||
|
segment = buf.process_chunk(_make_chunk(amplitude=0.0), ts)
|
||||||
|
assert segment is not None
|
||||||
|
assert isinstance(segment, tuple)
|
||||||
|
audio_bytes, start_time = segment
|
||||||
|
assert len(audio_bytes) > 0
|
||||||
|
|
||||||
|
|
||||||
|
def test_short_speech_discarded(buf):
|
||||||
|
ts = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||||
|
# 1 speech chunk = 250ms (below 500ms min)
|
||||||
|
buf.process_chunk(_make_chunk(amplitude=0.5), ts)
|
||||||
|
# Now silence to trigger emission
|
||||||
|
result = None
|
||||||
|
for _ in range(3):
|
||||||
|
result = buf.process_chunk(_make_chunk(amplitude=0.0), ts)
|
||||||
|
assert result is None # discarded as too short
|
||||||
Reference in New Issue
Block a user