wehub-resource-sync
116 lines
3.5 KiB
Python
116 lines
3.5 KiB
Python
import importlib.util
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
|
|
|
|
def load_generate_subtitle_module():
|
|
script_path = Path(__file__).resolve().parents[1] / "examples" / "subtitle" / "generate_subtitle.py"
|
|
spec = importlib.util.spec_from_file_location("generate_subtitle", script_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def test_generate_subtitle_requests_sentence_timestamps_and_writes_segmented_srt(
|
|
monkeypatch, tmp_path
|
|
):
|
|
module = load_generate_subtitle_module()
|
|
captured = {}
|
|
|
|
class FakeAutoModel:
|
|
def __init__(self, **kwargs):
|
|
captured["model_kwargs"] = kwargs
|
|
|
|
def generate(self, **kwargs):
|
|
captured["generate_kwargs"] = kwargs
|
|
if kwargs.get("sentence_timestamp") and kwargs.get("output_timestamp"):
|
|
return [
|
|
{
|
|
"text": "<|zh|>First segment Second segment",
|
|
"sentence_info": [
|
|
{"start": 0, "end": 3500, "text": "<|zh|>First segment"},
|
|
{"start": 3500, "end": 7200, "sentence": "Second segment", "spk": 1},
|
|
],
|
|
}
|
|
]
|
|
return [{"text": "<|zh|>First segment Second segment"}]
|
|
|
|
fake_funasr = types.ModuleType("funasr")
|
|
fake_funasr.AutoModel = FakeAutoModel
|
|
monkeypatch.setitem(sys.modules, "funasr", fake_funasr)
|
|
|
|
input_path = tmp_path / "input.mp4"
|
|
input_path.write_bytes(b"fake media")
|
|
output_path = tmp_path / "sub.srt"
|
|
monkeypatch.setattr(
|
|
sys,
|
|
"argv",
|
|
[
|
|
"generate_subtitle.py",
|
|
str(input_path),
|
|
"-o",
|
|
str(output_path),
|
|
"--lang",
|
|
"zh",
|
|
"--device",
|
|
"cpu",
|
|
"--spk",
|
|
],
|
|
)
|
|
|
|
module.main()
|
|
|
|
assert captured["model_kwargs"]["punc_model"] == "ct-punc"
|
|
assert captured["generate_kwargs"]["language"] == "zh"
|
|
assert captured["generate_kwargs"]["sentence_timestamp"] is True
|
|
assert captured["generate_kwargs"]["output_timestamp"] is True
|
|
assert captured["generate_kwargs"]["return_time_stamps"] is True
|
|
assert output_path.read_text(encoding="utf-8") == (
|
|
"1\n"
|
|
"00:00:00,000 --> 00:00:03,500\n"
|
|
"First segment\n\n"
|
|
"2\n"
|
|
"00:00:03,500 --> 00:00:07,200\n"
|
|
"[Speaker 1] Second segment\n\n"
|
|
)
|
|
|
|
|
|
def test_generate_subtitle_falls_back_to_timestamp_bounds_when_sentence_info_missing(
|
|
monkeypatch, tmp_path
|
|
):
|
|
module = load_generate_subtitle_module()
|
|
|
|
class FakeAutoModel:
|
|
def __init__(self, **kwargs):
|
|
pass
|
|
|
|
def generate(self, **kwargs):
|
|
return [
|
|
{
|
|
"text": "<|zh|>Full text",
|
|
"timestamp": [[250, 1000], [1000, 3200]],
|
|
}
|
|
]
|
|
|
|
fake_funasr = types.ModuleType("funasr")
|
|
fake_funasr.AutoModel = FakeAutoModel
|
|
monkeypatch.setitem(sys.modules, "funasr", fake_funasr)
|
|
|
|
input_path = tmp_path / "input.mp4"
|
|
input_path.write_bytes(b"fake media")
|
|
output_path = tmp_path / "sub.srt"
|
|
monkeypatch.setattr(
|
|
sys,
|
|
"argv",
|
|
["generate_subtitle.py", str(input_path), "-o", str(output_path), "--device", "cpu"],
|
|
)
|
|
|
|
module.main()
|
|
|
|
assert output_path.read_text(encoding="utf-8") == (
|
|
"1\n"
|
|
"00:00:00,250 --> 00:00:03,200\n"
|
|
"Full text\n\n"
|
|
)
|