[whisper] fix decoding maxlen (#2380)

wenet-e2e · Mar 4, 2024 · 8f8cedc · 8f8cedc
1 parent 8179fe1
commit 8f8cedc
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/wenet/transformer/search.py b/wenet/transformer/search.py
@@ -292,6 +292,8 @@ def attention_beam_search(
     cache: Optional[List[torch.Tensor]] = None
     if model.decoder.use_sdpa:
         encoder_mask = mask_to_bias(encoder_mask, encoder_out.dtype)
+    if hasattr(model, 'decode_maxlen'):
+        maxlen = model.decode_maxlen
     # 2. Decoder forward step by step
     for i in range(prefix_len, maxlen + 1):
         # Stop if all batch and all beam produce eos

diff --git a/wenet/whisper/whisper.py b/wenet/whisper/whisper.py
@@ -46,6 +46,7 @@ def __init__(
         assert reverse_weight == 0.0
         self.sos = special_tokens["sot"]
         self.eos = special_tokens["eot"]
+        self.decode_maxlen = self.decoder.embed[1].max_len
 
     # TODO(xcsong): time align
     def set_alignment_heads(self, dump: bytes):