Update modeling_conformer.py

modeling_conformer.py

@@ -1,9 +1,8 @@
 from huggingface_hub import hf_hub_download
 from torch import nn
-from transformers import Wav2Vec2ConformerModel
+from transformers import Wav2Vec2ConformerModel , Wav2Vec2CTCTokenizer
 from safetensors.torch import load_file
 from torch_state_bridge import state_bridge
-import json
 import torch
 import torch.nn.functional as F
 import torchaudio
@@ -21,58 +20,28 @@ class Op(nn.Module):
         return self.func(x)
 
 class Wav2Vec2ConformerRNNT(Wav2Vec2ConformerModel):
-    def __init__(self, config):
-        self.language = config.languages[0]
-        if len(config.languages) > 1:
-            config.hidden_size = 1024
-            config.num_hidden_layers = 24
-            config.conv_depthwise_kernel_size = 9
-            config.conv_stride = [2,2,2]
-            config.conv_kernel = [3,3,3]
-            config.conv_dim = [256,256,256]
-            config.feat_extract_norm = "group"
-            config.intermediate_size = 4096
-            config.num_feat_extract_layers = len(config.conv_dim)
-            config.lstm_layer = 2
-
-        self.cache_length = None
-        self.hop, self.preemph, self.eps, self.pad_to = 160, 0.97, 2**-24, 16
-        self.denorm = (2 ** config.num_feat_extract_layers) * self.hop / config.sampling_rate
-        self.scaler = config.hidden_size ** (1/2)
-        super().__init__(config)
-        self.eval()
 
     def init_weights(self):
         del self.encoder.pos_conv_embed
         config = self.config
+        self.cache_length = None
         self.enc = nn.Linear(config.hidden_size, config.joint_hidden)
         self.pred = nn.Linear(config.pred_hidden, config.joint_hidden)
