310 lines
11 KiB
Python
310 lines
11 KiB
Python
"""Release gates for Hermes' disposable tiny rolling transcription model."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import struct
|
|
import sys
|
|
import threading
|
|
import time
|
|
import wave
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
SERVER = ROOT / "dockerfiles" / "hermes-jetson-stt-server.py"
|
|
DOCKERFILE = ROOT / "dockerfiles" / "Dockerfile.hermes-jetson-stt"
|
|
PIPELINE = ROOT / "ci" / "Jenkinsfile.hermes-voice-image"
|
|
|
|
|
|
def _load_server(monkeypatch):
|
|
"""Import the server without requiring image-only CUDA dependencies."""
|
|
spec = importlib.util.spec_from_file_location("hermes_stt_dual_model", SERVER)
|
|
assert spec and spec.loader
|
|
module = importlib.util.module_from_spec(spec)
|
|
monkeypatch.setitem(sys.modules, "cgi", SimpleNamespace())
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"torch",
|
|
SimpleNamespace(
|
|
cuda=SimpleNamespace(is_available=lambda: False, synchronize=lambda: None)
|
|
),
|
|
)
|
|
monkeypatch.setitem(sys.modules, "whisper", SimpleNamespace())
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class _Connection:
|
|
"""Capture asynchronous stream responses."""
|
|
|
|
def __init__(self) -> None:
|
|
self.messages: list[dict] = []
|
|
self.condition = threading.Condition()
|
|
|
|
def send_json(self, payload: dict) -> None:
|
|
with self.condition:
|
|
self.messages.append(payload)
|
|
self.condition.notify_all()
|
|
|
|
def wait_for(self, message_type: str, timeout: float = 1.0) -> dict:
|
|
deadline = time.monotonic() + timeout
|
|
with self.condition:
|
|
while time.monotonic() < deadline:
|
|
for message in self.messages:
|
|
if message.get("type") == message_type:
|
|
return message
|
|
self.condition.wait(deadline - time.monotonic())
|
|
raise AssertionError(f"missing {message_type}: {self.messages}")
|
|
|
|
|
|
class _Model:
|
|
"""Return deterministic Whisper output while recording model ownership."""
|
|
|
|
def __init__(self, text: str) -> None:
|
|
self.text = text
|
|
self.calls: list[dict] = []
|
|
|
|
def transcribe(self, path: str, **options: object) -> dict:
|
|
self.calls.append({"path": path, "options": options})
|
|
return {
|
|
"text": self.text,
|
|
"language": "en",
|
|
"segments": [
|
|
{
|
|
"text": self.text,
|
|
"no_speech_prob": 0.01,
|
|
"avg_logprob": -0.1,
|
|
}
|
|
],
|
|
}
|
|
|
|
|
|
class _BlockingModel(_Model):
|
|
"""Hold one authoritative decode to exercise commit/speculation races."""
|
|
|
|
def __init__(self, text: str = "hello") -> None:
|
|
super().__init__(text)
|
|
self.entered = threading.Event()
|
|
self.release = threading.Event()
|
|
|
|
def transcribe(self, path: str, **options: object) -> dict:
|
|
self.entered.set()
|
|
assert self.release.wait(2.0)
|
|
return super().transcribe(path, **options)
|
|
|
|
|
|
class _SequencedModel(_Model):
|
|
"""Return evolving text and retain the audio size decoded each time."""
|
|
|
|
def __init__(self, texts: list[str]) -> None:
|
|
super().__init__(texts[0])
|
|
self.texts = texts
|
|
self.frame_counts: list[int] = []
|
|
|
|
def transcribe(self, path: str, **options: object) -> dict:
|
|
with wave.open(path, "rb") as wav_file:
|
|
self.frame_counts.append(wav_file.getnframes())
|
|
self.text = self.texts[min(len(self.calls), len(self.texts) - 1)]
|
|
return super().transcribe(path, **options)
|
|
|
|
|
|
def _start(session) -> None:
|
|
session.start(
|
|
{
|
|
"type": "start",
|
|
"turn_id": "dual-model-turn",
|
|
"format": "pcm_s16le",
|
|
"sample_rate": 16_000,
|
|
"language": "auto",
|
|
}
|
|
)
|
|
|
|
|
|
def _speech(samples: int = 1_600) -> bytes:
|
|
return struct.pack(f"<{samples}h", *([12_000] * samples))
|
|
|
|
|
|
def _silence(samples: int = 12_000) -> bytes:
|
|
return bytes(samples * 2)
|
|
|
|
|
|
def _wait_idle(session) -> None:
|
|
deadline = time.monotonic() + 1.0
|
|
with session._lock:
|
|
while session._rolling_inflight and time.monotonic() < deadline:
|
|
session._lock.wait(deadline - time.monotonic())
|
|
assert not session._rolling_inflight
|
|
|
|
|
|
def test_rolling_uses_tiny_once_while_final_remains_small(monkeypatch):
|
|
module = _load_server(monkeypatch)
|
|
monkeypatch.setattr(module, "ROLLING_MIN_AUDIO_MS", 100)
|
|
monkeypatch.setattr(module, "ROLLING_MAX_RESULT_LAG_MS", 5_000)
|
|
final_model = _Model("authoritative small")
|
|
rolling_model = _Model("provisional tiny")
|
|
connection = _Connection()
|
|
session = module.StreamingTranscription(connection, final_model, rolling_model)
|
|
_start(session)
|
|
|
|
session.append(_speech())
|
|
partial = connection.wait_for("partial")
|
|
assert partial["transcript"] == "provisional tiny"
|
|
assert partial["stable_transcript"] == ""
|
|
assert len(rolling_model.calls) == 1
|
|
assert final_model.calls == []
|
|
|
|
session.append(_speech())
|
|
_wait_idle(session)
|
|
assert len(rolling_model.calls) == 1
|
|
session.commit()
|
|
assert connection.wait_for("final")["transcript"] == "authoritative small"
|
|
assert len(final_model.calls) == 1
|
|
|
|
|
|
def test_trailing_silence_after_speculation_keeps_exact_cache_key(monkeypatch):
|
|
"""Recorder shutdown silence must not force a second Whisper-small decode."""
|
|
module = _load_server(monkeypatch)
|
|
connection = _Connection()
|
|
model = _Model("hello")
|
|
session = module.StreamingTranscription(connection, model)
|
|
_start(session)
|
|
|
|
session.append(_speech(samples=20_800))
|
|
session.speculate()
|
|
assert connection.wait_for("partial")["transcript"] == "hello"
|
|
session.append(_silence(samples=8_000))
|
|
session.commit()
|
|
|
|
assert connection.wait_for("final")["transcript"] == "hello"
|
|
assert len(model.calls) == 1
|
|
telemetry = module.STREAM_TELEMETRY.snapshot()
|
|
assert telemetry["speculative_cache_hits"] == 1
|
|
assert telemetry["direct_final_decodes"] == 0
|
|
assert telemetry["last_commit_cache_hit"] is True
|
|
assert telemetry["last_commit_audio_ms"] == 1_300.0
|
|
|
|
|
|
def test_commit_waits_for_one_inflight_exact_speculation_without_redecoding(
|
|
monkeypatch,
|
|
):
|
|
"""Exercise the browser timing race while Whisper-small is still running."""
|
|
module = _load_server(monkeypatch)
|
|
connection = _Connection()
|
|
model = _BlockingModel()
|
|
session = module.StreamingTranscription(connection, model)
|
|
_start(session)
|
|
|
|
session.append(_speech(samples=20_800))
|
|
session.speculate()
|
|
assert model.entered.wait(1.0)
|
|
session.append(_silence(samples=8_000))
|
|
committed = threading.Event()
|
|
worker = threading.Thread(target=lambda: (session.commit(), committed.set()))
|
|
worker.start()
|
|
time.sleep(0.02)
|
|
assert not committed.is_set()
|
|
model.release.set()
|
|
worker.join(timeout=1.0)
|
|
|
|
assert committed.is_set()
|
|
assert connection.wait_for("final")["transcript"] == "hello"
|
|
assert len(model.calls) == 1
|
|
telemetry = module.STREAM_TELEMETRY.snapshot()
|
|
assert telemetry["last_commit_cache_hit"] is True
|
|
assert telemetry["last_commit_pending_wait_ms"] >= 10
|
|
|
|
|
|
def test_resumed_speech_invalidates_frozen_eos_and_decodes_full_turn(monkeypatch):
|
|
"""Freezing EOS cannot discard an intentional barge-in continuation."""
|
|
module = _load_server(monkeypatch)
|
|
connection = _Connection()
|
|
model = _SequencedModel(["first thought", "complete thought"])
|
|
session = module.StreamingTranscription(connection, model)
|
|
_start(session)
|
|
|
|
session.append(_speech(samples=8_000))
|
|
session.speculate()
|
|
connection.wait_for("partial")
|
|
session.append(_silence(samples=4_000))
|
|
session.resume()
|
|
session.append(_speech(samples=4_000))
|
|
session.commit()
|
|
|
|
assert connection.wait_for("final")["transcript"] == "complete thought"
|
|
assert model.frame_counts == [8_000, 16_000]
|
|
telemetry = module.STREAM_TELEMETRY.snapshot()
|
|
assert telemetry["speculative_cache_hits"] == 0
|
|
assert telemetry["direct_final_decodes"] == 1
|
|
assert telemetry["last_commit_cache_hit"] is False
|
|
|
|
|
|
def test_eos_does_not_rearm_but_resumed_speech_does(monkeypatch):
|
|
module = _load_server(monkeypatch)
|
|
monkeypatch.setattr(module, "ROLLING_MIN_AUDIO_MS", 100)
|
|
monkeypatch.setattr(module, "ROLLING_MAX_RESULT_LAG_MS", 5_000)
|
|
final_model = _Model("final")
|
|
rolling_model = _Model("rolling")
|
|
session = module.StreamingTranscription(_Connection(), final_model, rolling_model)
|
|
_start(session)
|
|
|
|
session.append(_speech())
|
|
_wait_idle(session)
|
|
assert session._rolling_attempted is True
|
|
session.append(bytes(24_000))
|
|
assert session._at_eos is True
|
|
assert session._rolling_attempted is True
|
|
|
|
session.resume()
|
|
assert session._rolling_attempted is False
|
|
session.append(_speech())
|
|
_wait_idle(session)
|
|
assert len(rolling_model.calls) == 2
|
|
|
|
|
|
def test_dual_warmup_and_separate_timing_are_declared(monkeypatch):
|
|
module = _load_server(monkeypatch)
|
|
model = _Model("warm")
|
|
synchronized: list[bool] = []
|
|
monkeypatch.setattr(
|
|
module.torch.cuda, "synchronize", lambda: synchronized.append(True)
|
|
)
|
|
|
|
elapsed = module._warm_model(model)
|
|
|
|
assert elapsed >= 0
|
|
assert module.ROLLING_MIN_AUDIO_MS == 1_800
|
|
assert len(model.calls) == 1
|
|
assert model.calls[0]["options"]["language"] is None
|
|
assert synchronized == [True]
|
|
assert not Path(model.calls[0]["path"]).exists()
|
|
source = SERVER.read_text()
|
|
rolling_warm = "STARTUP_ROLLING_WARMUP_MS = _warm_model(rolling_model)"
|
|
final_warm = "STARTUP_WARMUP_MS = _warm_model(model)"
|
|
assert source.index(rolling_warm) < source.index(final_warm)
|
|
assert source.index(final_warm) < source.index("server = ThreadingHTTPServer")
|
|
assert 'kind="rolling"' in source
|
|
assert 'kind="speculative"' in source
|
|
assert '"last_rolling_decode_ms"' in source
|
|
assert '"last_speculative_decode_ms"' in source
|
|
|
|
|
|
def test_tiny_model_is_checksum_pinned_read_only_in_the_runtime_image():
|
|
dockerfile = DOCKERFILE.read_text()
|
|
digest = "65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9"
|
|
assert f"--checksum=sha256:{digest} --chmod=0444" in dockerfile
|
|
assert f"/models/{digest}/tiny.pt" in dockerfile
|
|
assert "HERMES_STT_ROLLING_MODEL=tiny" in dockerfile
|
|
assert "('large-v3-turbo.pt','small.pt','tiny.pt')" in dockerfile
|
|
|
|
|
|
def test_dual_model_contract_is_an_exact_voice_release_gate():
|
|
pipeline = PIPELINE.read_text()
|
|
assert "testing/tests/test_hermes_stt_rolling_model.py" in pipeline
|
|
source = SERVER.read_text()
|
|
assert '"path": "/v1/audio/transcriptions/stream"' in source
|
|
assert '"inference": INFERENCE_GATE.snapshot()' in source
|
|
assert '"rolling_model": ROLLING_MODEL_NAME' in source
|
|
assert "STARTUP_ROLLING_WARMUP_MS is not None" in source
|