Update decoding.py (openai#1155)

* Update decoding.py Following the suggestions of @Jeronymous in openai#914 and openai#924, it solves the problem of endless loop. * Removed blank line and whitespaces in empty lines. * Suggested changes according to the linter --------- Co-authored-by: Jong Wook Kim <[email protected]>
zackees · May 5, 2023 · 8179d68 · 8179d68
1 parent 09da266
commit 8179d68
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/whisper/decoding.py b/whisper/decoding.py
@@ -471,6 +471,13 @@ def apply(self, logits: Tensor, tokens: Tensor):
  # timestamps shouldn't decrease; forbid timestamp tokens smaller than the last
  logits[k, self.tokenizer.timestamp_begin : timestamps[-1]] = -np.inf
 
+ # to force that timestamps are strictly increasing
+ if last_was_timestamp and not penultimate_was_timestamp:
+ timestamp_last = timestamps[-1]
+ else:
+ timestamp_last = timestamps[-1] + 1
+ logits[k, self.tokenizer.timestamp_begin : timestamp_last] = -np.inf
+
  if tokens.shape[1] == self.sample_begin:
  # suppress generating non-timestamp tokens at the beginning
  logits[:, : self.tokenizer.timestamp_begin] = -np.inf