Update modeling_conformer.py

modeling_conformer.py

@@ -1,6 +1,6 @@
 from huggingface_hub import hf_hub_download
 from torch import nn
-from transformers import Wav2Vec2ConformerModel , Wav2Vec2CTCTokenizer
+from transformers import Wav2Vec2ConformerModel
 from safetensors.torch import load_file
 from torch_state_bridge import state_bridge
 import torch
@@ -30,7 +30,7 @@ class Wav2Vec2ConformerRNNT(Wav2Vec2ConformerModel):
         self.joint = nn.Linear(config.joint_hidden, config.vocab_size // len(config.languages) + 1)
         self.embed = nn.Embedding(config.vocab_size+1, config.pred_hidden, padding_idx=config.vocab_size)
         self.lstm = nn.LSTM(config.pred_hidden, config.pred_hidden, config.lstm_layer, batch_first=True)
-        self.act = nn.ReLU(inplace=True)
+        self.act = nn.ReLU()
         self.spec = torchaudio.transforms.Spectrogram(n_fft=512, hop_length=160, win_length=400, center=False)
         self.mask_layer = Op(lambda self_obj,x : x.masked_fill(self_obj.cache_pad_mask.unsqueeze(1), 0),True)
         self.register_buffer("mel_fb",torch.tensor(librosa.filters.mel(sr=config.sampling_rate,n_fft=512,n_mels=80)))
@@ -51,7 +51,7 @@ class Wav2Vec2ConformerRNNT(Wav2Vec2ConformerModel):
         self.encoder.layer_norm = nn.Identity()
         if config.multilingual:
             self.lang_joint_net = nn.ModuleDict({l: nn.Linear(config.joint_hidden, config.vocab_size // len(config.languages) + 1) for l in config.languages})
-        self.preemph, self.eps, self.pad_to = 0.97, 2**-24, 16
+        self.eps = 2**-24
         self.denorm = (2 ** config.num_feat_extract_layers) * self.spec.hop_length / config.sampling_rate
         self.scaler = config.hidden_size ** (1/2)
         return super().init_weights()
@@ -70,7 +70,7 @@ class Wav2Vec2ConformerRNNT(Wav2Vec2ConformerModel):
     def preprocessing(self, x):
         x, l = x
         l = (l // self.spec.hop_length + 1).long()
-        x = torch.cat((x[:, :1], x[:, 1:] - self.preemph * x[:, :-1]), 1)
+        x = torch.cat((x[:, :1], x[:, 1:] - self.config.preemph * x[:, :-1]), 1)
         x = (self.mel_fb @ self.spec(x) + self.eps).log()
         T = x.size(-1)
         m = torch.arange(T, device=x.device)[None] >= l[:, None]
@@ -80,76 +80,51 @@ class Wav2Vec2ConformerRNNT(Wav2Vec2ConformerModel):
         σ = (((x - μ[..., None])**2).sum(-1) / denom + 1e-5).sqrt()
         x = ((x - μ[..., None]) / σ[..., None]).masked_fill(m[:, None], 0)
         self.cache_length = self.calc_length(l, repeat_num=self.config.num_feat_extract_layers).long()
-        return F.pad(x, (0, (-T) % self.pad_to)).transpose(1, 2)
+        return F.pad(x, (0, (-T) % self.config.pad_to)).transpose(1, 2)
 
     def forward(self, input_values):
-        return self._greedy_decode(super().forward(self.preprocessing(input_values)).last_hidden_state)
-
-    def _greedy_decode(self, enc_out: torch.Tensor):
-        B, T, _ = enc_out.size()
-        device = enc_out.device
-        enc_proj = self.enc(enc_out)
-        max_symbols = self.config.max_symbols_per_step
-        max_len = T * max_symbols
-        token_buffer = torch.full(
-            (B, max_len),
-            -1,
-            dtype=torch.long,
-            device=device
-        )
-        start_buffer = torch.zeros(
-            (B, max_len),
-            device=device
-        )
-        lengths = torch.zeros(B, dtype=torch.long, device=device)
-        last = torch.full(
-            (B, 1),
-            self.config.blank_id,
-            dtype=torch.long,
-            device=device
-        )
-
-        h = None
+        return self.postprocessing(super().forward(self.preprocessing(input_values)).last_hidden_state)
 
-        for t in range(T):
-            e = enc_proj[:, t:t+1]
-
-            for _ in range(max_symbols):
-                p, h2 = self.lstm(self.embed(last), h)
-                joint = self.joint(self.act(e + self.pred(p))).squeeze(1)
+    def postprocessing(self, enc_out):
+        B, T, _ = enc_out.shape
+        H = self.lstm.hidden_size
+        blank = self.config.blank_id
+        pad = self.config.pad_id
+        max_len = T * self.config.max_symbols_per_step
 
-                n = joint.argmax(-1)
-                blank = n.eq(self.config.blank_id)
-                emit_mask = ~blank
+        tokens  = torch.full((B, max_len), pad,  dtype=torch.long,  device=enc_out.device)
+        starts  = torch.full((B, max_len), -1.0, dtype=enc_out.dtype, device=enc_out.device)
+        lengths = torch.zeros(B, dtype=torch.long, device=enc_out.device)
+        hx      = torch.zeros(1, B, H, dtype=enc_out.dtype, device=enc_out.device)
+        cx      = torch.zeros_like(hx)
+        last    = torch.full((B, 1), blank, dtype=torch.long, device=enc_out.device)
 
-                if not emit_mask.any():
-                    break
+        enc_proj = self.enc(enc_out)  # (B, T, D)
 
-                pos = lengths[emit_mask]
-
-                token_buffer[emit_mask, pos] = n[emit_mask]
-                start_buffer[emit_mask, pos] = t * self.denorm
+        for t in range(T):
+            e      = enc_proj[:, t:t+1]
+            t_sec  = torch.full((B, 1), t * self.denorm, dtype=enc_out.dtype, device=enc_out.device)
 
-                lengths[emit_mask] += 1
+            for _ in range(self.config.max_symbols_per_step):
+                hx_prev, cx_prev = hx, cx
 
-                last = torch.where(emit_mask[:, None], n[:, None], last)
+                p, (hx, cx) = self.lstm(self.embed(last), (hx, cx))
+                n            = self.joint(self.act(e + self.pred(p))).squeeze(1).argmax(-1)  # (B,)
+                emitted      = n.ne(blank)
 
-                if h is None:
-                    h = h2
-                else:
-                    keep_mask = blank.view(1, -1, 1)
-                    h = (
-                        torch.where(keep_mask, h[0], h2[0]),
-                        torch.where(keep_mask, h[1], h2[1]),
-                    )
+                # revert hidden for blanks
+                mask = emitted.view(1, B, 1)
+                hx   = torch.where(mask, hx, hx_prev)
+                cx   = torch.where(mask, cx, cx_prev)
+                last = torch.where(emitted.unsqueeze(1), n.unsqueeze(1), last)
 
-        tokens = []
-        starts = []
+                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))
 
-        for b in range(B):
-            L = lengths[b]
-            tokens.append(token_buffer[b, :L])
-            starts.append(start_buffer[b, :L])
+                tokens  = tokens.scatter(1, pos, fill_t)
+                starts  = starts.scatter(1, pos, fill_s)
+                lengths = lengths + emitted.long()
 
         return tokens, starts