wehub-resource-sync
297 lines
10 KiB
Python
297 lines
10 KiB
Python
"""
|
|
Do speech recognition on a long audio file and compare the result with the expected transcript
|
|
"""
|
|
|
|
import asyncio
|
|
import math
|
|
import time
|
|
from collections.abc import Callable
|
|
|
|
import pytest
|
|
from dotenv import load_dotenv
|
|
|
|
from livekit import agents, rtc
|
|
from livekit.agents import inference, stt
|
|
from livekit.agents.stt.stt import STT, RecognizeStream, SpeechData, SpeechEvent
|
|
from livekit.plugins import (
|
|
assemblyai,
|
|
aws,
|
|
azure,
|
|
cartesia,
|
|
deepgram,
|
|
elevenlabs,
|
|
fal,
|
|
fireworksai,
|
|
gladia,
|
|
google,
|
|
gradium,
|
|
mistralai,
|
|
nvidia,
|
|
openai,
|
|
sarvam,
|
|
soniox,
|
|
speechmatics,
|
|
)
|
|
|
|
from .utils import make_test_speech, wer
|
|
|
|
pytestmark = pytest.mark.stt
|
|
|
|
SAMPLE_RATE = 24000
|
|
WER_THRESHOLD = 0.25
|
|
MAX_RETRIES = 2
|
|
|
|
|
|
def parameter_factory(plugin):
|
|
return pytest.param(lambda: plugin.STT(), id=plugin.__name__)
|
|
|
|
|
|
STTs: list[Callable[[], stt.STT]] = [
|
|
parameter_factory(plugin)
|
|
for plugin in [
|
|
deepgram,
|
|
assemblyai,
|
|
speechmatics,
|
|
elevenlabs,
|
|
fireworksai,
|
|
gladia,
|
|
fal,
|
|
mistralai,
|
|
nvidia,
|
|
openai,
|
|
cartesia,
|
|
soniox,
|
|
google,
|
|
inference,
|
|
azure,
|
|
aws,
|
|
sarvam,
|
|
# rtzr,
|
|
# TODO: only Business account allowed outside South Korea
|
|
# clova,
|
|
# TODO: https://github.com/spi-tch/spitch-python/issues/162
|
|
# spitch,
|
|
]
|
|
] + [
|
|
pytest.param(lambda: cartesia.STT(model="ink-whisper"), id="livekit.plugins.cartesia._legacy"),
|
|
pytest.param(lambda: deepgram.STTv2(), id="livekit.plugins.deepgram.STTv2"),
|
|
pytest.param(
|
|
lambda: gradium.STT(model_endpoint="wss://us.api.gradium.ai/api/speech/asr"),
|
|
id="livekit.plugins.gradium.STT",
|
|
),
|
|
]
|
|
|
|
# entries whose recognize() path is identical to an existing STTs entry, so they only
|
|
# add value to test_stream (openai realtime shares the REST path with openai.STT())
|
|
STREAM_ONLY_STTs: list[Callable[[], stt.STT]] = [
|
|
pytest.param(lambda: openai.STT(use_realtime=True), id="livekit.plugins.openai.realtime"),
|
|
]
|
|
|
|
|
|
@pytest.fixture(scope="session", autouse=True)
|
|
def load_env():
|
|
load_dotenv()
|
|
|
|
|
|
async def batch_recognize(
|
|
stt: stt.STT, frames: list[rtc.AudioFrame], n_batches: int = 1
|
|
) -> SpeechEvent:
|
|
if n_batches == 1:
|
|
return await stt.recognize(buffer=frames)
|
|
if n_batches > len(frames):
|
|
raise ValueError("n_batches must be less than or equal to the number of frames")
|
|
|
|
batch_size: int = len(frames) // n_batches
|
|
events: list[SpeechEvent] = []
|
|
for i in range(n_batches):
|
|
batch = frames[i * batch_size : (i + 1) * batch_size]
|
|
events.append(await stt.recognize(buffer=batch))
|
|
|
|
assert len(events) > 0
|
|
return SpeechEvent(
|
|
type=agents.stt.SpeechEventType.FINAL_TRANSCRIPT,
|
|
request_id=events[0].request_id,
|
|
alternatives=[
|
|
SpeechData(
|
|
text=" ".join(
|
|
[
|
|
event.alternatives[0].text
|
|
for event in events
|
|
if event.alternatives[0].text is not None
|
|
]
|
|
),
|
|
language=events[0].alternatives[0].language,
|
|
)
|
|
],
|
|
)
|
|
|
|
|
|
@pytest.mark.usefixtures("job_process")
|
|
@pytest.mark.parametrize("stt_factory", STTs)
|
|
async def test_recognize(stt_factory: Callable[[], stt.STT], request):
|
|
plugin_id = request.node.callspec.id.split("-")[0] # e.g., "livekit.plugins.deepgram"
|
|
sample_rate = SAMPLE_RATE
|
|
frames, transcript, duration = await make_test_speech(sample_rate=sample_rate)
|
|
|
|
# TODO: differentiate missing key vs other errors
|
|
try:
|
|
stt_instance = stt_factory()
|
|
except ValueError as e:
|
|
pytest.skip(f"{plugin_id}: {e}")
|
|
|
|
async with stt_instance as stt:
|
|
label = f"{stt.model}@{stt.provider}"
|
|
if not stt.capabilities.offline_recognize:
|
|
pytest.skip(f"{label} does not support batch recognition")
|
|
|
|
for attempt in range(2):
|
|
try:
|
|
start_time = time.time()
|
|
|
|
# WARN: Sarvam only supports <30s audio chunks
|
|
if stt.provider == "Sarvam" and duration > 30:
|
|
frames, *_ = await make_test_speech(
|
|
sample_rate=sample_rate, chunk_duration_ms=5 * 1000
|
|
)
|
|
n_batches = math.ceil(duration / 30)
|
|
else:
|
|
n_batches = 1
|
|
event = await batch_recognize(stt, frames, n_batches)
|
|
text = event.alternatives[0].text
|
|
dt = time.time() - start_time
|
|
|
|
print(f"WER: {wer(text, transcript)} for {stt} in {dt:.2f}s")
|
|
|
|
# Relaxed WER threshold for some providers
|
|
if stt.provider in {
|
|
"Gladia",
|
|
}:
|
|
assert len(text) > 0 and wer(text, transcript) <= 1.0
|
|
else:
|
|
assert wer(text, transcript) <= WER_THRESHOLD
|
|
assert event.type == agents.stt.SpeechEventType.FINAL_TRANSCRIPT
|
|
return
|
|
except (AssertionError, Exception):
|
|
if attempt < MAX_RETRIES - 1:
|
|
print(f"Attempt {attempt + 1} failed for {label}, retrying...")
|
|
continue
|
|
else:
|
|
raise
|
|
|
|
|
|
@pytest.mark.usefixtures("job_process")
|
|
@pytest.mark.parametrize("stt_factory", STTs + STREAM_ONLY_STTs)
|
|
async def test_stream(stt_factory: Callable[[], STT], request):
|
|
sample_rate = SAMPLE_RATE
|
|
plugin_id = request.node.callspec.id.split("-")[0]
|
|
frames, transcript, _ = await make_test_speech(chunk_duration_ms=10, sample_rate=sample_rate)
|
|
|
|
# TODO: differentiate missing key vs other errors
|
|
try:
|
|
stt_instance: STT = stt_factory()
|
|
except ValueError as e:
|
|
pytest.skip(f"{plugin_id}: {e}")
|
|
|
|
async with stt_instance as stt:
|
|
label = f"{stt.model}@{stt.provider}"
|
|
if not stt.capabilities.streaming:
|
|
pytest.skip(f"{label} does not support streaming")
|
|
|
|
for attempt in range(MAX_RETRIES):
|
|
try:
|
|
state = {"closing": False}
|
|
|
|
async def _stream_input(
|
|
frames: list[rtc.AudioFrame], stream: RecognizeStream, state: dict = state
|
|
):
|
|
for frame in frames:
|
|
stream.push_frame(frame)
|
|
await asyncio.sleep(0.005)
|
|
|
|
stream.end_input()
|
|
state["closing"] = True
|
|
|
|
async def _stream_output(stream: RecognizeStream, state: dict = state):
|
|
text = ""
|
|
# make sure the events are sent in the right order
|
|
recv_start, recv_end = False, True
|
|
start_time = time.time()
|
|
got_final_transcript = False
|
|
sos_count, final_count = 0, 0
|
|
|
|
async for event in stream:
|
|
if event.type == agents.stt.SpeechEventType.START_OF_SPEECH:
|
|
assert recv_end, (
|
|
"START_OF_SPEECH recv but no END_OF_SPEECH has been sent before"
|
|
)
|
|
assert not recv_start
|
|
recv_end = False
|
|
recv_start = True
|
|
sos_count += 1
|
|
continue
|
|
|
|
if event.type == agents.stt.SpeechEventType.FINAL_TRANSCRIPT:
|
|
if text != "":
|
|
text += " "
|
|
text += event.alternatives[0].text
|
|
# ensure STT is tagging languages correctly
|
|
language = event.alternatives[0].language
|
|
if stt.provider not in {"FireworksAI", "RTZR", "livekit"}:
|
|
assert language is not None
|
|
assert language.lower().startswith("en")
|
|
got_final_transcript = True
|
|
final_count += 1
|
|
# Some providers don't send END_OF_SPEECH, break after final transcript
|
|
if state["closing"]:
|
|
break
|
|
|
|
if event.type == agents.stt.SpeechEventType.END_OF_SPEECH:
|
|
recv_start = False
|
|
recv_end = True
|
|
await asyncio.sleep(1)
|
|
# some providers emit END_OF_SPEECH before the segment's final transcript
|
|
if state["closing"] and final_count >= sos_count:
|
|
break
|
|
|
|
dt = time.time() - start_time
|
|
print(f"WER: {wer(text, transcript)} for streamed {stt} in {dt:.2f}s")
|
|
# Relaxed WER threshold for some providers
|
|
if stt.provider in {
|
|
"RTZR", # RTZR defaults to Korean
|
|
"Deepgram",
|
|
"Sarvam",
|
|
"FireworksAI",
|
|
}:
|
|
assert len(text) > 0 and wer(text, transcript) <= 1.0
|
|
else:
|
|
assert got_final_transcript, "No FINAL_TRANSCRIPT received"
|
|
assert wer(text, transcript) <= WER_THRESHOLD
|
|
|
|
timed_out = False
|
|
|
|
async def _run_test():
|
|
nonlocal timed_out
|
|
stream = None
|
|
try:
|
|
async with asyncio.timeout(60):
|
|
stream = stt.stream()
|
|
await asyncio.gather(
|
|
_stream_input(frames, stream), _stream_output(stream)
|
|
)
|
|
except TimeoutError:
|
|
timed_out = True
|
|
finally:
|
|
if stream is not None:
|
|
await stream.aclose()
|
|
|
|
await _run_test()
|
|
if timed_out:
|
|
pytest.fail(f"{label} streaming timed out after 60 seconds")
|
|
return
|
|
except (AssertionError, Exception):
|
|
if attempt < MAX_RETRIES - 1:
|
|
print(f"Attempt {attempt + 1} failed for {label}, retrying...")
|
|
continue
|
|
else:
|
|
raise
|