389 lines
13 KiB
Python
389 lines
13 KiB
Python
import io
|
|
import wave
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from sentence_api.transcription import (
|
|
MossTranscriber,
|
|
TranscriptionSegment,
|
|
WordTimestamp,
|
|
_parse_json_segments,
|
|
_plan_chunks,
|
|
_parse_word_timestamps,
|
|
split_sentences_at_punctuation,
|
|
_wav_bytes,
|
|
_wav_duration_seconds,
|
|
)
|
|
|
|
|
|
SAMPLE_RATE = 16000
|
|
|
|
|
|
def _speech(seconds: float, sample_rate: int = SAMPLE_RATE) -> np.ndarray:
|
|
t = np.arange(int(seconds * sample_rate)) / sample_rate
|
|
return 0.25 * np.sin(2 * np.pi * 220 * t)
|
|
|
|
|
|
def _silence(seconds: float, sample_rate: int = SAMPLE_RATE) -> np.ndarray:
|
|
return np.zeros(int(seconds * sample_rate))
|
|
|
|
|
|
def _audio(*parts: np.ndarray, sample_rate: int = SAMPLE_RATE) -> np.ndarray:
|
|
return np.concatenate([np.asarray(part) for part in parts]).astype(np.float32)
|
|
|
|
|
|
def test_plan_chunks_cuts_at_silence_gaps():
|
|
samples = _audio(
|
|
_speech(10),
|
|
_silence(0.8),
|
|
_speech(29.2),
|
|
_silence(1.0),
|
|
_speech(29.0),
|
|
_silence(0.6),
|
|
_speech(24.4),
|
|
)
|
|
chunks = _plan_chunks(samples, SAMPLE_RATE)
|
|
assert chunks[0][0] == pytest.approx(0.0)
|
|
assert chunks[-1][1] == pytest.approx(95.0)
|
|
for (start, end), (next_start, _) in zip(chunks, chunks[1:]):
|
|
assert end == pytest.approx(next_start)
|
|
assert end - start <= 34.0
|
|
assert 10.0 <= chunks[0][1] <= 10.8
|
|
assert 40.0 <= chunks[1][1] <= 41.0
|
|
assert 70.0 <= chunks[2][1] <= 70.6
|
|
|
|
|
|
def test_plan_chunks_falls_back_to_fixed_windows():
|
|
chunks = _plan_chunks(_speech(95), SAMPLE_RATE)
|
|
assert chunks == [(0.0, 30.0), (30.0, 60.0), (60.0, 90.0), (90.0, 95.0)]
|
|
|
|
|
|
def test_plan_chunks_merges_short_tail():
|
|
chunks = _plan_chunks(_speech(93), SAMPLE_RATE)
|
|
assert chunks == [(0.0, 30.0), (30.0, 60.0), (60.0, 93.0)]
|
|
|
|
|
|
def test_parse_json_segments_filters_hallucinated_segments():
|
|
raw = [
|
|
{"start": 0.0, "end": 3.0, "text": "hello world", "compression_ratio": 1.4},
|
|
{"start": 3.5, "end": 6.0, "text": "hello world", "compression_ratio": 12.0},
|
|
{"start": 6.5, "end": 9.0, "text": "no ratio field"},
|
|
]
|
|
filtered = _parse_json_segments(raw, compression_limit=2.4)
|
|
assert [segment.text for segment in filtered] == ["hello world", "no ratio field"]
|
|
assert len(_parse_json_segments(raw)) == 3
|
|
|
|
|
|
def test_wav_bytes_roundtrip(tmp_path):
|
|
samples = _speech(2.0)
|
|
data = _wav_bytes(samples, SAMPLE_RATE)
|
|
assert data.startswith(b"RIFF")
|
|
path = tmp_path / "chunk.wav"
|
|
path.write_bytes(data)
|
|
assert _wav_duration_seconds(path) == pytest.approx(2.0, abs=0.05)
|
|
with wave.open(str(path), "rb") as wav:
|
|
assert wav.getnchannels() == 1
|
|
assert wav.getsampwidth() == 2
|
|
assert wav.getframerate() == SAMPLE_RATE
|
|
|
|
|
|
def test_wav_duration_returns_none_for_non_wav(tmp_path):
|
|
path = tmp_path / "note.txt"
|
|
path.write_text("not audio")
|
|
assert _wav_duration_seconds(path) is None
|
|
|
|
|
|
def test_chunked_transcribe_offsets_timestamps(monkeypatch):
|
|
requests = []
|
|
call_count = {"n": 0}
|
|
|
|
def fake_post(endpoint, data=None, files=None, timeout=None):
|
|
call_count["n"] += 1
|
|
requests.append(data)
|
|
audio = files["file"][1].read()
|
|
with wave.open(io.BytesIO(audio), "rb") as wav:
|
|
duration = wav.getnframes() / wav.getframerate()
|
|
if call_count["n"] == 1:
|
|
segments = [
|
|
{"start": 1.0, "end": 4.0, "text": "first part", "compression_ratio": 1.4},
|
|
{"start": 27.0, "end": 30.0, "text": "second part", "compression_ratio": 1.5},
|
|
]
|
|
elif call_count["n"] == 2:
|
|
segments = [
|
|
{
|
|
"start": 1.0,
|
|
"end": 4.0,
|
|
"text": "third part",
|
|
"compression_ratio": 1.4,
|
|
"words": [
|
|
{"word": "third", "start": 1.0, "end": 2.5},
|
|
{"word": " part", "start": 2.6, "end": 4.0},
|
|
],
|
|
},
|
|
{"start": 27.0, "end": 30.0, "text": "fourth part", "compression_ratio": 1.5},
|
|
]
|
|
else:
|
|
segments = [{"start": 0.5, "end": 3.5, "text": "tail part", "compression_ratio": 1.3}]
|
|
assert abs(duration - 30.0) < 1.0 or abs(duration - 5.0) < 1.0
|
|
|
|
class Response:
|
|
status_code = 200
|
|
|
|
def raise_for_status(self):
|
|
pass
|
|
|
|
def json(self):
|
|
return {"text": "", "segments": segments}
|
|
|
|
return Response()
|
|
|
|
monkeypatch.setattr("sentence_api.transcription.httpx.post", fake_post)
|
|
transcriber = MossTranscriber(endpoint="http://whisper:9000", model="whisper")
|
|
result = transcriber._transcribe_chunked(_speech(65), SAMPLE_RATE, None)
|
|
assert [(s.start_seconds, s.end_seconds, s.text) for s in result.segments] == [
|
|
(1.0, 4.0, "first part"),
|
|
(27.0, 30.0, "second part"),
|
|
(31.0, 34.0, "third part"),
|
|
(57.0, 60.0, "fourth part"),
|
|
(60.5, 63.5, "tail part"),
|
|
]
|
|
assert result.text == "first part second part third part fourth part tail part"
|
|
assert [(w.start_seconds, w.end_seconds, w.text) for w in result.segments[2].words] == [
|
|
(31.0, 32.5, "third"),
|
|
(32.6, 34.0, "part"),
|
|
]
|
|
assert result.segments[0].words is None
|
|
assert len(requests) == 3
|
|
assert all(request.get("condition_on_previous_text") == "false" for request in requests)
|
|
assert all(request.get("timestamp_granularities[]") == "word" for request in requests)
|
|
|
|
|
|
def test_parse_json_segments_parses_word_timestamps():
|
|
raw = [
|
|
{
|
|
"start": 0.0,
|
|
"end": 4.0,
|
|
"text": "Hello world. Good day.",
|
|
"words": [
|
|
{"word": "Hello", "start": 0.0, "end": 0.6},
|
|
{"word": " world.", "start": 0.7, "end": 1.5},
|
|
{"word": " Good", "start": 1.8, "end": 2.4},
|
|
{"word": " day.", "start": 2.5, "end": 3.2},
|
|
],
|
|
},
|
|
{"start": 4.5, "end": 6.0, "text": "no words"},
|
|
]
|
|
segments = _parse_json_segments(raw)
|
|
assert [word.text for word in segments[0].words] == ["Hello", "world.", "Good", "day."]
|
|
assert segments[1].words is None
|
|
|
|
|
|
def test_parse_word_timestamps_ignores_invalid_entries():
|
|
words = _parse_word_timestamps(
|
|
[
|
|
{"word": "ok", "start": 0.0, "end": 0.5},
|
|
{"word": "bad"},
|
|
{"word": "", "start": 1.0, "end": 1.5},
|
|
{"word": "flat", "start": 2.0, "end": 2.0},
|
|
]
|
|
)
|
|
assert [(word.text, word.start_seconds, word.end_seconds) for word in words] == [
|
|
("ok", 0.0, 0.5)
|
|
]
|
|
assert _parse_word_timestamps(None) is None
|
|
assert _parse_word_timestamps("nope") is None
|
|
|
|
|
|
def test_split_sentences_merges_segments_without_punctuation():
|
|
segments = [
|
|
TranscriptionSegment(
|
|
start_seconds=0.0,
|
|
end_seconds=8.0,
|
|
text=(
|
|
"So I've been single for about four years now and I find it hard "
|
|
"to meet a guy especially in"
|
|
),
|
|
),
|
|
TranscriptionSegment(
|
|
start_seconds=8.0,
|
|
end_seconds=9.5,
|
|
text="California.",
|
|
words=[WordTimestamp(8.0, 9.5, "California.")],
|
|
),
|
|
TranscriptionSegment(
|
|
start_seconds=9.5,
|
|
end_seconds=20.0,
|
|
text=(
|
|
"I've tried dating apps, I've met friends through friends, "
|
|
"I've done online dating"
|
|
),
|
|
),
|
|
TranscriptionSegment(
|
|
start_seconds=20.0,
|
|
end_seconds=24.0,
|
|
text=(
|
|
"and I just feel like most of the guys that I meet aren't serious "
|
|
"about a relationship."
|
|
),
|
|
words=[WordTimestamp(20.0, 24.0, "relationship.")],
|
|
),
|
|
]
|
|
sentences = split_sentences_at_punctuation(segments)
|
|
assert [(s.text, s.start_seconds, s.end_seconds) for s in sentences] == [
|
|
(
|
|
"So I've been single for about four years now and I find it hard to meet "
|
|
"a guy especially in California.",
|
|
0.0,
|
|
9.5,
|
|
),
|
|
(
|
|
"I've tried dating apps, I've met friends through friends, I've done "
|
|
"online dating and I just feel like most of the guys that I meet aren't "
|
|
"serious about a relationship.",
|
|
9.5,
|
|
24.0,
|
|
),
|
|
]
|
|
|
|
|
|
def test_split_sentences_uses_word_timestamps_within_segment():
|
|
segment = TranscriptionSegment(
|
|
start_seconds=0.0,
|
|
end_seconds=4.0,
|
|
text="Hello world. Good day. Nice to meet you.",
|
|
words=[
|
|
WordTimestamp(0.0, 0.6, "Hello"),
|
|
WordTimestamp(0.7, 1.5, "world."),
|
|
WordTimestamp(1.8, 2.4, "Good"),
|
|
WordTimestamp(2.5, 3.2, "day."),
|
|
WordTimestamp(3.3, 3.8, "Nice"),
|
|
WordTimestamp(3.9, 4.0, "you."),
|
|
],
|
|
)
|
|
sentences = split_sentences_at_punctuation([segment])
|
|
assert [(s.text, s.start_seconds, s.end_seconds) for s in sentences] == [
|
|
("Hello world.", 0.0, 1.5),
|
|
("Good day.", 1.5, 3.2),
|
|
("Nice to meet you.", 3.2, 4.0),
|
|
]
|
|
|
|
|
|
def test_split_sentences_splits_at_question_marks():
|
|
segment = TranscriptionSegment(
|
|
start_seconds=0.0,
|
|
end_seconds=10.0,
|
|
text="Really? Wait a minute. What about now?",
|
|
)
|
|
sentences = split_sentences_at_punctuation([segment])
|
|
assert [s.text for s in sentences] == [
|
|
"Really?",
|
|
"Wait a minute.",
|
|
"What about now?",
|
|
]
|
|
assert sentences[0].start_seconds == 0.0
|
|
assert sentences[0].end_seconds == sentences[1].start_seconds
|
|
assert sentences[1].end_seconds == sentences[2].start_seconds
|
|
assert sentences[2].end_seconds == 10.0
|
|
|
|
|
|
def test_split_sentences_falls_back_to_proportional():
|
|
segments = [
|
|
TranscriptionSegment(
|
|
start_seconds=10.0, end_seconds=14.0, text="This is the first"
|
|
),
|
|
TranscriptionSegment(
|
|
start_seconds=14.0,
|
|
end_seconds=18.0,
|
|
text="sentence. And this is the second one.",
|
|
),
|
|
]
|
|
sentences = split_sentences_at_punctuation(segments)
|
|
assert [s.text for s in sentences] == [
|
|
"This is the first sentence.",
|
|
"And this is the second one.",
|
|
]
|
|
assert sentences[0].start_seconds == 10.0
|
|
assert sentences[0].end_seconds == sentences[1].start_seconds
|
|
assert sentences[1].end_seconds == 18.0
|
|
|
|
|
|
def test_split_sentences_falls_back_when_word_ends_are_invalid():
|
|
segment = TranscriptionSegment(
|
|
start_seconds=0.0,
|
|
end_seconds=2.0,
|
|
text="One. Two.",
|
|
words=[
|
|
WordTimestamp(0.0, 2.5, "One."),
|
|
WordTimestamp(2.6, 3.0, "Two."),
|
|
],
|
|
)
|
|
sentences = split_sentences_at_punctuation([segment])
|
|
assert [s.text for s in sentences] == ["One.", "Two."]
|
|
assert sentences[0].end_seconds == sentences[1].start_seconds
|
|
assert sentences[1].end_seconds == 2.0
|
|
|
|
|
|
def test_split_sentences_keeps_incomplete_tail_as_sentence():
|
|
segment = TranscriptionSegment(
|
|
start_seconds=1.0, end_seconds=2.0, text="no period here"
|
|
)
|
|
sentences = split_sentences_at_punctuation([segment])
|
|
assert [(s.text, s.start_seconds, s.end_seconds) for s in sentences] == [
|
|
("no period here", 1.0, 2.0)
|
|
]
|
|
assert split_sentences_at_punctuation([]) == []
|
|
|
|
|
|
def test_split_sentences_handles_ellipsis_and_dots_only():
|
|
segment = TranscriptionSegment(
|
|
start_seconds=0.0,
|
|
end_seconds=2.0,
|
|
text="Wait... What? ...",
|
|
)
|
|
sentences = split_sentences_at_punctuation([segment])
|
|
assert [s.text for s in sentences] == ["Wait...", "What?"]
|
|
assert sentences[0].start_seconds == 0.0
|
|
assert sentences[0].end_seconds == sentences[1].start_seconds
|
|
assert sentences[1].end_seconds == 2.0
|
|
assert split_sentences_at_punctuation(
|
|
[TranscriptionSegment(0.0, 1.0, "...")]
|
|
) == []
|
|
|
|
|
|
def test_post_audio_requests_word_timestamps_and_falls_back(monkeypatch):
|
|
calls = []
|
|
|
|
class RejectedResponse:
|
|
status_code = 400
|
|
|
|
def raise_for_status(self):
|
|
pass
|
|
|
|
def json(self):
|
|
return {"error": "unknown parameter"}
|
|
|
|
class OkResponse:
|
|
status_code = 200
|
|
|
|
def raise_for_status(self):
|
|
pass
|
|
|
|
def json(self):
|
|
return {"text": "hi.", "segments": []}
|
|
|
|
def fake_post(endpoint, data=None, files=None, timeout=None):
|
|
calls.append(dict(data))
|
|
if "timestamp_granularities[]" in data:
|
|
return RejectedResponse()
|
|
return OkResponse()
|
|
|
|
monkeypatch.setattr("sentence_api.transcription.httpx.post", fake_post)
|
|
transcriber = MossTranscriber(endpoint="http://whisper:9000", model="whisper")
|
|
result = transcriber._post_audio(
|
|
io.BytesIO(b"fake-audio"), "clip.wav", "audio/wav", "en"
|
|
)
|
|
assert result["text"] == "hi."
|
|
assert calls[0]["timestamp_granularities[]"] == "word"
|
|
assert "timestamp_granularities[]" not in calls[1]
|