f9feb0b393
Leftover ffmpeg-era imports from the earlier draft of the CPU transcribe smoke test; the final numpy-array-bypass version doesn't need them. Ruff autofix.
138 lines
4.6 KiB
Python
138 lines
4.6 KiB
Python
"""Unit + integration coverage for second_brain.transcribe.
|
|
|
|
Unit-level: resolve_settings + write_srt are pure-Python and run anywhere.
|
|
|
|
Integration: an end-to-end CPU/int8 transcribe of a synthesized audio
|
|
clip — auto-skipped when ffmpeg or faster-whisper is missing (the dev
|
|
side doesn't install faster-whisper by default). Only the *tower* with
|
|
GPU large-v3 can validate the production hot path; this test proves the
|
|
wiring + the SRT contract + the lazy-import path on cheap CPU.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
|
|
import pytest
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pure helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_resolve_settings_picks_int8_for_cpu_by_default():
|
|
from second_brain.transcribe import resolve_settings
|
|
|
|
out = resolve_settings({"whisper_device": "cpu"})
|
|
assert out["device"] == "cpu"
|
|
assert out["compute_type"] == "int8"
|
|
assert out["model"] == "large-v3"
|
|
|
|
|
|
def test_resolve_settings_picks_float16_for_cuda_by_default():
|
|
from second_brain.transcribe import resolve_settings
|
|
|
|
out = resolve_settings({"whisper_device": "cuda"})
|
|
assert out["device"] == "cuda"
|
|
assert out["compute_type"] == "float16"
|
|
|
|
|
|
def test_resolve_settings_env_overrides_block(monkeypatch):
|
|
from second_brain.transcribe import resolve_settings
|
|
|
|
monkeypatch.setenv("WHISPER_MODEL", "small")
|
|
monkeypatch.setenv("WHISPER_DEVICE", "cpu")
|
|
monkeypatch.setenv("WHISPER_COMPUTE_TYPE", "int8_float16")
|
|
out = resolve_settings({"whisper_device": "cuda"})
|
|
assert out["model"] == "small"
|
|
assert out["device"] == "cpu"
|
|
assert out["compute_type"] == "int8_float16"
|
|
|
|
|
|
def test_write_srt_roundtrip(tmp_path):
|
|
from second_brain.transcribe import write_srt
|
|
|
|
segs = [
|
|
{"start": 0.0, "end": 1.5, "text": "hello"},
|
|
{"start": 1.5, "end": 3.25, "text": "world"},
|
|
]
|
|
out = tmp_path / "x.srt"
|
|
write_srt(segs, out)
|
|
body = out.read_text()
|
|
assert "1\n00:00:00,000 --> 00:00:01,500\nhello" in body
|
|
assert "2\n00:00:01,500 --> 00:00:03,250\nworld" in body
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Integration: synthesized audio → CPU/int8 faster-whisper
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _have_faster_whisper() -> bool:
|
|
try:
|
|
import faster_whisper # type: ignore[import-not-found] # noqa: F401
|
|
except ImportError:
|
|
return False
|
|
return True
|
|
|
|
|
|
def _have_numpy() -> bool:
|
|
try:
|
|
import numpy # noqa: F401
|
|
except ImportError:
|
|
return False
|
|
return True
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not (_have_faster_whisper() and _have_numpy()),
|
|
reason="needs faster-whisper (uv sync --extra tower) + numpy to exercise CPU path",
|
|
)
|
|
def test_cpu_transcribe_smoke(tmp_path):
|
|
"""Hand faster-whisper a synthesized audio array on CPU/int8 with `tiny`.
|
|
|
|
Pure-tone input gives empty / minimal recognition — the point is the
|
|
wiring: model loads, the segments iterator drains, an SRT is writable.
|
|
Using `tiny` to keep this under ~30s coldcache. Skipped on dev where
|
|
faster-whisper isn't installed; only the tower's `uv sync --extra
|
|
tower` brings it in.
|
|
|
|
Bypassing ffmpeg by passing a numpy array — faster-whisper.transcribe
|
|
accepts (path | file-like | np.ndarray). This exercises the same code
|
|
path the worker uses, only with zero external binaries.
|
|
"""
|
|
import numpy as np
|
|
|
|
from second_brain.transcribe import write_srt
|
|
|
|
# Build 2s of 440Hz mono at 16kHz, the model's native sample rate.
|
|
sr = 16000
|
|
duration = 2.0
|
|
t = np.arange(0, int(sr * duration)) / sr
|
|
audio = (0.2 * np.sin(2 * np.pi * 440 * t)).astype(np.float32)
|
|
|
|
# Load `tiny` to keep this cheap; we're not asserting transcription
|
|
# quality, just that the pipeline executes and returns the documented
|
|
# shape.
|
|
from faster_whisper import WhisperModel # type: ignore[import-not-found]
|
|
|
|
model = WhisperModel("tiny", device="cpu", compute_type="int8")
|
|
segments_iter, _ = model.transcribe(audio, language="en", beam_size=1)
|
|
|
|
# Materialize to mirror what Transcriber.transcribe does internally.
|
|
segments: list[dict] = []
|
|
text_parts: list[str] = []
|
|
for seg in segments_iter:
|
|
s = (seg.text or "").strip()
|
|
segments.append({"start": float(seg.start), "end": float(seg.end), "text": s})
|
|
if s:
|
|
text_parts.append(s)
|
|
text = " ".join(text_parts)
|
|
|
|
assert isinstance(text, str)
|
|
assert isinstance(segments, list)
|
|
|
|
srt = tmp_path / "out.srt"
|
|
write_srt(segments, srt)
|
|
assert srt.exists()
|