100 lines
3.7 KiB
Python
100 lines
3.7 KiB
Python
from whisperlivekit.timed_objects import Line, format_time, SpeakerSegment, Silence
|
|
from time import time
|
|
|
|
|
|
class TokensAlignment:
|
|
|
|
def __init__(self, state, args, sep):
|
|
self.state = state
|
|
self.diarization = args.diarization
|
|
self._tokens_index = 0
|
|
self._diarization_index = 0
|
|
self._translation_index = 0
|
|
|
|
self.all_tokens = []
|
|
self.all_diarization_segments = []
|
|
self.all_translation_segments = []
|
|
|
|
self.new_tokens = []
|
|
self.new_translation = []
|
|
self.new_diarization = []
|
|
self.new_tokens_buffer = []
|
|
self.sep = ' '
|
|
|
|
def update(self):
|
|
self.new_tokens, self.state.new_tokens = self.state.new_tokens, []
|
|
self.new_diarization, self.state.new_diarization = self.state.new_diarization, []
|
|
self.new_translation, self.state.new_translation = self.state.new_translation, []
|
|
self.new_tokens_buffer, self.state.new_tokens_buffer = self.state.new_tokens_buffer, []
|
|
|
|
self.all_tokens.extend(self.new_tokens)
|
|
self.all_diarization_segments.extend(self.new_diarization)
|
|
self.all_translation_segments.extend(self.new_translation)
|
|
|
|
def create_lines_from_tokens(self, current_silence, beg_loop):
|
|
lines = []
|
|
current_line_tokens = []
|
|
for token in self.all_tokens:
|
|
if type(token) == Silence:
|
|
if current_line_tokens:
|
|
lines.append(Line().build_from_tokens(current_line_tokens))
|
|
current_line_tokens = []
|
|
end_silence = token.end if token.has_ended else time() - beg_loop
|
|
if lines and lines[-1].speaker == -2:
|
|
lines[-1].end = end_silence
|
|
else:
|
|
lines.append(Line(
|
|
speaker = -2,
|
|
text = '',
|
|
start = token.start,
|
|
end = end_silence
|
|
))
|
|
else:
|
|
current_line_tokens.append(token)
|
|
if current_line_tokens:
|
|
lines.append(Line().build_from_tokens(current_line_tokens))
|
|
if current_silence:
|
|
end_silence = current_silence.end if current_silence.has_ended else time() - beg_loop
|
|
if lines and lines[-1].speaker == -2:
|
|
lines[-1].end = end_silence
|
|
else:
|
|
lines.append(Line(
|
|
speaker = -2,
|
|
text = '',
|
|
start = current_silence.start,
|
|
end = end_silence
|
|
))
|
|
|
|
return lines
|
|
|
|
def align_tokens(self):
|
|
if not self.diarization:
|
|
pass
|
|
# return self.all_tokens
|
|
|
|
def compute_punctuations_segments(self):
|
|
punctuations_breaks = []
|
|
new_tokens = self.state.tokens[self.state.last_validated_token:]
|
|
for i in range(len(new_tokens)):
|
|
token = new_tokens[i]
|
|
if token.is_punctuation():
|
|
punctuations_breaks.append({
|
|
'token_index': i,
|
|
'token': token,
|
|
'start': token.start,
|
|
'end': token.end,
|
|
})
|
|
punctuations_segments = []
|
|
for i, break_info in enumerate(punctuations_breaks):
|
|
start = punctuations_breaks[i - 1]['end'] if i > 0 else 0.0
|
|
end = break_info['end']
|
|
punctuations_segments.append({
|
|
'start': start,
|
|
'end': end,
|
|
'token_index': break_info['token_index'],
|
|
'token': break_info['token']
|
|
})
|
|
return punctuations_segments
|
|
|
|
def concatenate_diar_segments(self):
|
|
diarization_segments = self.state.diarization_segments
|