from huggingface_hub import hf_hub_download from torch import nn from transformers import Wav2Vec2ConformerModel from safetensors.torch import load_file from torch_state_bridge import state_bridge import torch import torch.nn.functional as F import torchaudio import librosa class Op(nn.Module): def __init__(self, func,allow_self=False): super().__init__() self.func = func self.allow_self = allow_self def forward(self, x): if self.allow_self: return self.func(self,x) return self.func(x) class Wav2Vec2ConformerRNNT(Wav2Vec2ConformerModel): 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 // 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() 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=self.spec.n_fft,n_mels=80))) for idx,l in enumerate(self.feature_extractor.conv_layers): 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] * self.calc_length(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) l.conv_module.pointwise_conv1.bias = nn.Parameter(torch.empty(l.conv_module.pointwise_conv1.out_channels)) 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 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.values()}) 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() def _mask_hidden_states(self, hidden_states, mask_time_indices = None, attention_mask = None): hidden_states = hidden_states * self.scaler 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, padding=1, kernel_size=3, stride=2, repeat_num=1): for _ in range(repeat_num): lengths = (lengths + 2 * padding - kernel_size) // stride + 1 return lengths def preprocessing(self, x): x, l = x l = (l // self.spec.hop_length + 1).long() 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] x = x.masked_fill(m[:, None], 0) μ = x.sum(-1) / l[:, None] 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 = self.calc_length(l, repeat_num=self.config.num_feat_extract_layers).long() return F.pad(x, (0, (-T) % self.config.pad_to)).transpose(1, 2) def forward(self, input_values): return self.postprocessing(super().forward(self.preprocessing(input_values)).last_hidden_state) 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 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(self.config.lstm_layer, 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) enc_proj = self.enc(enc_out) # (B, T, D) 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) for _ in range(self.config.max_symbols_per_step): hx_prev, cx_prev = hx, cx 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) # 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) 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 def change_language(self,language): self.joint.load_state_dict(self.lang_joint_net[language].state_dict()) @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) @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\n" changes += f"lang_joint_net.{model.config.language},joint\n" else: changes += "encoder.pre_encode.conv_module.{n},encoder.pre_encode.conv_module.{(n-2)}\n" changes += "encoder.pre_encode.conv_module.{n},feature_extractor.conv_layers.{(n//3+1)}.conv.{(n%3)}\n" 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)