Update modeling_conformer.py

modeling_conformer.py

@@ -1,5 +1,4 @@
 from datetime import timedelta
-import gc
 import json
 from huggingface_hub import hf_hub_download
 import torch
@@ -15,56 +14,83 @@ import webrtcvad
 from torch.utils.data import Dataset , DataLoader
 import srt
 
-def calc_length(lengths, all_paddings=2, kernel_size=3, stride=2, repeat_num=1):
-    add_pad = all_paddings - kernel_size
-    for _ in range(repeat_num):
-        lengths = torch.floor((lengths.float() + add_pad) / stride + 1)
-    return lengths
-
 class ChunkedData(Dataset):
     def __init__(self, wav, sr):
-        if sr != 16000: wav = torchaudio.functional.resample(wav, sr, 16000)
-        wav = wav.mean(0, keepdim=True)
-        self.data, self.ts = self.make_chunks(wav)
-
-    def __len__(self): return len(self.data)
-    def __getitem__(self, i): return self.data[i], self.ts[i]
-
-    def make_chunks(self, wav, sr=16000, ag=2, min_s=10, max_s=15, ms=30):
-        w = (wav * 32768).clamp(-32768, 32767).short().squeeze(0)
-        fl = int(sr * ms / 1000)
-        nf = len(w) // fl
-        w = w[: nf * fl]
-        fr = w.view(nf, fl)
+        if sr != 16000:
+            wav = torchaudio.functional.resample(wav, sr, 16000)
+
+        self.wav = wav.mean(0, keepdim=True)
+        self.sr = 16000
+        
+        # Sirf timestamps store karo, actual chunk nahi
+        self.ts = self.make_chunk_timestamps(self.wav)
+
+    def __len__(self):
+        return len(self.ts)
+
+    def __getitem__(self, i):
+        st, ed = self.ts[i]
+        st_i = int(st * self.sr)
+        ed_i = int(ed * self.sr)
+        chunk = self.wav[:, st_i:ed_i].squeeze()
+        return chunk, self.ts[i]
+
+    def make_chunk_timestamps(self, wav, sr=16000, ag=2, min_s=10, max_s=15, ms=30):
+
+        wav_int16 = (wav * 32768).clamp(-32768, 32767).short().squeeze(0)
+
+        frame_len = int(sr * ms / 1000)
+        num_frames = len(wav_int16) // frame_len
+        wav_int16 = wav_int16[: num_frames * frame_len]
+
+        frames = wav_int16.view(num_frames, frame_len)
+
         vad = webrtcvad.Vad(ag)
-        sp = torch.zeros(nf, dtype=torch.bool)
-        for i, f in enumerate(fr):
-            try: sp[i] = vad.is_speech(f.cpu().numpy().tobytes(), sr)
-            except: pass
-        seg, s = [], None
-        for i, v in enumerate(sp):
-            if v and s is None: s = i
-            elif not v and s is not None: seg.append((s, i)); s = None
-        if s is not None: seg.append((s, len(sp)))
-        cs, ts, st = [], [], 0
-        mn, mx, N = int(min_s * sr), int(max_s * sr), len(w)
-        while st < N:
-            ed = min(st + mx, N)
-            f = ed // fl
-            while f < len(sp) and sp[f]:
-                f += 1; ed = min(f * fl, N)
-                if ed - st > mx * 1.5: break
-            if ed - st < mn and ed < N: ed = min(st + mn, N)
-            cs.append(wav[:, st:ed].squeeze())
-            ts.append([round(st / sr, 2), round(ed / sr, 2)])
+        speech = torch.tensor(
+            [vad.is_speech(frame.numpy().tobytes(), sr) for frame in frames],
+            dtype=torch.bool
+        )
+
+        timestamps = []
+        total_samples = len(wav_int16)
+
+        min_len = int(min_s * sr)
+        max_len = int(max_s * sr)
+
+        st = 0
+
+        while st < total_samples:
+            ed = min(st + max_len, total_samples)
+
+            if ed - st < min_len and ed < total_samples:
+                ed = min(st + min_len, total_samples)
+
+            timestamps.append((
+                round(st / sr, 2),
+                round(ed / sr, 2)
+            ))
+
             st = ed
-        return cs, torch.tensor(ts)
+
+        return timestamps
 
 
 
 def padding_audio(batch):
     audios, times = zip(*batch)
-    return pad_sequence(audios, batch_first=True), torch.tensor([audio.numel() for audio in audios]), torch.stack(times)
+
+    lengths = torch.tensor([audio.numel() for audio in audios])
+    times = torch.tensor(times, dtype=torch.float32)
+
+    padded = pad_sequence(audios, batch_first=True)
+
+    return padded, lengths, times
+
+def calc_length(lengths, all_paddings=2, kernel_size=3, stride=2, repeat_num=1):
+    add_pad = all_paddings - kernel_size
+    for _ in range(repeat_num):
+        lengths = torch.floor((lengths.float() + add_pad) / stride + 1)
+    return lengths
 
 class Op(nn.Module):
     def __init__(self, func,allow_self=False):
