shape.py

transcription_pipeline/shape.py

1. """From timed, speaker-tagged words (WhisperX segments) to sentences and speaker turns.
3. A sentence starts at its first aligned word. With none (no alignment model for the
4. language, or a token such as a number that never aligns) its time is interpolated
5. across the segment by character position. A sentence is also split where the
6. speaker changes, since speech is not always punctuated at a change of voice; a run
7. shorter than MIN_RUN_WORDS counts as diarization noise. Raw labels become
8. "Speaker 1", "Speaker 2", ... in order of first appearance.
9. """
11. from __future__ import annotations
13. from collections import defaultdict
15. from .sentences import split_sentences
17. # A speaker run shorter than this is treated as diarization noise.
18. MIN_RUN_WORDS = 2
21. def _num(x):
22. try:
23. v = float(x)
24. except (TypeError, ValueError):
25. return None
26. return v if v == v else None # drop NaN
29. def _word_offsets(text: str, words: list[dict]) -> list[int | None]:
30. """Locate each word in the segment text, in order. None when not found."""
31. offsets, cursor = [], 0
32. low = text.lower()
33. for w in words:
34. token = str(w.get("word", "")).strip().lower()
35. if not token:
36. offsets.append(None)
37. continue
38. at = low.find(token, cursor)
39. if at < 0:
40. offsets.append(None)
41. continue
42. offsets.append(at)
43. cursor = at + len(token)
44. return offsets
47. def _majority_speaker(words: list[dict]) -> str | None:
48. weight: dict[str, float] = defaultdict(float)
49. first_seen: dict[str, int] = {}
50. for i, w in enumerate(words):
51. spk = w.get("speaker")
52. if not spk:
53. continue
54. s, e = _num(w.get("start")), _num(w.get("end"))
55. weight[spk] += (e - s) if (s is not None and e is not None and e > s) else 0.01
56. first_seen.setdefault(spk, i)
57. if not weight:
58. return None
59. # Highest total duration wins; a tie goes to whoever spoke first.
60. return max(weight, key=lambda k: (weight[k], -first_seen[k]))
63. def _speaker_runs(items: list[tuple[int, dict]], min_words: int = MIN_RUN_WORDS) -> list[list]:
64. """Split (offset, word) pairs into runs of one speaker each.
66. Words with no speaker join the run they sit in. A run shorter than
67. `min_words` is folded into its neighbour, so a single misattributed word
68. does not break a sentence in two.
69. """
70. runs: list[list] = []
71. for off, w in items:
72. spk = w.get("speaker")
73. if runs and (spk is None or spk == runs[-1][0] or runs[-1][0] is None):
74. if runs[-1][0] is None:
75. runs[-1][0] = spk
76. runs[-1][1].append((off, w))
77. else:
78. runs.append([spk, [(off, w)]])
79. changed = True
80. while changed and len(runs) > 1:
81. changed = False
82. for i, (_, words) in enumerate(runs):
83. if len(words) < min_words:
84. j = i - 1 if i > 0 else i + 1
85. if i > 0:
86. runs[j][1].extend(words)
87. else:
88. runs[j][1][:0] = words
89. del runs[i]
90. # Neighbours may now share a speaker: merge them.
91. k = 0
92. while k + 1 < len(runs):
93. if runs[k][0] == runs[k + 1][0]:
94. runs[k][1].extend(runs.pop(k + 1)[1])
95. else:
96. k += 1
97. changed = True
98. break
99. return runs
102. def _piece(text: str, words: list[dict], start_fb: float, end_fb: float, speaker) -> dict:
103. starts = [_num(w.get("start")) for w in words if _num(w.get("start")) is not None]
104. ends = [_num(w.get("end")) for w in words if _num(w.get("end")) is not None]
105. start, timed = (min(starts), True) if starts else (start_fb, False)
106. end = max(ends) if ends else end_fb
107. return {"text": text, "start": round(start, 3), "end": round(max(end, start), 3),
108. "speaker": speaker, "timed": timed}
111. def sentences_from_segments(segments: list[dict]) -> list[dict]:
112. """[{text, start, end, speaker (raw label or None), timed (bool)}] in order."""
113. out: list[dict] = []
114. for seg in segments:
115. text = (seg.get("text") or "").strip()
116. if not text:
117. continue
118. seg_start = _num(seg.get("start")) or 0.0
119. seg_end = _num(seg.get("end"))
120. span = seg_end - seg_start if (seg_end is not None and seg_end > seg_start) else 0.0
121. words = seg.get("words") or []
122. offsets = _word_offsets(text, words)
123. sentences = split_sentences(text)
125. def at(pos: int) -> float:
126. return seg_start + span * min(pos, len(text)) / max(len(text), 1)
128. for k, s in enumerate(sentences):
129. lo = s["offset"]
130. hi = sentences[k + 1]["offset"] if k + 1 < len(sentences) else len(text) + 1
131. mine = [(off, w) for w, off in zip(words, offsets)
132. if off is not None and lo <= off < hi]
133. runs = _speaker_runs(mine) if any(w.get("speaker") for _, w in mine) else []
134. if len(runs) <= 1:
135. speaker = _majority_speaker([w for _, w in mine]) or seg.get("speaker")
136. out.append(_piece(s["text"], [w for _, w in mine], at(lo),
137. at(hi) if span else at(lo), speaker))
138. continue
139. # Cut the sentence where each later run's first word begins.
140. cuts = [lo] + [r[1][0][0] for r in runs[1:]] + [lo + len(s["text"])]
141. for (spk, run), a, b in zip(runs, cuts, cuts[1:]):
142. piece = text[a:b].strip()
143. if piece:
144. out.append(_piece(piece, [w for _, w in run], at(a), at(b), spk))
145. return out
148. def rename_speakers(sentences: list[dict]) -> tuple[list[dict], list[str]]:
149. """Relabel raw speakers as 'Speaker N' by first appearance. Returns (sentences, names)."""
150. mapping: dict[str, str] = {}
151. for s in sentences:
152. raw = s.get("speaker")
153. if raw and raw not in mapping:
154. mapping[raw] = f"Speaker {len(mapping) + 1}"
155. renamed = [{**s, "speaker": mapping.get(s.get("speaker"))} for s in sentences]
156. return renamed, list(mapping.values())
159. def turns(sentences: list[dict]) -> list[dict]:
160. """Merge consecutive sentences by the same speaker into turns."""
161. out: list[dict] = []
162. for s in sentences:
163. if out and out[-1]["speaker"] == s["speaker"]:
164. out[-1]["sentences"].append(s)
165. out[-1]["end"] = s["end"]
166. else:
167. out.append({"speaker": s["speaker"], "start": s["start"], "end": s["end"],
168. "sentences": [s]})
169. return out