59 lines · 1.9 KB
python
| 1 | from __future__ import annotations |
| 2 | |
| 3 | import sys |
| 4 | from pathlib import Path |
| 5 | from unittest import mock |
| 6 | |
| 7 | sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| 8 | |
| 9 | from space_tape import transcribe as tr |
| 10 | |
| 11 | |
| 12 | def test_longform_guards_passed_and_greedy() -> None: |
| 13 | calls: list[dict] = [] |
| 14 | |
| 15 | def fake_asr(path, **kwargs): |
| 16 | calls.append(kwargs) |
| 17 | return {"chunks": [{"timestamp": (0.0, 1.0), "text": "hi"}]} |
| 18 | |
| 19 | with mock.patch.object(tr, "_pipeline", return_value=fake_asr): |
| 20 | cues = tr.transcribe("x.wav", chunk_offset=300.0) |
| 21 | gen = calls[0]["generate_kwargs"] |
| 22 | assert gen["condition_on_prev_tokens"] is False |
| 23 | assert gen["no_speech_threshold"] == 0.6 |
| 24 | # A temperature tuple triggers sampling, which crashes on transformers 5.17. |
| 25 | assert gen["temperature"] == 0.0 |
| 26 | assert cues == [{"start": 300.0, "end": 301.0, "text": "hi"}] |
| 27 | |
| 28 | |
| 29 | def _raising_asr(message: str): |
| 30 | def asr(path, **kwargs): |
| 31 | raise ValueError(message) |
| 32 | return asr |
| 33 | |
| 34 | |
| 35 | def test_silent_chunk_with_word_timestamps_is_empty() -> None: |
| 36 | silent = _raising_asr("torch.cat(): expected a non-empty list of Tensors") |
| 37 | with mock.patch.object(tr, "_pipeline", return_value=silent): |
| 38 | assert tr.transcribe("x.wav", word_timestamps=True) == [] |
| 39 | |
| 40 | |
| 41 | def test_other_value_errors_still_raise() -> None: |
| 42 | for word_timestamps, message in ( |
| 43 | (True, "some other failure"), |
| 44 | (False, "torch.cat(): expected a non-empty list of Tensors"), |
| 45 | ): |
| 46 | with mock.patch.object(tr, "_pipeline", return_value=_raising_asr(message)): |
| 47 | try: |
| 48 | tr.transcribe("x.wav", word_timestamps=word_timestamps) |
| 49 | except ValueError: |
| 50 | continue |
| 51 | raise AssertionError(f"swallowed: {word_timestamps} {message}") |
| 52 | |
| 53 | |
| 54 | if __name__ == "__main__": |
| 55 | test_longform_guards_passed_and_greedy() |
| 56 | test_silent_chunk_with_word_timestamps_is_empty() |
| 57 | test_other_value_errors_still_raise() |
| 58 | print("ok") |