@@ -110,7 +136,16 @@ class Wav2Vec2ConformerRNNT(Wav2Vec2ConformerModel):
         self.act = nn.ReLU(inplace=True)
         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.mel_fb = nn.Parameter(torch.tensor(librosa.filters.mel(sr=self.config.sampling_rate, n_fft=512, n_mels=80)),False)
+        self.register_buffer(
+            "mel_fb",
+            torch.tensor(
+                librosa.filters.mel(
+                    sr=self.config.sampling_rate,
+                    n_fft=512,
+                    n_mels=80
+                )
+            )
+        )
 
         for idx,l in enumerate(self.feature_extractor.conv_layers):
             if len(self.config.languages) == 1 or idx == 0:
@@ -146,13 +181,14 @@ class Wav2Vec2ConformerRNNT(Wav2Vec2ConformerModel):
         m = torch.arange(T, device=x.device)[None] >= l[:, None]
         x = x.masked_fill(m[:, None], 0)
         μ = x.sum(-1) / l[:, None]
-        σ = (((x - μ[..., None])**2).sum(-1) / (l[:, None] - 1) + 1e-5).sqrt()
+        denom = torch.clamp(l[:, None] - 1, min=1)
+        σ = (((x - μ[..., None])**2).sum(-1) / denom + 1e-5).sqrt()
         x = ((x - μ[..., None]) / σ[..., None]).masked_fill(m[:, None], 0)
         self.cache_length = calc_length(l, repeat_num=self.config.num_feat_extract_layers).long()
         return F.pad(x, (0, (-T) % self.pad_to)).transpose(1, 2)
 
     def forward(self, input_values):
-        return self.postprocessing(super().forward(self.preprocessing(input_values)).last_hidden_state)
+        return self._greedy_decode(super().forward(self.preprocessing(input_values)).last_hidden_state)
 
     @torch.inference_mode()
     def transcribe(self,wav,sr,batch_size):
@@ -164,12 +200,13 @@ class Wav2Vec2ConformerRNNT(Wav2Vec2ConformerModel):
             timestamp = timestamp.to(device)
             subtitles.extend(self.make_srt(self.forward((batch, lengths)),timestamp))
             yield srt.compose(subtitles)
-            torch.cuda.empty_cache()
-            gc.collect()
+            del batch
+            del lengths
 
     def load_state_dict(self, state_dict, strict=True, assign=False):
