Spaces:
Running on Zero
Running on Zero
| """ | |
| Align Whisper words with speaker diarization turns. | |
| Strategy: give every word the speaker whose turns overlap it the most, tidy | |
| the speaker changes so they fall on sentence boundaries, then join consecutive | |
| words from the same speaker into readable paragraphs. | |
| """ | |
| from __future__ import annotations | |
| from typing import List | |
| from modules.models import ( | |
| AlignedSegment, | |
| DiarizationSegment, | |
| TranscriptionSegment, | |
| ) | |
| # A pause this long starts a new paragraph even if the speaker is the same. | |
| MAX_PAUSE_SECONDS = 1.5 | |
| # After this long, a paragraph ends at the next sentence boundary. | |
| PARAGRAPH_SECONDS = 20.0 | |
| SENTENCE_ENDINGS = (".", "?", "!", "।", "。") | |
| # Word timestamps and turn edges are each off by a few tenths of a second, so | |
| # up to this many words on either side of a speaker change may be mislabelled. | |
| BOUNDARY_WORDS = 3 | |
| def align_transcript_with_speakers( | |
| transcription_segments: List[TranscriptionSegment], | |
| diarization_segments: List[DiarizationSegment], | |
| ) -> List[AlignedSegment]: | |
| """ | |
| Merge word timestamps with speaker labels. | |
| *transcription_segments* are words as produced by TranscriberService | |
| (leading space kept), so their text is concatenated as-is. | |
| """ | |
| words = transcription_segments | |
| speakers = [_find_dominant_speaker(w, diarization_segments) for w in words] | |
| _snap_to_sentences(words, speakers) | |
| aligned: list[AlignedSegment] = [] | |
| for word, speaker in zip(words, speakers): | |
| current = aligned[-1] if aligned else None | |
| if ( | |
| current | |
| and current.speaker == speaker | |
| and word.start - current.end < MAX_PAUSE_SECONDS | |
| and not ( | |
| current.end - current.start > PARAGRAPH_SECONDS | |
| and _ends_sentence(current.text) | |
| ) | |
| ): | |
| current.text += word.text | |
| current.end = word.end | |
| else: | |
| aligned.append( | |
| AlignedSegment( | |
| start=word.start, | |
| end=word.end, | |
| speaker=speaker, | |
| text=word.text, | |
| ) | |
| ) | |
| for seg in aligned: | |
| seg.text = seg.text.strip() | |
| return aligned | |
| # --------------------------------------------------------------------------- | |
| # Internal helpers | |
| # --------------------------------------------------------------------------- | |
| def _ends_sentence(text: str) -> bool: | |
| return text.rstrip().endswith(SENTENCE_ENDINGS) | |
| def _snap_to_sentences(words: List[TranscriptionSegment], speakers: List[str]) -> None: | |
| """ | |
| Move speaker changes that land mid-sentence onto the nearest sentence | |
| boundary, editing *speakers* in place. People mostly change at sentence | |
| ends; a change one or two words off is timing error, not an interruption. | |
| ponytail: relies on Whisper's punctuation, so it does nothing for | |
| unpunctuated output, and a genuine one-word interruption without | |
| punctuation is folded into the surrounding speaker. | |
| """ | |
| n = len(words) | |
| # 1. A few words of one speaker inside another speaker's sentence | |
| # ("... prohibition, [such] as Andhra Pradesh"). | |
| i = 0 | |
| while i < n: | |
| j = i | |
| while j + 1 < n and speakers[j + 1] == speakers[i]: | |
| j += 1 | |
| if ( | |
| 0 < i and j < n - 1 | |
| and j - i + 1 < BOUNDARY_WORDS | |
| and speakers[i - 1] == speakers[j + 1] | |
| and not _ends_sentence(words[i - 1].text) | |
| and not _ends_sentence(words[j].text) | |
| ): | |
| speakers[i: j + 1] = [speakers[i - 1]] * (j - i + 1) | |
| i = j + 1 | |
| # 2. A change just before or just after a sentence end. | |
| for i in range(1, n): | |
| previous, current = speakers[i - 1], speakers[i] | |
| if previous == current or _ends_sentence(words[i - 1].text): | |
| continue | |
| # The sentence finishes within the new speaker's first words: they | |
| # belong to the previous speaker ("... still is | banned. When was"). | |
| run_end = i | |
| while run_end + 1 < n and speakers[run_end + 1] == current: | |
| run_end += 1 | |
| head = next( | |
| (j for j in range(i, min(i + BOUNDARY_WORDS, run_end + 1)) if _ends_sentence(words[j].text)), | |
| None, | |
| ) | |
| if head is not None and head < run_end: | |
| speakers[i: head + 1] = [previous] * (head + 1 - i) | |
| continue | |
| # The previous speaker's last words start a new sentence: they belong | |
| # to the new speaker ("... New Jersey. And I'm | Sheila in Texas"). | |
| tail = i - 1 | |
| while tail > 0 and speakers[tail - 1] == previous and not _ends_sentence(words[tail - 1].text): | |
| tail -= 1 | |
| if ( | |
| i - tail <= BOUNDARY_WORDS | |
| and tail > 0 | |
| and speakers[tail - 1] == previous | |
| ): | |
| speakers[tail:i] = [current] * (i - tail) | |
| def _find_dominant_speaker( | |
| t_seg: TranscriptionSegment, | |
| d_segments: List[DiarizationSegment], | |
| ) -> str: | |
| """ | |
| Return the speaker with the most overlap for *t_seg*. A word that falls | |
| in a gap between turns goes to the nearest turn. | |
| """ | |
| if not d_segments: | |
| return "Speaker 1" | |
| # ponytail: scans every turn for every word, O(words x turns). Fine up to | |
| # a few hours of audio; switch to a two-pointer sweep if it ever shows up. | |
| speaker_overlap: dict[str, float] = {} | |
| for d_seg in d_segments: | |
| overlap = _calculate_overlap( | |
| t_seg.start, t_seg.end, d_seg.start, d_seg.end | |
| ) | |
| if overlap > 0: | |
| speaker_overlap[d_seg.speaker] = ( | |
| speaker_overlap.get(d_seg.speaker, 0.0) + overlap | |
| ) | |
| if speaker_overlap: | |
| return max(speaker_overlap, key=speaker_overlap.get) | |
| nearest = min( | |
| d_segments, | |
| key=lambda d: max(d.start - t_seg.end, t_seg.start - d.end), | |
| ) | |
| return nearest.speaker | |
| def _calculate_overlap( | |
| start1: float, end1: float, start2: float, end2: float | |
| ) -> float: | |
| """Return the overlap duration between two intervals (≥ 0).""" | |
| return max(0.0, min(end1, end2) - max(start1, start2)) | |