heimoshuiyu/whisper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

suppress generating non-timestamp tokens at the beginning (#532)

Browse Source
This commit is contained in:
jumon
2022-11-16 04:44:36 +09:00
committed by GitHub
parent 9f70a352f9
commit 76148a56c5
1 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
whisper/decoding.py
View File

@@ -423,10 +423,14 @@ class ApplyTimestampRules(LogitFilter):
else: # cannot be normal text tokens else: # cannot be normal text tokens
logits[k, : self.tokenizer.eot] = -np.inf logits[k, : self.tokenizer.eot] = -np.inf
# apply the `max_initial_timestamp` option if tokens.shape[1] == self.sample_begin:
if tokens.shape[1] == self.sample_begin and self.max_initial_timestamp_index is not None: # suppress generating non-timestamp tokens at the beginning
last_allowed = self.tokenizer.timestamp_begin + self.max_initial_timestamp_index logits[:, : self.tokenizer.timestamp_begin] = -np.inf
logits[:, last_allowed + 1 :] = -np.inf
# apply the `max_initial_timestamp` option
if self.max_initial_timestamp_index is not None:
last_allowed = self.tokenizer.timestamp_begin + self.max_initial_timestamp_index
logits[:, last_allowed + 1 :] = -np.inf
# if sum of probability over timestamps is above any other token, sample timestamp # if sum of probability over timestamps is above any other token, sample timestamp
logprobs = F.log_softmax(logits.float(), dim=-1) logprobs = F.log_softmax(logits.float(), dim=-1)