-        del state_dict['ctc_decoder.decoder_layers.0.bias']
-        del state_dict['ctc_decoder.decoder_layers.0.weight']
+        state_dict.pop('ctc_decoder.decoder_layers.0.bias', None)
+        state_dict.pop('ctc_decoder.decoder_layers.0.weight', None)
+
         state_dict['preprocessor.featurizer.fb'] = state_dict['preprocessor.featurizer.fb'].squeeze(0)
         changes = """
 preprocessor.featurizer.fb,mel_fb
@@ -207,61 +244,161 @@ encoder.pre_encode.conv_module.{n},feature_extractor.conv_layers.{(n//3+1)}.conv
             state_dict = {k: v for k, v in state_dict.items() if "lang_joint_net" not in k}
         return super().load_state_dict(state_dict, strict, assign)
 
-    def postprocessing(self, x):
-        if len(self.config.languages) > 1:
-            self.joint.load_state_dict(self.lang_joint_net[self.language].state_dict())
-        B = x.size(0)
-        last = x.new_full((B, 1), self.config.blank_id, dtype=torch.long)
-        h, tok, st = None, [[] for _ in range(B)], [[] for _ in range(B)]
-        for t, e in enumerate(x.unbind(1)):
-            v = t < self.cache_length
-            if not v.any(): break
-            e = e[:, None]
-            for _ in range(self.config.max_symbols_per_step):
+    @torch.jit.export
+    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
+
+        for t in range(T):
+            e = enc_proj[:, t:t+1]
+
+            for _ in range(max_symbols):
                 p, h2 = self.lstm(self.embed(last), h)
-                lg = self.joint(self.act(self.enc(e) + self.pred(p))).squeeze(1)
-                n = torch.where(v, lg.argmax(-1), self.config.blank_id)
-                b = n.eq(self.config.blank_id)
-                if b.all(): break
-                a = v & ~b
-                for i in a.nonzero().flatten().tolist():
-                    tok[i].append(n[i]); st[i].append(t * self.denorm)
-                last = torch.where(a[:, None], n[:, None], last)
-                if h is None: h = h2
+                joint = self.joint(self.act(e + self.pred(p))).squeeze(1)
+
+                n = joint.argmax(-1)
+                blank = n.eq(self.config.blank_id)
+                emit_mask = ~blank
+
+                if not emit_mask.any():
+                    break
+
+                pos = lengths[emit_mask]
+
+                token_buffer[emit_mask, pos] = n[emit_mask]
+                start_buffer[emit_mask, pos] = t * self.denorm
+
+                lengths[emit_mask] += 1
+
+                last = torch.where(emit_mask[:, None], n[:, None], last)
+
+                if h is None:
+                    h = h2
                 else:
-                    k = (b | ~v).view(1, -1, 1)
-                    h = (torch.where(k, h[0], h2[0]), torch.where(k, h[1], h2[1]))
-        self.cache_length = None
-        device = next(self.parameters()).device
-        return [torch.tensor(i,device=device) for i in tok], [torch.tensor(i,device=device) for i in st]
-
-    def make_srt(self, x, ts):
-        t , s = x
-        start_token_segment = self.config.languages.index(self.language) * self.joint.out_features
-        all_tokens, all_starts, all_ends = [], [], []
-        device = t[0].device
-        for tokens, starts, (s, e) in zip(t,s, ts):
-            tokens += start_token_segment
-            starts += s
+                    keep_mask = blank.view(1, -1, 1)
+                    h = (
+                        torch.where(keep_mask, h[0], h2[0]),
+                        torch.where(keep_mask, h[1], h2[1]),
+                    )
+
+        tokens = []
+        starts = []
+
+        for b in range(B):
+            L = lengths[b]
+            tokens.append(token_buffer[b, :L])
+            starts.append(start_buffer[b, :L])
+
+        return tokens, starts
+
+    def make_srt(self, decoded, ts):
+
+        tokens_list, starts_list = decoded
+
+        start_token_segment = (
+            self.config.languages.index(self.language)
+            * self.joint.out_features
+        )
+
+        all_tokens = []
+        all_starts = []
+        all_ends = []
+
+        device = tokens_list[0].device
+
+        for tokens, starts, (seg_start, seg_end) in zip(
+                tokens_list, starts_list, ts):
+
+            tokens = tokens + start_token_segment
+            starts = starts + seg_start
+
             all_tokens.append(tokens)
             all_starts.append(starts)
-            all_ends.append(torch.cat([starts[1:], e[None]]))
-            all_tokens.append(torch.tensor([-1],device=device))
-            all_starts.append(torch.tensor([e],device=device))
-            all_ends.append(torch.tensor([e + 0.005],device=device))
-        return [srt.Subtitle(i,timedelta(seconds=float(st)),timedelta(seconds=float(en)),"<line>" if tok == -1 else self.config.vocab[int(tok)]) for i, (tok, st, en) in enumerate(zip(torch.cat(all_tokens), torch.cat(all_starts), torch.cat(all_ends)), 1)]
+            all_ends.append(torch.cat([starts[1:], seg_end[None]]))
+
+            # newline marker
+            all_tokens.append(torch.tensor([-1], device=device))
+            all_starts.append(torch.tensor([seg_end], device=device))
+            all_ends.append(torch.tensor([seg_end + 0.005], device=device))
+
+        return [
+            srt.Subtitle(
+                i,
+                timedelta(seconds=float(st)),
+                timedelta(seconds=float(en)),
+                "<line>" if tok == -1 else self.config.vocab[int(tok)]
+            )
+            for i, (tok, st, en) in enumerate(
+                zip(
+                    torch.cat(all_tokens),
+                    torch.cat(all_starts),
+                    torch.cat(all_ends)
+                ), 1
+            )
+        ]
 
 
     @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path, config = None, language=None,**kwargs):
+    def from_pretrained(
+            cls,
+            pretrained_model_name_or_path,
+            config=None,
+            language=None,
+            use_jit=False,
+            use_quantization=False):
+
+        if config is None:
+            raise ValueError("config must be provided")
+
         if language:
             config.languages = [language]
-            config.vocab = ['<unk>'] + json.load(open(hf_hub_download(pretrained_model_name_or_path, "vocab.json")))['small'][language]
-        else:
-            temp_vocab = json.load(open(hf_hub_download(pretrained_model_name_or_path, "vocab.json")))['large']
-            config.vocab = []
-            for i in sorted(config.languages):
-                config.vocab.extend(['<unk>'] + temp_vocab[i])
+
+            vocab_file = hf_hub_download(
+                pretrained_model_name_or_path,
+                "vocab.json"
+            )
+
+            vocab_json = json.load(open(vocab_file))
+            config.vocab = ['<unk>'] + vocab_json['small'][language]
+
         model = cls(config)
-        model.load_state_dict(load_file(hf_hub_download(pretrained_model_name_or_path, f"{language or 'all'}.safetensors")))
+
+        weight_file = hf_hub_download(
+            pretrained_model_name_or_path,
+            f"{language or 'all'}.safetensors"
+        )
+
+        model.load_state_dict(load_file(weight_file))
+        if use_quantization:
+            model = torch.quantization.quantize_dynamic(model)
+
         return model