File size: 6,068 Bytes
313814b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
94c7543
 
 
 
 
 
 
 
313814b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
# TODO: rename module
from __future__ import annotations

import re
from dataclasses import dataclass

from speaches.config import config


# TODO: use the `Segment` from `faster-whisper.transcribe` instead
@dataclass
class Segment:
    text: str
    start: float = 0.0
    end: float = 0.0

    @property
    def is_eos(self) -> bool:
        if self.text.endswith("..."):
            return False
        for punctuation_symbol in ".?!":
            if self.text.endswith(punctuation_symbol):
                return True
        return False

    def offset(self, seconds: float) -> None:
        self.start += seconds
        self.end += seconds


# TODO: use the `Word` from `faster-whisper.transcribe` instead
@dataclass
class Word(Segment):
    probability: float = 0.0

    @classmethod
    def common_prefix(cls, a: list[Word], b: list[Word]) -> list[Word]:
        i = 0
        while (
            i < len(a)
            and i < len(b)
            and canonicalize_word(a[i].text) == canonicalize_word(b[i].text)
        ):
            i += 1
        return a[:i]


class Transcription:
    def __init__(self, words: list[Word] = []) -> None:
        self.words: list[Word] = []
        self.extend(words)

    @property
    def text(self) -> str:
        return " ".join(word.text for word in self.words).strip()

    @property
    def start(self) -> float:
        return self.words[0].start if len(self.words) > 0 else 0.0

    @property
    def end(self) -> float:
        return self.words[-1].end if len(self.words) > 0 else 0.0

    @property
    def duration(self) -> float:
        return self.end - self.start

    def after(self, seconds: float) -> Transcription:
        return Transcription(
            words=[word for word in self.words if word.start > seconds]
        )

    def extend(self, words: list[Word]) -> None:
        self._ensure_no_word_overlap(words)
        self.words.extend(words)

    def _ensure_no_word_overlap(self, words: list[Word]) -> None:
        if len(self.words) > 0 and len(words) > 0:
            if (
                words[0].start + config.word_timestamp_error_margin
                <= self.words[-1].end
            ):
                raise ValueError(
                    f"Words overlap: {self.words[-1]} and {words[0]}. Error margin: {config.word_timestamp_error_margin}"
                )
        for i in range(1, len(words)):
            if words[i].start + config.word_timestamp_error_margin <= words[i - 1].end:
                raise ValueError(
                    f"Words overlap: {words[i - 1]} and {words[i]}. All words: {words}"
                )


def test_segment_is_eos():
    assert not Segment("Hello").is_eos
    assert not Segment("Hello...").is_eos
    assert Segment("Hello.").is_eos
    assert Segment("Hello!").is_eos
    assert Segment("Hello?").is_eos
    assert not Segment("Hello. Yo").is_eos
    assert not Segment("Hello. Yo...").is_eos
    assert Segment("Hello. Yo.").is_eos


def to_full_sentences(words: list[Word]) -> list[Segment]:
    sentences: list[Segment] = [Segment("")]
    for word in words:
        sentences[-1] = Segment(
            start=sentences[-1].start,
            end=word.end,
            text=sentences[-1].text + word.text,
        )
        if word.is_eos:
            sentences.append(Segment(""))
    if len(sentences) > 0 and not sentences[-1].is_eos:
        sentences.pop()
    return sentences


def tests_to_full_sentences():
    assert to_full_sentences([]) == []
    assert to_full_sentences([Word(text="Hello")]) == []
    assert to_full_sentences([Word(text="Hello..."), Word(" world")]) == []
    assert to_full_sentences([Word(text="Hello..."), Word(" world.")]) == [
        Segment(text="Hello... world.")
    ]
    assert to_full_sentences(
        [Word(text="Hello..."), Word(" world."), Word(" How")]
    ) == [Segment(text="Hello... world.")]


def to_text(words: list[Word]) -> str:
    return "".join(word.text for word in words)


def to_text_w_ts(words: list[Word]) -> str:
    return "".join(f"{word.text}({word.start:.2f}-{word.end:.2f})" for word in words)


def canonicalize_word(text: str) -> str:
    text = text.lower()
    # Remove non-alphabetic characters using regular expression
    text = re.sub(r"[^a-z]", "", text)
    return text.lower().strip().strip(".,?!")


def test_canonicalize_word():
    assert canonicalize_word("ABC") == "abc"
    assert canonicalize_word("...ABC?") == "abc"
    assert canonicalize_word("... AbC  ...") == "abc"


def common_prefix(a: list[Word], b: list[Word]) -> list[Word]:
    i = 0
    while (
        i < len(a)
        and i < len(b)
        and canonicalize_word(a[i].text) == canonicalize_word(b[i].text)
    ):
        i += 1
    return a[:i]


def test_common_prefix():
    def word(text: str) -> Word:
        return Word(text=text, start=0.0, end=0.0, probability=0.0)

    a = [word("a"), word("b"), word("c")]
    b = [word("a"), word("b"), word("c")]
    assert common_prefix(a, b) == [word("a"), word("b"), word("c")]

    a = [word("a"), word("b"), word("c")]
    b = [word("a"), word("b"), word("d")]
    assert common_prefix(a, b) == [word("a"), word("b")]

    a = [word("a"), word("b"), word("c")]
    b = [word("a")]
    assert common_prefix(a, b) == [word("a")]

    a = [word("a")]
    b = [word("a"), word("b"), word("c")]
    assert common_prefix(a, b) == [word("a")]

    a = [word("a")]
    b = []
    assert common_prefix(a, b) == []

    a = []
    b = [word("a")]
    assert common_prefix(a, b) == []

    a = [word("a"), word("b"), word("c")]
    b = [word("b"), word("c")]
    assert common_prefix(a, b) == []


def test_common_prefix_and_canonicalization():
    def word(text: str) -> Word:
        return Word(text=text, start=0.0, end=0.0, probability=0.0)

    a = [word("A...")]
    b = [word("a?"), word("b"), word("c")]
    assert common_prefix(a, b) == [word("A...")]

    a = [word("A..."), word("B?"), word("C,")]
    b = [word("a??"), word("  b"), word(" ,c")]
    assert common_prefix(a, b) == [word("A..."), word("B?"), word("C,")]