-        self.joint = nn.Linear(config.joint_hidden, config.vocab_size // 22 + 1)
+        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.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=self.config.sampling_rate,
-                    n_fft=512,
-                    n_mels=80
-                )
-            )
-        )
-
+        self.register_buffer("mel_fb",torch.tensor(librosa.filters.mel(sr=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:
+            if not(config.multilingual) or idx == 0:
                 l.conv = nn.Conv2d(l.conv.in_channels,l.conv.out_channels,l.conv.kernel_size[0],l.conv.stride,1)
                 l.layer_norm = nn.Identity()
             else:
                 l.conv = nn.Sequential(nn.Conv2d(l.conv.in_channels,l.conv.out_channels,l.conv.kernel_size[0],l.conv.stride,1,groups=l.conv.out_channels),nn.Conv2d(l.conv.in_channels,l.conv.out_channels, 1))
-
         self.feature_extractor.conv_layers.append(Op(lambda x : x.transpose(1, 2)))
-        self.feature_projection.projection = nn.Linear(config.conv_dim[-1] * int(self.calc_length(torch.tensor(80.),repeat_num=self.config.num_feat_extract_layers)),config.hidden_size)
+        self.feature_projection.projection = nn.Linear(config.conv_dim[-1] * int(self.calc_length(torch.tensor(80.),repeat_num=config.num_feat_extract_layers)),config.hidden_size)
         self.feature_projection.layer_norm = Op(lambda x:x.permute(0, 2, 1, 3).flatten(2))
         for l in self.encoder.layers:
             l.conv_module.glu = nn.Sequential(l.conv_module.glu,self.mask_layer)
@@ -80,8 +49,11 @@ class Wav2Vec2ConformerRNNT(Wav2Vec2ConformerModel):
             l.conv_module.pointwise_conv2.bias = nn.Parameter(torch.empty(l.conv_module.pointwise_conv2.out_channels))
             l.conv_module.depthwise_conv.bias = nn.Parameter(torch.empty(l.conv_module.depthwise_conv.out_channels))
         self.encoder.layer_norm = nn.Identity()
-        if len(self.config.languages) > 1:
-            self.lang_joint_net = nn.ModuleDict({l: nn.Linear(config.joint_hidden, config.vocab_size // 22 + 1) for l in config.languages})
+        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.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()
 
     def _mask_hidden_states(self, hidden_states, mask_time_indices = None, attention_mask = None):
@@ -97,7 +69,7 @@ class Wav2Vec2ConformerRNNT(Wav2Vec2ConformerModel):
 
     def preprocessing(self, x):
         x, l = x
-        l = (l // self.hop + 1).long()
+        l = (l // self.spec.hop_length + 1).long()
         x = torch.cat((x[:, :1], x[:, 1:] - self.preemph * x[:, :-1]), 1)
         x = (self.mel_fb @ self.spec(x) + self.eps).log()
         T = x.size(-1)
@@ -113,72 +85,23 @@ class Wav2Vec2ConformerRNNT(Wav2Vec2ConformerModel):
     def forward(self, input_values):
         return self._greedy_decode(super().forward(self.preprocessing(input_values)).last_hidden_state)
 
-    def load_state_dict(self, state_dict, strict=True, assign=False):
-        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
-preprocessor.featurizer.window,spec.window
-norm_feed_forward1,ffn1_layer_norm
-norm_feed_forward2,ffn2_layer_norm
-feed_forward1.linear1,ffn1.intermediate_dense
-feed_forward1.linear2,ffn1.output_dense
-feed_forward2.linear1,ffn2.intermediate_dense
-feed_forward2.linear2,ffn2.output_dense
-norm_self_att,self_attn_layer_norm
-norm_out,final_layer_norm
-norm_conv,conv_module.layer_norm
-.conv.,.conv_module.
-decoder.prediction.dec_rnn.lstm,lstm
-decoder.prediction.embed,embed
-joint.enc,enc
-joint.pred,pred
-joint.joint_net.2,lang_joint_net
-encoder.pre_encode.conv_module.0,feature_extractor.conv_layers.0.conv
-encoder.pre_encode.out,feature_projection.projection
-"""
-        if len(self.config.languages) == 1:
-            changes += f"""lang_joint_net.{self.language},joint
-encoder.pre_encode.conv_module.{{n}},feature_extractor.conv_layers.{{(n/2)}}.conv"""
-        else:
-            state_dict["joint.weight"] = self.joint.weight.clone()
-            state_dict["joint.bias"] = self.joint.bias.clone()
-            changes += """encoder.pre_encode.conv_module.{n},encoder.pre_encode.conv_module.{(n-2)}
-encoder.pre_encode.conv_module.{n},feature_extractor.conv_layers.{(n//3+1)}.conv.{(n%3)}
-"""
-            # replicate many changes for complex maths
-        state_dict = state_bridge(state_dict, changes)
-        if len(self.config.languages) == 1:
-            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)
-
-    @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,
@@ -230,37 +153,60 @@ encoder.pre_encode.conv_module.{n},feature_extractor.conv_layers.{(n//3+1)}.conv
 
         return tokens, starts
 
-    @classmethod
-    def from_pretrained(
-            cls,
-            pretrained_model_name_or_path,
-            config=None,
-            language=None,
-            use_quantization=False):
-
-        if config is None:
-            raise ValueError("config must be provided")
-
-        if language:
-            config.languages = [language]
-
-            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)
-
-        weight_file = hf_hub_download(
-            pretrained_model_name_or_path,
-            f"{language or 'all'}.safetensors"
-        )
+    def change_language(self,language):
+        self.joint.load_state_dict(self.lang_joint_net[language].state_dict())
 
-        model.load_state_dict(load_file(weight_file))
-        if use_quantization:
-            model = torch.quantization.quantize_dynamic(model)
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, config = None, cache_dir = None, ignore_mismatched_sizes = False, force_download = False, local_files_only = False, token = None, revision = "main", use_safetensors = None, weights_only = True, **kwargs):
+        config.language = kwargs.pop("language",None)
+        config.multilingual = not(config.language)
+        if config.multilingual:
+            config.hidden_size = 1024
+            config.num_hidden_layers = 24
+            config.conv_depthwise_kernel_size = 9
+            config.conv_stride = [2,2,2]
+            config.conv_kernel = [3,3,3]
+            config.conv_dim = [256,256,256]
+            config.feat_extract_norm = "group"
+            config.intermediate_size = config.hidden_size * 4
+            config.num_feat_extract_layers = len(config.conv_dim)
+            config.lstm_layer = 2
+        kwargs['state_dict'] = load_file(hf_hub_download(pretrained_model_name_or_path,f"{config.language or 'all'}.safetensors"))
+        return super().from_pretrained(None, *model_args, config=config, cache_dir=cache_dir, ignore_mismatched_sizes=ignore_mismatched_sizes, force_download=force_download, local_files_only=local_files_only, token=token, revision=revision, use_safetensors=use_safetensors, weights_only=weights_only, **kwargs)
 
-        return model
+    @staticmethod
+    def _load_pretrained_model(model, state_dict, checkpoint_files, load_config):
+        changes = """
+preprocessor.featurizer.fb,mel_fb
+preprocessor.featurizer.window,spec.window
+norm_feed_forward1,ffn1_layer_norm
+norm_feed_forward2,ffn2_layer_norm
+feed_forward1.linear1,ffn1.intermediate_dense
+feed_forward1.linear2,ffn1.output_dense
+feed_forward2.linear1,ffn2.intermediate_dense
+feed_forward2.linear2,ffn2.output_dense
+norm_self_att,self_attn_layer_norm
+norm_out,final_layer_norm
+norm_conv,conv_module.layer_norm
+.conv.,.conv_module.
+decoder.prediction.dec_rnn.lstm,lstm
+decoder.prediction.embed,embed
+joint.enc,enc
+joint.pred,pred
+joint.joint_net.2,lang_joint_net
+encoder.pre_encode.conv_module.0,feature_extractor.conv_layers.0.conv
+encoder.pre_encode.out,feature_projection.projection
+"""
+        if not model.config.multilingual:
+            changes += "encoder.pre_encode.conv_module.{n},feature_extractor.conv_layers.{(n/2)}.conv"
+            changes += f"lang_joint_net.{model.config.language},joint"
+        else:
+            changes += "encoder.pre_encode.conv_module.{n},encoder.pre_encode.conv_module.{(n-2)}"
+            changes += "encoder.pre_encode.conv_module.{n},feature_extractor.conv_layers.{(n//3+1)}.conv.{(n%3)}"
+        state_dict = state_bridge(state_dict, changes)
+        if not model.config.multilingual:
+            state_dict = {k: v for k, v in state_dict.items() if "lang_joint_net" not in k}
+        state_dict['mel_fb'] = state_dict['mel_fb'].squeeze(0)
+        state_dict.pop('ctc_decoder.decoder_layers.0.bias', None)
+        state_dict.pop('ctc_decoder.decoder_layers.0.weight', None)
+        return super()._load_pretrained_model(model, state_dict, checkpoint_files, load_config)