From f5040f950aef3de6439ee6e34530d76cd57b97df Mon Sep 17 00:00:00 2001 From: Jong Wook Kim Date: Mon, 6 Mar 2023 14:32:32 -0500 Subject: [PATCH] Decoding improvements (#1033) * suppress task tokens (transcribe/translate) * not ignoring the last segment ending with one timestamp --- whisper/decoding.py | 8 ++++++- whisper/tokenizer.py | 8 +++++++ whisper/transcribe.py | 52 +++++++++++++------------------------------ 3 files changed, 31 insertions(+), 37 deletions(-) diff --git a/whisper/decoding.py b/whisper/decoding.py index c0e3496c..50b93aa5 100644 --- a/whisper/decoding.py +++ b/whisper/decoding.py @@ -553,7 +553,13 @@ def _get_suppress_tokens(self) -> Tuple[int]: assert isinstance(suppress_tokens, list), "suppress_tokens must be a list" 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: # no-speech probability is collected separately diff --git a/whisper/tokenizer.py b/whisper/tokenizer.py index 7b4605f3..7efa2d4a 100644 --- a/whisper/tokenizer.py +++ b/whisper/tokenizer.py @@ -160,6 +160,14 @@ def decode_with_timestamps(self, tokens) -> str: def eot(self) -> int: 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 def sot(self) -> int: return self._get_single_token_id("<|startoftranscript|>") diff --git a/whisper/transcribe.py b/whisper/transcribe.py index bc910c69..d155b20d 100644 --- a/whisper/transcribe.py +++ b/whisper/transcribe.py @@ -146,7 +146,7 @@ def decode_with_fallback(segment: torch.Tensor) -> DecodingResult: initial_prompt_tokens = [] def add_segment( - *, start: float, end: float, text_tokens: torch.Tensor, result: DecodingResult, encoder_embeddings, decoder_embeddings + *, start: float, end: float, text_tokens: torch.Tensor, result: DecodingResult ): text = tokenizer.decode([token for token in text_tokens if token < tokenizer.eot]) if len(text.strip()) == 0: # skip empty text output @@ -164,8 +164,6 @@ def add_segment( "avg_logprob": result.avg_logprob, "compression_ratio": result.compression_ratio, "no_speech_prob": result.no_speech_prob, - "encoder_embeddings":encoder_embeddings, - "decoder_embeddings":decoder_embeddings } ) if verbose: @@ -199,59 +197,41 @@ def add_segment( timestamp_tokens: torch.Tensor = tokens.ge(tokenizer.timestamp_begin) 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 ended_with_single_timestamp := timestamp_tokens[-2:].tolist() == [False, True]: + consecutive = consecutive.tolist() + [len(tokens)] last_slice = 0 for current_slice in consecutive: sliced_tokens = tokens[last_slice:current_slice] - start_timestamp_position = ( - sliced_tokens[0].item() - tokenizer.timestamp_begin - ) - end_timestamp_position = min( - sliced_tokens[-1].item() - tokenizer.timestamp_begin, - np.ceil((num_frames - seek) / input_stride) - 1 - ) - encoder_embeddings = result.encoder_embeddings[:, :, - start_timestamp_position:int(end_timestamp_position)] - decoder_embeddings = result.decoder_embeddings[:,:, int(last_slice)+1:int(current_slice)-1] - + start_timestamp_pos = sliced_tokens[0].item() - tokenizer.timestamp_begin + end_timestamp_pos = sliced_tokens[-1].item() - tokenizer.timestamp_begin 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], result=result, - encoder_embeddings=encoder_embeddings, - decoder_embeddings=decoder_embeddings ) 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()) else: duration = segment_duration timestamps = tokens[timestamp_tokens.nonzero().flatten()] if len(timestamps) > 0 and timestamps[-1].item() != tokenizer.timestamp_begin: # 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 = min( - timestamps[-1].item() - tokenizer.timestamp_begin, - np.ceil((num_frames - seek) / input_stride) - 1 - ) - duration = last_timestamp_position * time_precision + last_timestamp_pos = timestamps[-1].item() - tokenizer.timestamp_begin + duration = last_timestamp_pos * time_precision - start_timestamp_position = ( - timestamps[0].item() - tokenizer.timestamp_begin - ) - encoder_embeddings = result.encoder_embeddings[:, :, - start_timestamp_position:int(last_timestamp_position)] - decoder_embeddings = result.decoder_embeddings[:, :, 1:-1] add_segment( start=timestamp_offset, end=timestamp_offset + duration, text_tokens=tokens, result=result, - encoder_embeddings=encoder_embeddings, - decoder_embeddings=decoder_embeddings ) seek += segment.shape[-1]