Merge branch 'main' into word-level-timestamps

README.md

@@ -5,14 +5,14 @@
 [[Model card]](https://github.com/openai/whisper/blob/main/model-card.md)
 [[Colab example]](https://colab.research.google.com/github/openai/whisper/blob/master/notebooks/LibriSpeech.ipynb)
 
-Whisper is a general-purpose speech recognition model. It is trained on a large dataset of diverse audio and is also a multi-task model that can perform multilingual speech recognition as well as speech translation and language identification.
+Whisper is a general-purpose speech recognition model. It is trained on a large dataset of diverse audio and is also a multitasking model that can perform multilingual speech recognition, speech translation, and language identification.
 
 
 ## Approach
 
 ![Approach](https://raw.githubusercontent.com/openai/whisper/main/approach.png)
 
-A Transformer sequence-to-sequence model is trained on various speech processing tasks, including multilingual speech recognition, speech translation, spoken language identification, and voice activity detection. All of these tasks are jointly represented as a sequence of tokens to be predicted by the decoder, allowing for a single model to replace many different stages of a traditional speech processing pipeline. The multitask training format uses a set of special tokens that serve as task specifiers or classification targets.
+A Transformer sequence-to-sequence model is trained on various speech processing tasks, including multilingual speech recognition, speech translation, spoken language identification, and voice activity detection. These tasks are jointly represented as a sequence of tokens to be predicted by the decoder, allowing a single model to replace many stages of a traditional speech-processing pipeline. The multitask training format uses a set of special tokens that serve as task specifiers or classification targets.
 
 
 ## Setup
@@ -68,9 +68,9 @@ There are five model sizes, four with English-only versions, offering speed and
 | medium | 769 M | `medium.en` | `medium` | ~5 GB | ~2x |
 | large | 1550 M | N/A | `large` | ~10 GB | 1x |
 
-For English-only applications, the `.en` models tend to perform better, especially for the `tiny.en` and `base.en` models. We observed that the difference becomes less significant for the `small.en` and `medium.en` models.
+The `.en` models for English-only applications tend to perform better, especially for the `tiny.en` and `base.en` models. We observed that the difference becomes less significant for the `small.en` and `medium.en` models.
 
-Whisper's performance varies widely depending on the language. The figure below shows a WER (Word Error Rate) breakdown by languages of Fleurs dataset, using the `large-v2` model. More WER and BLEU scores corresponding to the other models and datasets can be found in Appendix D in [the paper](https://arxiv.org/abs/2212.04356). The smaller is better.
+Whisper's performance varies widely depending on the language. The figure below shows a WER (Word Error Rate) breakdown by languages of the Fleurs dataset using the `large-v2` model. More WER and BLEU scores corresponding to the other models and datasets can be found in Appendix D in [the paper](https://arxiv.org/abs/2212.04356). The smaller, the better.
 
 ![WER breakdown by language](https://raw.githubusercontent.com/openai/whisper/main/language-breakdown.svg)
 
@@ -144,4 +144,4 @@ Please use the [🙌 Show and tell](https://github.com/openai/whisper/discussion
 
 ## License
 
-The code and the model weights of Whisper are released under the MIT License. See [LICENSE](https://github.com/openai/whisper/blob/main/LICENSE) for further details.
+Whisper's code and model weights are released under the MIT License. See [LICENSE](https://github.com/openai/whisper/blob/main/LICENSE) for further details.

whisper/decoding.py

@@ -549,7 +549,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

@@ -161,6 +161,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

@@ -216,36 +216,36 @@ def new_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 = (
- 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
  current_segments.append(new_segment(
- start=time_offset + start_timestamp_position * time_precision,
- end=time_offset + end_timestamp_position * time_precision,
+ start=timestamp_offset + start_timestamp_pos * time_precision,
+ end=timestamp_offset + end_timestamp_pos * time_precision,
  tokens=sliced_tokens,
  result=result,
  ))
  current_tokens.append(sliced_tokens.tolist())
  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 = 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
 
  current_segments.append(new_segment(
  start=time_offset,