|
@@ -197,35 +197,35 @@ def transcribe(
|
|
timestamp_tokens: torch.Tensor = tokens.ge(tokenizer.timestamp_begin)
|
|
timestamp_tokens: torch.Tensor = tokens.ge(tokenizer.timestamp_begin)
|
|
consecutive = torch.where(timestamp_tokens[:-1] & timestamp_tokens[1:])[0].add_(1)
|
|
consecutive = torch.where(timestamp_tokens[:-1] & timestamp_tokens[1:])[0].add_(1)
|
|
if len(consecutive) > 0: # if the output contains two consecutive timestamp tokens
|
|
if len(consecutive) > 0: # if the output contains two consecutive timestamp tokens
|
|
|
|
+ if ended_with_single_timestamp := timestamp_tokens[-2:].tolist() == [False, True]:
|
|
|
|
+ consecutive = consecutive.tolist() + [len(tokens)]
|
|
last_slice = 0
|
|
last_slice = 0
|
|
for current_slice in consecutive:
|
|
for current_slice in consecutive:
|
|
sliced_tokens = tokens[last_slice:current_slice]
|
|
sliced_tokens = tokens[last_slice:current_slice]
|
|
- start_timestamp_position = (
|
|
|
|
- sliced_tokens[0].item() - tokenizer.timestamp_begin
|
|
|
|
- )
|
|
|
|
- end_timestamp_position = (
|
|
|
|
- sliced_tokens[-1].item() - tokenizer.timestamp_begin
|
|
|
|
- )
|
|
|
|
|
|
+ start_timestamp_pos = sliced_tokens[0].item() - tokenizer.timestamp_begin
|
|
|
|
+ end_timestamp_pos = sliced_tokens[-1].item() - tokenizer.timestamp_begin
|
|
add_segment(
|
|
add_segment(
|
|
- start=timestamp_offset + start_timestamp_position * time_precision,
|
|
|
|
- end=timestamp_offset + end_timestamp_position * time_precision,
|
|
|
|
|
|
+ start=timestamp_offset + start_timestamp_pos * time_precision,
|
|
|
|
+ end=timestamp_offset + end_timestamp_pos * time_precision,
|
|
text_tokens=sliced_tokens[1:-1],
|
|
text_tokens=sliced_tokens[1:-1],
|
|
result=result,
|
|
result=result,
|
|
)
|
|
)
|
|
last_slice = current_slice
|
|
last_slice = current_slice
|
|
- last_timestamp_position = (
|
|
|
|
- tokens[last_slice - 1].item() - tokenizer.timestamp_begin
|
|
|
|
- )
|
|
|
|
- seek += last_timestamp_position * input_stride
|
|
|
|
|
|
+ if ended_with_single_timestamp:
|
|
|
|
+ # single timestamp at the end means no speech after the last timestamp.
|
|
|
|
+ seek += segment.shape[-1]
|
|
|
|
+ else:
|
|
|
|
+ # otherwise, ignore the unfinished segment and seek to the last timestamp
|
|
|
|
+ last_timestamp_pos = tokens[last_slice - 1].item() - tokenizer.timestamp_begin
|
|
|
|
+ seek += last_timestamp_pos * input_stride
|
|
all_tokens.extend(tokens[: last_slice + 1].tolist())
|
|
all_tokens.extend(tokens[: last_slice + 1].tolist())
|
|
else:
|
|
else:
|
|
duration = segment_duration
|
|
duration = segment_duration
|
|
timestamps = tokens[timestamp_tokens.nonzero().flatten()]
|
|
timestamps = tokens[timestamp_tokens.nonzero().flatten()]
|
|
if len(timestamps) > 0 and timestamps[-1].item() != tokenizer.timestamp_begin:
|
|
if len(timestamps) > 0 and timestamps[-1].item() != tokenizer.timestamp_begin:
|
|
# no consecutive timestamps but it has a timestamp; use the last one.
|
|
# no consecutive timestamps but it has a timestamp; use the last one.
|
|
- # single timestamp at the end means no speech after the last timestamp.
|
|
|
|
- last_timestamp_position = timestamps[-1].item() - tokenizer.timestamp_begin
|
|
|
|
- duration = last_timestamp_position * time_precision
|
|
|
|
|
|
+ last_timestamp_pos = timestamps[-1].item() - tokenizer.timestamp_begin
|
|
|
|
+ duration = last_timestamp_pos * time_precision
|
|
|
|
|
|
add_segment(
|
|
add_segment(
|
|
start=timestamp_offset,
|
|
start=timestamp_offset,
|