mirror of
https://github.com/yusufipk/dikte.git
synced 2026-09-11 10:56:10 +00:00
yusufipk
7a507c52eb
The two switches were written to disk by the Save button at the far end of the window, so a file transcribed with timestamps and cleanup was transcribed without either the next time. They belong to the run rather than to the form: they go to disk as they are ticked now, and the folder the last file came from goes with them. Stop only set a flag that was looked at between chunks, and a file under ten minutes is one chunk, so for most files it was looked at after the work it was meant to stop had already finished. Nothing that blocks is reached by a flag. The request is inside urlopen, ffmpeg is inside communicate, and a whisper on this machine is a process of ours that would grind on to the end of the chunk with nobody left to hand the answer to. So the socket is shut down under the read, ffmpeg is killed, and a local server is stopped and left for the next run to start again. Shutting the socket down rather than closing it is the point: close() alone leaves a thread already inside recv() waiting for bytes that are never coming now. The connection is registered before it has a socket, so a stop landing in the few lines between making a connection and blocking on it refuses the connection rather than missing it and letting urllib quietly open another.
273 lines
11 KiB
Python
273 lines
11 KiB
Python
"""Transcribing a file: splitting it, stamping it, and writing subtitles.
|
|
|
|
to_srt is the awkward one. The text is the authority on wording and the segments
|
|
on timing, and they meet at a whole-second stamp that a cleanup model was asked
|
|
to leave alone. It has to survive a model that wrapped a line, dropped one, or
|
|
made up a stamp nobody recorded.
|
|
"""
|
|
|
|
import contextlib
|
|
import time
|
|
import unittest
|
|
import wave
|
|
from unittest import mock
|
|
|
|
import api
|
|
import filetranscribe as ft
|
|
from tests.support import DikteTest, make_wav, silence, tone
|
|
|
|
|
|
class Timestamps(unittest.TestCase):
|
|
def test_under_an_hour(self):
|
|
self.assertEqual(ft.format_timestamp(0), "00:00")
|
|
self.assertEqual(ft.format_timestamp(65.9), "01:05")
|
|
self.assertEqual(ft.format_timestamp(599), "09:59")
|
|
|
|
def test_past_an_hour_the_hours_show(self):
|
|
self.assertEqual(ft.format_timestamp(3600), "1:00:00")
|
|
self.assertEqual(ft.format_timestamp(3725), "1:02:05")
|
|
|
|
def test_srt_wants_milliseconds_and_a_comma(self):
|
|
self.assertEqual(ft.srt_timestamp(0), "00:00:00,000")
|
|
self.assertEqual(ft.srt_timestamp(1.5), "00:00:01,500")
|
|
self.assertEqual(ft.srt_timestamp(3725.25), "01:02:05,250")
|
|
|
|
def test_a_negative_start_is_pulled_up_to_zero(self):
|
|
self.assertEqual(ft.srt_timestamp(-3), "00:00:00,000")
|
|
|
|
|
|
class ToSrt(unittest.TestCase):
|
|
def test_nothing_to_do(self):
|
|
self.assertEqual(ft.to_srt("", []), "")
|
|
self.assertEqual(ft.to_srt("no stamps here", []), "")
|
|
|
|
def test_a_cue_takes_its_timing_from_the_segment_it_came_from(self):
|
|
srt = ft.to_srt("[00:01] Hello there.", [(1.25, 2.75, "hello there")])
|
|
self.assertIn("00:00:01,250 --> 00:00:02,750", srt)
|
|
self.assertIn("Hello there.", srt)
|
|
|
|
def test_the_text_wins_on_wording(self):
|
|
"""Cleanup edits survive; only the timing comes from the segments."""
|
|
srt = ft.to_srt("[00:01] Hello there.", [(1.0, 2.0, "uh hello uh there")])
|
|
self.assertIn("Hello there.", srt)
|
|
self.assertNotIn("uh", srt)
|
|
|
|
def test_cues_are_numbered_from_one(self):
|
|
srt = ft.to_srt("[00:00] One\n[00:02] Two",
|
|
[(0.0, 1.0, "One"), (2.0, 3.0, "Two")])
|
|
self.assertTrue(srt.startswith("1\n"))
|
|
self.assertIn("\n2\n", srt)
|
|
|
|
def test_a_wrapped_line_joins_the_cue_above_it(self):
|
|
srt = ft.to_srt("[00:01] Hello there,\nand welcome.",
|
|
[(1.0, 4.0, "hello there and welcome")])
|
|
self.assertIn("Hello there, and welcome.", srt)
|
|
self.assertEqual(srt.count(" --> "), 1)
|
|
|
|
def test_a_stamp_nobody_recorded_still_gets_timing(self):
|
|
srt = ft.to_srt("[00:05] Invented.", [])
|
|
self.assertIn("00:00:05,000 --> ", srt)
|
|
|
|
def test_a_cue_with_no_end_runs_until_the_next_one(self):
|
|
srt = ft.to_srt("[00:00] One\n[00:04] Two", [])
|
|
self.assertIn("00:00:00,000 --> 00:00:04,000", srt)
|
|
|
|
def test_the_last_cue_gets_a_minimum_length(self):
|
|
srt = ft.to_srt("[00:10] Last words.", [])
|
|
self.assertIn("00:00:10,000 --> 00:00:11,500", srt)
|
|
|
|
def test_a_cue_is_cut_short_when_the_next_one_starts_first(self):
|
|
"""Whisper's end times overlap now and then; subtitles must not."""
|
|
srt = ft.to_srt("[00:00] One\n[00:02] Two",
|
|
[(0.0, 9.0, "One"), (2.0, 3.0, "Two")])
|
|
self.assertIn("00:00:00,000 --> 00:00:02,000", srt)
|
|
|
|
def test_blank_lines_and_empty_cues_are_dropped(self):
|
|
srt = ft.to_srt("[00:00] One\n\n[00:02]\n[00:03] Three", [])
|
|
self.assertEqual(srt.count(" --> "), 2)
|
|
|
|
def test_the_hour_form_of_a_stamp_is_understood(self):
|
|
srt = ft.to_srt("[1:02:05] Late.", [])
|
|
self.assertIn("01:02:05,000", srt)
|
|
|
|
def test_the_file_ends_with_a_newline(self):
|
|
self.assertTrue(ft.to_srt("[00:00] One", []).endswith("\n"))
|
|
|
|
|
|
class SplitText(unittest.TestCase):
|
|
def test_short_text_stays_whole(self):
|
|
self.assertEqual(ft.split_text("hello", False), ["hello"])
|
|
|
|
def test_a_long_transcript_is_broken_up(self):
|
|
text = " ".join(["word"] * 8000)
|
|
blocks = ft.split_text(text, False)
|
|
self.assertGreater(len(blocks), 1)
|
|
for block in blocks:
|
|
self.assertLessEqual(len(block), ft.CLEANUP_CHUNK_CHARS)
|
|
|
|
def test_nothing_is_lost_in_the_splitting(self):
|
|
text = " ".join(f"word{index}" for index in range(4000))
|
|
self.assertEqual(" ".join(ft.split_text(text, False)), text)
|
|
|
|
def test_a_timestamped_transcript_is_never_broken_mid_line(self):
|
|
text = "\n".join(f"[00:{index:02d}] a line of some length here"
|
|
for index in range(600))
|
|
blocks = ft.split_text(text, True)
|
|
self.assertGreater(len(blocks), 1)
|
|
for block in blocks:
|
|
for line in block.splitlines():
|
|
self.assertTrue(line.startswith("["))
|
|
|
|
def test_a_single_line_longer_than_the_limit_is_kept_whole(self):
|
|
text = "x" * (ft.CLEANUP_CHUNK_CHARS + 100)
|
|
self.assertEqual(ft.split_text(text, True), [text])
|
|
|
|
|
|
class SplitWav(DikteTest):
|
|
def wav(self, seconds, name="in.wav"):
|
|
return make_wav(self.path(name), silence(seconds))
|
|
|
|
def test_a_short_file_is_handed_back_as_it_is(self):
|
|
path = self.wav(2)
|
|
self.assertEqual(ft.split_wav(path, self.root), [(path, 0.0)])
|
|
|
|
def test_a_long_file_is_cut_at_the_chunk_length(self):
|
|
path = self.wav(5)
|
|
with mock.patch.object(ft, "CHUNK_SECONDS", 2):
|
|
chunks = ft.split_wav(path, self.root)
|
|
self.assertEqual([offset for _, offset in chunks], [0, 2, 4])
|
|
|
|
def test_the_chunks_add_up_to_the_original(self):
|
|
path = self.wav(5)
|
|
with mock.patch.object(ft, "CHUNK_SECONDS", 2):
|
|
chunks = ft.split_wav(path, self.root)
|
|
total = 0
|
|
for chunk_path, _ in chunks:
|
|
with contextlib.closing(wave.open(chunk_path, "rb")) as wav:
|
|
total += wav.getnframes()
|
|
self.assertEqual(wav.getframerate(), 16000)
|
|
self.assertEqual(total, 5 * 16000)
|
|
|
|
def test_the_chunks_do_not_write_over_each_other(self):
|
|
path = self.wav(5)
|
|
with mock.patch.object(ft, "CHUNK_SECONDS", 2):
|
|
chunks = ft.split_wav(path, self.root)
|
|
self.assertEqual(len({chunk for chunk, _ in chunks}), len(chunks))
|
|
|
|
|
|
class Transcriber(DikteTest):
|
|
"""The chain, with ffmpeg and both API calls faked."""
|
|
|
|
def setUp(self):
|
|
super().setUp()
|
|
self.source = make_wav(self.path("input.wav"), tone(1.0))
|
|
self.conf = self.config(openrouter_api_key="sk-or-test")
|
|
|
|
def run_chain(self, timestamps=False, cleanup=False, transcript="raw text",
|
|
segments=None, cleaned="clean text", fail=None):
|
|
worker = ft.FileTranscriber(self.conf)
|
|
done, failures, progress = [], [], []
|
|
worker.finished.connect(lambda *args: done.append(args))
|
|
worker.failed.connect(failures.append)
|
|
worker.progress.connect(progress.append)
|
|
|
|
def to_wav(path, workdir, aborter=None):
|
|
return make_wav(self.path("converted.wav"), tone(1.0))
|
|
|
|
with mock.patch.object(ft, "_to_wav", side_effect=to_wav), \
|
|
mock.patch.object(ft.shutil, "which", return_value="/usr/bin/ffmpeg"), \
|
|
mock.patch.object(api, "transcribe",
|
|
side_effect=fail or (lambda *a, **k: transcript)), \
|
|
mock.patch.object(api, "transcribe_segments",
|
|
return_value=segments or [(0.0, 1.0, "raw text")]), \
|
|
mock.patch.object(api, "cleanup", return_value=cleaned) as cleanup_call:
|
|
# The chain is run here rather than through start(): its signals are
|
|
# emitted from the worker thread, and a queued connection would need
|
|
# an event loop to deliver them. This is the same code, one frame down.
|
|
worker._work(self.source, timestamps, cleanup)
|
|
return done, failures, progress, cleanup_call
|
|
|
|
def test_plain_text_out(self):
|
|
done, failures, _, _ = self.run_chain()
|
|
self.assertEqual(failures, [])
|
|
self.assertEqual(done[0][0], "raw text")
|
|
|
|
def test_cleanup_replaces_the_text(self):
|
|
done, _, _, _ = self.run_chain(cleanup=True)
|
|
self.assertEqual(done[0][0], "clean text")
|
|
|
|
def test_cleanup_is_told_it_is_writing_subtitles(self):
|
|
_, _, _, cleanup_call = self.run_chain(cleanup=True)
|
|
prompt = cleanup_call.call_args.args[3]
|
|
self.assertEqual(prompt, self.conf.cleanup_prompt(subtitles=True))
|
|
|
|
def test_timestamps_come_back_as_segments_and_as_stamped_lines(self):
|
|
done, _, _, _ = self.run_chain(
|
|
timestamps=True, segments=[(0.0, 1.0, "one"), (2.0, 3.0, "two")])
|
|
text, segments = done[0]
|
|
self.assertEqual(text, "[00:00] one\n[00:02] two")
|
|
self.assertEqual(len(segments), 2)
|
|
|
|
def test_no_ffmpeg_installed(self):
|
|
worker = ft.FileTranscriber(self.conf)
|
|
failures = []
|
|
worker.failed.connect(failures.append)
|
|
with mock.patch.object(ft.shutil, "which", return_value=None):
|
|
worker._work(self.source, False, False)
|
|
self.assertIn("ffmpeg", failures[0])
|
|
|
|
def test_an_api_failure_is_reported_rather_than_raised(self):
|
|
def boom(*args, **kwargs):
|
|
raise api.ApiError("OpenAI rejected the API key")
|
|
_, failures, _, _ = self.run_chain(fail=boom)
|
|
self.assertIn("rejected", failures[0])
|
|
|
|
def test_empty_text_is_not_sent_to_cleanup(self):
|
|
_, _, _, cleanup_call = self.run_chain(cleanup=True, transcript="")
|
|
cleanup_call.assert_not_called()
|
|
|
|
def test_a_stopped_run_is_not_a_failure(self):
|
|
def stopped(*args, **kwargs):
|
|
raise api.Aborted
|
|
done, failures, progress, _ = self.run_chain(fail=stopped)
|
|
self.assertEqual(failures, [])
|
|
self.assertEqual(done, [])
|
|
self.assertEqual(progress[-1], "Stopped.")
|
|
|
|
def test_the_request_is_handed_the_stop_to_watch(self):
|
|
worker = ft.FileTranscriber(self.conf)
|
|
with mock.patch.object(ft, "_to_wav", side_effect=lambda *a: self.source), \
|
|
mock.patch.object(ft.shutil, "which", return_value="/usr/bin/ffmpeg"), \
|
|
mock.patch.object(api, "transcribe", return_value="text") as call:
|
|
worker._work(self.source, False, False)
|
|
self.assertIs(call.call_args.kwargs["aborter"], worker._abort)
|
|
|
|
def test_stopping_a_local_run_stops_the_model_with_it(self):
|
|
"""Closing the socket is nothing to a process of ours: it would grind on
|
|
to the end of the chunk with nobody left to hand the answer to."""
|
|
worker = ft.FileTranscriber(self.conf)
|
|
worker._local = mock.Mock()
|
|
worker.stop()
|
|
self.assertTrue(worker._abort.aborted)
|
|
for _ in range(100):
|
|
if worker._local.stop.called:
|
|
break
|
|
time.sleep(0.01)
|
|
worker._local.stop.assert_called_once_with()
|
|
|
|
def test_a_run_that_is_over_leaves_the_model_alone(self):
|
|
worker = ft.FileTranscriber(self.conf)
|
|
worker.stop()
|
|
self.assertTrue(worker._abort.aborted)
|
|
|
|
def test_a_second_start_while_one_is_running_is_ignored(self):
|
|
worker = ft.FileTranscriber(self.conf)
|
|
worker._thread = mock.Mock(is_alive=lambda: True)
|
|
self.assertTrue(worker.busy)
|
|
worker.start(self.source, False, False)
|
|
self.assertTrue(worker.busy)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|