Mirroristas/whisper
2
0
Fork 0
mirror of https://github.com/openai/whisper.git synced 2025-11-24 14:35:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

Decoding improvements (#1033)

Browse Source 
* suppress task tokens (transcribe/translate)

* not ignoring the last segment ending with one timestamp
This commit is contained in:
Jong Wook Kim 2023-03-06 14:32:32 -05:00 committed by GitHub
parent 3e1780fd37
commit eab8d920ed
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 30 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
whisper/decoding.py
View File

@ -549,7 +549,13 @@ class DecodingTask:
assert isinstance(suppress_tokens, list), "suppress_tokens must be a list" assert isinstance(suppress_tokens, list), "suppress_tokens must be a list"
suppress_tokens.extend( suppress_tokens.extend(
[self.tokenizer.sot, self.tokenizer.sot_prev, self.tokenizer.sot_lm] [
self.tokenizer.transcribe,
self.tokenizer.translate,
self.tokenizer.sot,
self.tokenizer.sot_prev,
self.tokenizer.sot_lm
]
) )
if self.tokenizer.no_speech is not None: if self.tokenizer.no_speech is not None:
# no-speech probability is collected separately # no-speech probability is collected separately

8
whisper/tokenizer.py
View File

@ -160,6 +160,14 @@ class Tokenizer:
def eot(self) -> int: def eot(self) -> int:
return self.tokenizer.eos_token_id return self.tokenizer.eos_token_id
@cached_property
def transcribe(self) -> int:
return self._get_single_token_id("<|transcribe|>")
@cached_property
def translate(self) -> int:
return self._get_single_token_id("<|translate|>")
@cached_property @cached_property
def sot(self) -> int: def sot(self) -> int:
return self._get_single_token_id("<|startoftranscript|>") return self._get_single_token_id("<|startoftranscript|>")

30
whisper/transcribe.py
View File

@ -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 = ( start_timestamp_pos = sliced_tokens[0].item() - tokenizer.timestamp_begin
sliced_tokens[0].item() - tokenizer.timestamp_begin end_timestamp_pos = sliced_tokens[-1].item() - tokenizer.timestamp_begin
)
end_timestamp_position = (
sliced_tokens[-1].item() - tokenizer.timestamp_begin
)
add_segment( add_segment(
start=timestamp_offset + start_timestamp_position * time_precision, start=timestamp_offset + start_timestamp_pos * time_precision,
end=timestamp_offset + end_timestamp_position * 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 = ( if ended_with_single_timestamp:
tokens[last_slice - 1].item() - tokenizer.timestamp_begin # single timestamp at the end means no speech after the last timestamp.
) seek += segment.shape[-1]
seek += last_timestamp_position * input_stride 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_pos = timestamps[-1].item() - tokenizer.timestamp_begin
last_timestamp_position = timestamps[-1].item() - tokenizer.timestamp_begin duration = last_timestamp_pos * time_precision
duration = last_timestamp_position * time_precision
add_segment( add_segment(
start=timestamp_offset, start=timestamp_offset,