pipeline.py

transcription_pipeline/pipeline.py

1. """One recording through the whole pipeline.
3. normalize -> transcribe (local or DeepInfra) -> align (word timestamps)
4. -> diarize (who spoke when) -> sentences and speaker turns -> write outputs.
6. The model-backed steps are passed in as callables, so tests can drive the whole
7. flow with stand-ins and the CLI wires in the real ones (steps.py).
8. """
10. from __future__ import annotations
12. import datetime as dt
13. import tempfile
14. from dataclasses import dataclass
15. from pathlib import Path
16. from typing import Callable
18. from . import audio, render
19. from .engines.local import DEFAULT_MODEL
20. from .filing import write_unique
21. from .shape import rename_speakers, sentences_from_segments
23. Transcriber = Callable[[Path, Path, "Options"], dict] # (wav16k, source, opts) -> result
24. Aligner = Callable[[dict, Path, "Options"], dict] # (result, wav16k, opts) -> result with words
25. Diarizer = Callable[[dict, Path, "Options"], dict] # (result, wav16k, opts) -> result with speakers
28. @dataclass
29. class Options:
30. engine: str = "local" # local | deepinfra
31. language: str | None = None
32. diarize: bool = True
33. min_speakers: int | None = None
34. max_speakers: int | None = None
35. formats: tuple[str, ...] = ("txt", "md", "json")
36. timestamps: bool = True
37. model: str = DEFAULT_MODEL
40. @dataclass
41. class Result:
42. outputs: dict[str, Path]
43. speakers: list[str]
44. duration: float
45. sentences: int
46. language: str | None
49. def process(source: Path, out_dir: Path, opts: Options, transcriber: Transcriber,
50. aligner: Aligner | None, diarizer: Diarizer | None,
51. vault_inbox: Path | None = None) -> Result:
52. source = Path(source)
53. with tempfile.TemporaryDirectory(prefix="tp-") as tmp:
54. wav = audio.to_wav16k(source, Path(tmp) / "audio.wav")
55. duration = audio.wav_duration(wav)
56. result = transcriber(wav, source, opts)
57. if aligner is not None and result.get("segments"):
58. result = aligner(result, wav, opts)
59. if opts.diarize and diarizer is not None and result.get("segments"):
60. result = diarizer(result, wav, opts)
62. sentences, speakers = rename_speakers(sentences_from_segments(result.get("segments", [])))
63. stem = source.stem
64. recorded = dt.datetime.fromtimestamp(source.stat().st_mtime).date().isoformat() \
65. if source.exists() else dt.date.today().isoformat()
66. meta = {"title": stem, "source": source.name, "recorded": recorded,
67. "transcribed": dt.datetime.now().isoformat(timespec="seconds"),
68. "duration": audio.hms(duration), "engine": opts.engine,
69. "model": result.get("model"), "language": result.get("language"),
70. "speakers": speakers, "diarized": bool(opts.diarize and diarizer is not None)}
72. outputs: dict[str, Path] = {}
73. writers = {
74. "txt": lambda: render.to_txt(sentences, opts.timestamps),
75. "md": lambda: render.to_markdown(sentences, meta),
76. "json": lambda: render.to_json(sentences, result, meta),
77. "srt": lambda: render.to_srt(sentences),
78. "vtt": lambda: render.to_vtt(sentences),
79. }
80. for fmt in opts.formats:
81. if fmt not in render.FORMATS:
82. *rest, last = render.FORMATS
83. raise ValueError(f"Unknown format {fmt!r}: use {', '.join(rest)} or {last}.")
84. outputs[fmt] = write_unique(out_dir, f"{stem}.{fmt}", writers[fmt]())
85. if vault_inbox is not None:
86. outputs["vault"] = write_unique(vault_inbox, f"{stem}.md",
87. render.to_markdown(sentences, meta))
88. return Result(outputs, speakers, duration, len(sentences), result.get("language"))