Decoding improvements (openai#1033)

* suppress task tokens (transcribe/translate)

* not ignoring the last segment ending with one timestamp

Author: jongwook

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

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|>")

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]