Decoding improvements (#1033)

* suppress task tokens (transcribe/translate)

* not ignoring the last segment ending with one timestamp

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

@@ -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

@@ -197,35 +197,35 @@ 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 = (
- 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(
- 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,
  )
  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
 
  add_segment(
  start=timestamp_offset,