Update modeling_conformer.py

modeling_conformer.py

@@ -61,10 +61,9 @@ class Wav2Vec2ConformerRNNT(Wav2Vec2ConformerModel):
         self.mask_layer.cache_pad_mask = (torch.arange(hidden_states.size(1), device=hidden_states.device).unsqueeze(0) >= self.cache_length.unsqueeze(1))
         return super()._mask_hidden_states(hidden_states, mask_time_indices, attention_mask)
 
-    def calc_length(self, lengths, all_paddings=2, kernel_size=3, stride=2, repeat_num=1):
-        add_pad = all_paddings - kernel_size
+    def calc_length(self, lengths, padding=1, kernel_size=3, stride=2, repeat_num=1):
         for _ in range(repeat_num):
-            lengths = (lengths + add_pad) // stride + 1
+            lengths = (lengths + 2 * padding - kernel_size) // stride + 1
         return lengths
 
     def preprocessing(self, x):
@@ -118,13 +117,11 @@ class Wav2Vec2ConformerRNNT(Wav2Vec2ConformerModel):
                 cx   = torch.where(mask, cx, cx_prev)
                 last = torch.where(emitted.unsqueeze(1), n.unsqueeze(1), last)
 
-                pos    = lengths.unsqueeze(1).clamp(max=max_len - 1)
-                fill_t = torch.where(emitted.unsqueeze(1), n.unsqueeze(1),    torch.full_like(n.unsqueeze(1), pad))
-                fill_s = torch.where(emitted.unsqueeze(1), t_sec,             torch.full_like(t_sec, -1.0))
-
-                tokens  = tokens.scatter(1, pos, fill_t)
-                starts  = starts.scatter(1, pos, fill_s)
-                lengths = lengths + emitted.long()
+                if emitted.any():
+                    idx    = lengths[emitted].unsqueeze(1).clamp(max=max_len - 1)
+                    tokens[emitted] = tokens[emitted].scatter(1, idx, n[emitted].unsqueeze(1))
+                    starts[emitted] = starts[emitted].scatter(1, idx, t_sec[emitted])
+                    lengths[emitted] += 1
 
         return tokens, starts, lengths