diff --git a/almeval/models/stepaudio/cosyvoice/__init__.py b/almeval/models/stepaudio/cosyvoice/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/almeval/models/stepaudio/cosyvoice/cli/__init__.py b/almeval/models/stepaudio/cosyvoice/cli/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/almeval/models/stepaudio/cosyvoice/cli/cosyvoice.py b/almeval/models/stepaudio/cosyvoice/cli/cosyvoice.py new file mode 100644 index 0000000000000000000000000000000000000000..b8c0b4ba2a45c8d5bdeca5279fdaaac53d89b544 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/cli/cosyvoice.py @@ -0,0 +1,68 @@ +# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os +import uuid +import time +from tqdm import tqdm +import torch +import torchaudio +from hyperpyyaml import load_hyperpyyaml +from cosyvoice.cli.frontend import CosyVoiceFrontEnd +from cosyvoice.cli.model import CosyVoiceModel + + +class CosyVoice: + + def __init__( + self, + model_dir, + ): + self.model_dir = model_dir + with open("{}/cosyvoice.yaml".format(model_dir), "r") as f: + configs = load_hyperpyyaml(f) + self.frontend = CosyVoiceFrontEnd( + configs["feat_extractor"], + "{}/campplus.onnx".format(model_dir), + "{}/speech_tokenizer_v1.onnx".format(model_dir), + ) + self.model = CosyVoiceModel(configs["flow"], configs["hift"]) + self.model.load( + "{}/flow.pt".format(model_dir), + "{}/hift.pt".format(model_dir), + ) + self.model.flow = self.model.flow.to(torch.bfloat16) + del configs + + def token_to_wav_offline( + self, + speech_token, + speech_feat, + speech_feat_len, + prompt_token, + prompt_token_len, + embedding, + ): + tts_mel = self.model.flow.inference( + token=speech_token.to(self.model.device), + token_len=torch.tensor([speech_token.size(1)], dtype=torch.int32).to( + self.model.device + ), + prompt_token=prompt_token.to(self.model.device), + prompt_token_len=prompt_token_len.to(self.model.device), + prompt_feat=speech_feat.to(self.model.device), + prompt_feat_len=speech_feat_len.to(self.model.device), + embedding=embedding.to(self.model.device), + ) + tts_speech = self.model.hift.inference(mel=tts_mel.float())[0].cpu() + return tts_speech diff --git a/almeval/models/stepaudio/cosyvoice/cli/frontend.py b/almeval/models/stepaudio/cosyvoice/cli/frontend.py new file mode 100644 index 0000000000000000000000000000000000000000..d968551816eaa87f32547222470afdf2c64f4d68 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/cli/frontend.py @@ -0,0 +1,106 @@ +# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import onnxruntime +import torch +import numpy as np +import whisper +from typing import Callable +import torchaudio.compliance.kaldi as kaldi + + +class CosyVoiceFrontEnd: + + def __init__( + self, + feat_extractor: Callable, + campplus_model: str, + speech_tokenizer_model: str, + ): + self.feat_extractor = feat_extractor + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + option = onnxruntime.SessionOptions() + option.graph_optimization_level = ( + onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL + ) + option.intra_op_num_threads = 1 + self.campplus_session = onnxruntime.InferenceSession( + campplus_model, sess_options=option, providers=["CPUExecutionProvider"] + ) + self.speech_tokenizer_session = onnxruntime.InferenceSession( + speech_tokenizer_model, + sess_options=option, + providers=[ + ( + "CUDAExecutionProvider" + if torch.cuda.is_available() + else "CPUExecutionProvider" + ) + ], + ) + + def _extract_speech_token(self, speech): + assert ( + speech.shape[1] / 16000 <= 30 + ), "do not support extract speech token for audio longer than 30s" + feat = whisper.log_mel_spectrogram(speech, n_mels=128) + speech_token = ( + self.speech_tokenizer_session.run( + None, + { + self.speech_tokenizer_session.get_inputs()[0] + .name: feat.detach() + .cpu() + .numpy(), + self.speech_tokenizer_session.get_inputs()[1].name: np.array( + [feat.shape[2]], dtype=np.int32 + ), + }, + )[0] + .flatten() + .tolist() + ) + speech_token = torch.tensor([speech_token], dtype=torch.int32).to(self.device) + speech_token_len = torch.tensor([speech_token.shape[1]], dtype=torch.int32).to( + self.device + ) + return speech_token, speech_token_len + + def _extract_spk_embedding(self, speech): + feat = kaldi.fbank(speech, num_mel_bins=80, dither=0, sample_frequency=16000) + feat = feat - feat.mean(dim=0, keepdim=True) + embedding = ( + self.campplus_session.run( + None, + { + self.campplus_session.get_inputs()[0] + .name: feat.unsqueeze(dim=0) + .cpu() + .numpy() + }, + )[0] + .flatten() + .tolist() + ) + embedding = torch.tensor([embedding]).to(self.device) + return embedding + + def _extract_speech_feat(self, speech): + speech_feat = ( + self.feat_extractor(speech).squeeze(dim=0).transpose(0, 1).to(self.device) + ) + speech_feat = speech_feat.unsqueeze(dim=0) + speech_feat_len = torch.tensor([speech_feat.shape[1]], dtype=torch.int32).to( + self.device + ) + return speech_feat, speech_feat_len diff --git a/almeval/models/stepaudio/cosyvoice/cli/model.py b/almeval/models/stepaudio/cosyvoice/cli/model.py new file mode 100644 index 0000000000000000000000000000000000000000..b284d9e869b06d9d083431446295b56d2023076d --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/cli/model.py @@ -0,0 +1,32 @@ +# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import torch + + +class CosyVoiceModel: + + def __init__( + self, + flow: torch.nn.Module, + hift: torch.nn.Module, + ): + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.flow = flow + self.hift = hift + + def load(self, flow_model, hift_model): + self.flow.load_state_dict(torch.load(flow_model, map_location=self.device)) + self.flow.to(self.device).eval() + self.hift.load_state_dict(torch.load(hift_model, map_location=self.device)) + self.hift.to(self.device).eval() diff --git a/almeval/models/stepaudio/cosyvoice/flow/decoder.py b/almeval/models/stepaudio/cosyvoice/flow/decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..9fa6ba6dab254c8d8cff5f32003fbb202491a422 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/flow/decoder.py @@ -0,0 +1,238 @@ +# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Zhihao Du) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import torch +import torch.nn as nn +from einops import pack, rearrange, repeat +from cosyvoice.matcha.decoder import ( + SinusoidalPosEmb, + Block1D, + ResnetBlock1D, + Downsample1D, + TimestepEmbedding, + Upsample1D, +) +from cosyvoice.matcha.transformer import BasicTransformerBlock + + +class ConditionalDecoder(nn.Module): + def __init__( + self, + in_channels, + out_channels, + channels=(256, 256), + dropout=0.05, + attention_head_dim=64, + n_blocks=1, + num_mid_blocks=2, + num_heads=4, + act_fn="snake", + ): + """ + This decoder requires an input with the same shape of the target. So, if your text content + is shorter or longer than the outputs, please re-sampling it before feeding to the decoder. + """ + super().__init__() + channels = tuple(channels) + self.in_channels = in_channels + self.out_channels = out_channels + + self.time_embeddings = SinusoidalPosEmb(in_channels) + time_embed_dim = channels[0] * 4 + self.time_mlp = TimestepEmbedding( + in_channels=in_channels, + time_embed_dim=time_embed_dim, + act_fn="silu", + ) + self.down_blocks = nn.ModuleList([]) + self.mid_blocks = nn.ModuleList([]) + self.up_blocks = nn.ModuleList([]) + + output_channel = in_channels + for i in range(len(channels)): # pylint: disable=consider-using-enumerate + input_channel = output_channel + output_channel = channels[i] + is_last = i == len(channels) - 1 + resnet = ResnetBlock1D( + dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim + ) + transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + dim=output_channel, + num_attention_heads=num_heads, + attention_head_dim=attention_head_dim, + dropout=dropout, + activation_fn=act_fn, + ) + for _ in range(n_blocks) + ] + ) + downsample = ( + Downsample1D(output_channel) + if not is_last + else nn.Conv1d(output_channel, output_channel, 3, padding=1) + ) + self.down_blocks.append( + nn.ModuleList([resnet, transformer_blocks, downsample]) + ) + + for _ in range(num_mid_blocks): + input_channel = channels[-1] + out_channels = channels[-1] + resnet = ResnetBlock1D( + dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim + ) + + transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + dim=output_channel, + num_attention_heads=num_heads, + attention_head_dim=attention_head_dim, + dropout=dropout, + activation_fn=act_fn, + ) + for _ in range(n_blocks) + ] + ) + + self.mid_blocks.append(nn.ModuleList([resnet, transformer_blocks])) + + channels = channels[::-1] + (channels[0],) + for i in range(len(channels) - 1): + input_channel = channels[i] * 2 + output_channel = channels[i + 1] + is_last = i == len(channels) - 2 + resnet = ResnetBlock1D( + dim=input_channel, + dim_out=output_channel, + time_emb_dim=time_embed_dim, + ) + transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + dim=output_channel, + num_attention_heads=num_heads, + attention_head_dim=attention_head_dim, + dropout=dropout, + activation_fn=act_fn, + ) + for _ in range(n_blocks) + ] + ) + upsample = ( + Upsample1D(output_channel, use_conv_transpose=True) + if not is_last + else nn.Conv1d(output_channel, output_channel, 3, padding=1) + ) + self.up_blocks.append(nn.ModuleList([resnet, transformer_blocks, upsample])) + self.final_block = Block1D(channels[-1], channels[-1]) + self.final_proj = nn.Conv1d(channels[-1], self.out_channels, 1) + self.initialize_weights() + + def initialize_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv1d): + nn.init.kaiming_normal_(m.weight, nonlinearity="relu") + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.GroupNorm): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Linear): + nn.init.kaiming_normal_(m.weight, nonlinearity="relu") + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def forward(self, x, mask, mu, t, spks=None, cond=None): + """Forward pass of the UNet1DConditional model. + + Args: + x (torch.Tensor): shape (batch_size, in_channels, time) + mask (_type_): shape (batch_size, 1, time) + t (_type_): shape (batch_size) + spks (_type_, optional): shape: (batch_size, condition_channels). Defaults to None. + cond (_type_, optional): placeholder for future use. Defaults to None. + + Raises: + ValueError: _description_ + ValueError: _description_ + + Returns: + _type_: _description_ + """ + + t = self.time_embeddings(t).to(t.dtype) + t = self.time_mlp(t) + + x = pack([x, mu], "b * t")[0] + + if spks is not None: + spks = repeat(spks, "b c -> b c t", t=x.shape[-1]) + x = pack([x, spks], "b * t")[0] + if cond is not None: + x = pack([x, cond], "b * t")[0] + + hiddens = [] + masks = [mask] + for resnet, transformer_blocks, downsample in self.down_blocks: + mask_down = masks[-1] + x = resnet( + x.to(torch.bfloat16), mask_down.to(torch.bfloat16), t.to(torch.bfloat16) + ) + x = rearrange(x, "b c t -> b t c").contiguous() + # attn_mask = torch.matmul(mask_down.transpose(1, 2).contiguous(), mask_down) + for transformer_block in transformer_blocks: + x = transformer_block( + hidden_states=x, + # attention_mask=attn_mask, + timestep=t, + ) + x = rearrange(x, "b t c -> b c t").contiguous() + hiddens.append(x) # Save hidden states for skip connections + x = downsample(x * mask_down) + masks.append(mask_down[:, :, ::2]) + masks = masks[:-1] + mask_mid = masks[-1] + + for resnet, transformer_blocks in self.mid_blocks: + x = resnet(x, mask_mid, t) + x = rearrange(x, "b c t -> b t c").contiguous() + # attn_mask = torch.matmul(mask_mid.transpose(1, 2).contiguous(), mask_mid) + for transformer_block in transformer_blocks: + x = transformer_block( + hidden_states=x, + # attention_mask=attn_mask, + timestep=t, + ) + x = rearrange(x, "b t c -> b c t").contiguous() + + for resnet, transformer_blocks, upsample in self.up_blocks: + mask_up = masks.pop() + skip = hiddens.pop() + x = pack([x[:, :, : skip.shape[-1]], skip], "b * t")[0] + x = resnet(x, mask_up, t) + x = rearrange(x, "b c t -> b t c").contiguous() + # attn_mask = torch.matmul(mask_up.transpose(1, 2).contiguous(), mask_up) + for transformer_block in transformer_blocks: + x = transformer_block( + hidden_states=x, + # attention_mask=attn_mask, + timestep=t, + ) + x = rearrange(x, "b t c -> b c t").contiguous() + x = upsample(x * mask_up) + x = self.final_block(x, mask_up) + output = self.final_proj(x * mask_up) + return output * mask diff --git a/almeval/models/stepaudio/cosyvoice/flow/flow.py b/almeval/models/stepaudio/cosyvoice/flow/flow.py new file mode 100644 index 0000000000000000000000000000000000000000..ab9a8124ef86d542965aa4d2a1b8f97e6bd27e43 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/flow/flow.py @@ -0,0 +1,196 @@ +# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Zhihao Du) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import logging +import random +from typing import Dict, Optional +import torch +import torch.nn as nn +from torch.nn import functional as F +from omegaconf import DictConfig +from cosyvoice.utils.mask import make_pad_mask +import time + + +class MaskedDiffWithXvec(torch.nn.Module): + def __init__( + self, + input_size: int = 512, + output_size: int = 80, + spk_embed_dim: int = 192, + output_type: str = "mel", + vocab_size: int = 4096, + input_frame_rate: int = 50, + only_mask_loss: bool = True, + encoder: torch.nn.Module = None, + length_regulator: torch.nn.Module = None, + decoder: torch.nn.Module = None, + decoder_conf: Dict = { + "in_channels": 240, + "out_channel": 80, + "spk_emb_dim": 80, + "n_spks": 1, + "cfm_params": DictConfig( + { + "sigma_min": 1e-06, + "solver": "euler", + "t_scheduler": "cosine", + "training_cfg_rate": 0.2, + "inference_cfg_rate": 0.7, + "reg_loss_type": "l1", + } + ), + "decoder_params": { + "channels": [256, 256], + "dropout": 0.0, + "attention_head_dim": 64, + "n_blocks": 4, + "num_mid_blocks": 12, + "num_heads": 8, + "act_fn": "gelu", + }, + }, + mel_feat_conf: Dict = { + "n_fft": 1024, + "num_mels": 80, + "sampling_rate": 22050, + "hop_size": 256, + "win_size": 1024, + "fmin": 0, + "fmax": 8000, + }, + ): + super().__init__() + self.input_size = input_size + self.output_size = output_size + self.decoder_conf = decoder_conf + self.mel_feat_conf = mel_feat_conf + self.vocab_size = vocab_size + self.output_type = output_type + self.input_frame_rate = input_frame_rate + logging.info(f"input frame rate={self.input_frame_rate}") + self.input_embedding = nn.Embedding(vocab_size, input_size) + self.spk_embed_affine_layer = torch.nn.Linear(spk_embed_dim, output_size) + self.encoder = encoder + self.encoder_proj = torch.nn.Linear(self.encoder.output_size(), output_size) + self.decoder = decoder + self.length_regulator = length_regulator + self.only_mask_loss = only_mask_loss + + def forward( + self, + batch: dict, + device: torch.device, + ) -> Dict[str, Optional[torch.Tensor]]: + token = batch["speech_token"].to(device) + token_len = batch["speech_token_len"].to(device) + feat = batch["speech_feat"].to(device) + feat_len = batch["speech_feat_len"].to(device) + embedding = batch["embedding"].to(device) + + # xvec projection + embedding = F.normalize(embedding, dim=1) + embedding = self.spk_embed_affine_layer(embedding) + + # concat text and prompt_text + mask = (~make_pad_mask(token_len)).float().unsqueeze(-1).to(device) + token = self.input_embedding(torch.clamp(token, min=0)) * mask + + # text encode + h, h_lengths = self.encoder(token, token_len) + h = self.encoder_proj(h) + h, h_lengths = self.length_regulator(h, feat_len) + + # get conditions + conds = torch.zeros(feat.shape, device=token.device) + for i, j in enumerate(feat_len): + if random.random() < 0.5: + continue + index = random.randint(0, int(0.3 * j)) + conds[i, :index] = feat[i, :index] + conds = conds.transpose(1, 2) + + mask = (~make_pad_mask(feat_len)).to(h) + feat = F.interpolate( + feat.unsqueeze(dim=1), size=h.shape[1:], mode="nearest" + ).squeeze(dim=1) + loss, _ = self.decoder.compute_loss( + feat.transpose(1, 2).contiguous(), + mask.unsqueeze(1), + h.transpose(1, 2).contiguous(), + embedding, + cond=conds, + ) + return {"loss": loss} + + @torch.inference_mode() + def inference( + self, + token, + token_len, + prompt_token, + prompt_token_len, + prompt_feat, + prompt_feat_len, + embedding, + ): + assert token.shape[0] == 1 + # xvec projection + embedding = F.normalize(embedding, dim=1) + embedding = self.spk_embed_affine_layer(embedding) + + # concat text and prompt_text + token_len1, token_len2 = prompt_token.shape[1], token.shape[1] + # text encode + token, token_len = ( + torch.concat([prompt_token, token], dim=1), + prompt_token_len + token_len, + ) + token = self.input_embedding(torch.clamp(token, min=0)) + h, _ = self.encoder.inference(token, token_len) + h = self.encoder_proj(h) + mel_len1, mel_len2 = prompt_feat.shape[1], int( + token_len2 + / self.input_frame_rate + * self.mel_feat_conf["sampling_rate"] + / self.mel_feat_conf["hop_size"] + ) + + h, _ = self.length_regulator.inference( + h[:, :token_len1], + h[:, token_len1:], + mel_len1, + mel_len2, + ) + + # get conditions + conds = torch.zeros( + [1, mel_len1 + mel_len2, self.output_size], device=token.device + ) + conds[:, :mel_len1] = prompt_feat + conds = conds.transpose(1, 2) + + # mask = (~make_pad_mask(torch.tensor([mel_len1 + mel_len2]))).to(h) + mask = torch.ones( + [1, mel_len1 + mel_len2], device=h.device, dtype=torch.bfloat16 + ) + feat = self.decoder( + mu=h.transpose(1, 2).contiguous(), + mask=mask.unsqueeze(1), + spks=embedding, + cond=conds, + n_timesteps=10, + ) + feat = feat[:, :, mel_len1:] + assert feat.shape[2] == mel_len2 + return feat diff --git a/almeval/models/stepaudio/cosyvoice/flow/flow_matching.py b/almeval/models/stepaudio/cosyvoice/flow/flow_matching.py new file mode 100644 index 0000000000000000000000000000000000000000..d29673fcb4f21e43e65359df93ca119f707ca27b --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/flow/flow_matching.py @@ -0,0 +1,315 @@ +# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Zhihao Du) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import time +import torch +import torch.nn.functional as F +from cosyvoice.matcha.flow_matching import BASECFM + + +class ConditionalCFM(BASECFM): + def __init__( + self, + in_channels, + cfm_params, + n_spks=1, + spk_emb_dim=64, + estimator: torch.nn.Module = None, + ): + super().__init__( + n_feats=in_channels, + cfm_params=cfm_params, + n_spks=n_spks, + spk_emb_dim=spk_emb_dim, + ) + self.t_scheduler = cfm_params.t_scheduler + self.training_cfg_rate = cfm_params.training_cfg_rate + self.inference_cfg_rate = cfm_params.inference_cfg_rate + in_channels = in_channels + (spk_emb_dim if n_spks > 0 else 0) + # Just change the architecture of the estimator here + self.estimator = estimator + self.inference_graphs = {} + self.inference_buffers = {} + # self.capture_inference() + + @torch.inference_mode() + def forward( + self, + mu, + mask, + n_timesteps, + temperature=1.0, + spks=None, + cond=None, + ): + """Forward diffusion + + Args: + mu (torch.Tensor): output of encoder + shape: (batch_size, n_feats, mel_timesteps) + mask (torch.Tensor): output_mask + shape: (batch_size, 1, mel_timesteps) + n_timesteps (int): number of diffusion steps + temperature (float, optional): temperature for scaling noise. Defaults to 1.0. + spks (torch.Tensor, optional): speaker ids. Defaults to None. + shape: (batch_size, spk_emb_dim) + cond: Not used but kept for future purposes + + Returns: + sample: generated mel-spectrogram + shape: (batch_size, n_feats, mel_timesteps) + """ + z = torch.randn_like(mu) * temperature + t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) + if self.t_scheduler == "cosine": + t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) + return self.solve_euler( + z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond + ) + + @torch.inference_mode() + def capture_inference(self, seq_len_to_capture=list(range(128, 512, 8))): + start_time = time.time() + print( + f"capture_inference for ConditionalCFM solve euler, seq_len_to_capture: {seq_len_to_capture}" + ) + for seq_len in seq_len_to_capture: + static_z = torch.randn( + 1, 80, seq_len, device=torch.device("cuda"), dtype=torch.bfloat16 + ) + static_t_span = torch.linspace( + 0, 1, 11, device=torch.device("cuda"), dtype=torch.bfloat16 + ) # only capture at 10 steps + static_mu = torch.randn( + 1, 80, seq_len, device=torch.device("cuda"), dtype=torch.bfloat16 + ) + static_mask = torch.ones( + 1, 1, seq_len, device=torch.device("cuda"), dtype=torch.bfloat16 + ) + static_spks = torch.randn( + 1, 80, device=torch.device("cuda"), dtype=torch.bfloat16 + ) + static_cond = torch.randn( + 1, 80, seq_len, device=torch.device("cuda"), dtype=torch.float32 + ) + static_out = torch.randn( + 1, 80, seq_len, device=torch.device("cuda"), dtype=torch.bfloat16 + ) + + self._solve_euler_impl( + static_z, + t_span=static_t_span, + mu=static_mu, + mask=static_mask, + spks=static_spks, + cond=static_cond, + ) + torch.cuda.synchronize() + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + static_out = self._solve_euler_impl( + static_z, + t_span=static_t_span, + mu=static_mu, + mask=static_mask, + spks=static_spks, + cond=static_cond, + ) + + self.inference_buffers[seq_len] = { + "z": static_z, + "t_span": static_t_span, + "mu": static_mu, + "mask": static_mask, + "spks": static_spks, + "cond": static_cond, + "out": static_out, + } + self.inference_graphs[seq_len] = g + end_time = time.time() + print( + f"capture_inference for ConditionalCFM solve euler, time elapsed: {end_time - start_time}" + ) + + def solve_euler(self, x, t_span, mu, mask, spks, cond): + if hasattr(self, "inference_graphs") and len(self.inference_graphs) > 0: + curr_seq_len = x.shape[2] + + available_lengths = sorted(list(self.inference_graphs.keys())) + + if curr_seq_len <= max(available_lengths): + target_len = min(available_lengths, key=lambda x: abs(x - curr_seq_len)) + if target_len == curr_seq_len: + padded_x = x + padded_mu = mu + padded_mask = mask + if cond is not None: + padded_cond = cond + else: + padded_x = torch.randn( + (x.shape[0], x.shape[1], target_len), + dtype=x.dtype, + device=x.device, + ) + padded_x[:, :, :curr_seq_len] = x + + padded_mu = torch.randn( + (mu.shape[0], mu.shape[1], target_len), + dtype=mu.dtype, + device=mu.device, + ) + padded_mu[:, :, :curr_seq_len] = mu + + # FIXME(ys): uses zeros and maskgroupnorm + padded_mask = torch.ones( + (mask.shape[0], mask.shape[1], target_len), + dtype=mask.dtype, + device=mask.device, + ) + + if cond is not None: + padded_cond = torch.randn( + (cond.shape[0], cond.shape[1], target_len), + dtype=cond.dtype, + device=cond.device, + ) + padded_cond[:, :, :curr_seq_len] = cond + + buffer = self.inference_buffers[target_len] + buffer["z"].copy_(padded_x) + buffer["t_span"].copy_(t_span) + buffer["mu"].copy_(padded_mu) + buffer["mask"].copy_(padded_mask) + buffer["spks"].copy_(spks) + if cond is not None: + buffer["cond"].copy_(padded_cond) + + self.inference_graphs[target_len].replay() + + output = buffer["out"][:, :, :curr_seq_len] + return output + + return self._solve_euler_impl(x, t_span, mu, mask, spks, cond) + + def _solve_euler_impl(self, x, t_span, mu, mask, spks, cond): + """ + Fixed euler solver for ODEs. + Args: + x (torch.Tensor): random noise + t_span (torch.Tensor): n_timesteps interpolated + shape: (n_timesteps + 1,) + mu (torch.Tensor): output of encoder + shape: (batch_size, n_feats, mel_timesteps) + mask (torch.Tensor): output_mask + shape: (batch_size, 1, mel_timesteps) + spks (torch.Tensor, optional): speaker ids. Defaults to None. + shape: (batch_size, spk_emb_dim) + cond: Not used but kept for future purposes + """ + t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0] + t = t.unsqueeze(dim=0) + + # I am storing this because I can later plot it by putting a debugger here and saving it to a file + # Or in future might add like a return_all_steps flag + sol = [] + + for step in range(1, len(t_span)): + if self.inference_cfg_rate > 0: + x_double = torch.cat([x, x], dim=0) + mask_double = torch.cat([mask, mask], dim=0) + mu_double = torch.cat([mu, torch.zeros_like(mu)], dim=0) + t_double = torch.cat([t, t], dim=0) + spks_double = ( + torch.cat([spks, torch.zeros_like(spks)], dim=0) + if spks is not None + else None + ) + cond_double = torch.cat([cond, torch.zeros_like(cond)], dim=0) + + dphi_dt_double = self.forward_estimator( + x_double, mask_double, mu_double, t_double, spks_double, cond_double + ) + + dphi_dt, cfg_dphi_dt = torch.chunk(dphi_dt_double, 2, dim=0) + dphi_dt = ( + 1.0 + self.inference_cfg_rate + ) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt + else: + dphi_dt = self.forward_estimator(x, mask, mu, t, spks, cond) + + x = x + dt * dphi_dt + t = t + dt + sol.append(x) + if step < len(t_span) - 1: + dt = t_span[step + 1] - t + + return sol[-1] + + def forward_estimator(self, x, mask, mu, t, spks, cond): + if isinstance(self.estimator, torch.nn.Module): + return self.estimator.forward(x, mask, mu, t, spks, cond) + else: + ort_inputs = { + "x": x.cpu().numpy(), + "mask": mask.cpu().numpy(), + "mu": mu.cpu().numpy(), + "t": t.cpu().numpy(), + "spks": spks.cpu().numpy(), + "cond": cond.cpu().numpy(), + } + output = self.estimator.run(None, ort_inputs)[0] + return torch.tensor(output, dtype=x.dtype, device=x.device) + + def compute_loss(self, x1, mask, mu, spks=None, cond=None): + """Computes diffusion loss + + Args: + x1 (torch.Tensor): Target + shape: (batch_size, n_feats, mel_timesteps) + mask (torch.Tensor): target mask + shape: (batch_size, 1, mel_timesteps) + mu (torch.Tensor): output of encoder + shape: (batch_size, n_feats, mel_timesteps) + spks (torch.Tensor, optional): speaker embedding. Defaults to None. + shape: (batch_size, spk_emb_dim) + + Returns: + loss: conditional flow matching loss + y: conditional flow + shape: (batch_size, n_feats, mel_timesteps) + """ + b, _, t = mu.shape + + # random timestep + t = torch.rand([b, 1, 1], device=mu.device, dtype=mu.dtype) + if self.t_scheduler == "cosine": + t = 1 - torch.cos(t * 0.5 * torch.pi) + # sample noise p(x_0) + z = torch.randn_like(x1) + + y = (1 - (1 - self.sigma_min) * t) * z + t * x1 + u = x1 - (1 - self.sigma_min) * z + + # during training, we randomly drop condition to trade off mode coverage and sample fidelity + if self.training_cfg_rate > 0: + cfg_mask = torch.rand(b, device=x1.device) > self.training_cfg_rate + mu = mu * cfg_mask.view(-1, 1, 1) + spks = spks * cfg_mask.view(-1, 1) + cond = cond * cfg_mask.view(-1, 1, 1) + + pred = self.estimator(y, mask, mu, t.squeeze(), spks, cond) + loss = F.mse_loss(pred * mask, u * mask, reduction="sum") / ( + torch.sum(mask) * u.shape[1] + ) + return loss, y diff --git a/almeval/models/stepaudio/cosyvoice/flow/length_regulator.py b/almeval/models/stepaudio/cosyvoice/flow/length_regulator.py new file mode 100644 index 0000000000000000000000000000000000000000..19c07a0163327492332d1e5de12fa36bf99921b9 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/flow/length_regulator.py @@ -0,0 +1,65 @@ +# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Zhihao Du) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Tuple +import torch.nn as nn +import torch +from torch.nn import functional as F +from cosyvoice.utils.mask import make_pad_mask + + +class InterpolateRegulator(nn.Module): + def __init__( + self, + channels: int, + sampling_ratios: Tuple, + out_channels: int = None, + groups: int = 1, + ): + super().__init__() + self.sampling_ratios = sampling_ratios + out_channels = out_channels or channels + model = nn.ModuleList([]) + if len(sampling_ratios) > 0: + for _ in sampling_ratios: + module = nn.Conv1d(channels, channels, 3, 1, 1) + norm = nn.GroupNorm(groups, channels) + act = nn.Mish() + model.extend([module, norm, act]) + model.append(nn.Conv1d(channels, out_channels, 1, 1)) + self.model = nn.Sequential(*model) + + def forward(self, x, ylens=None): + # x in (B, T, D) + mask = (~make_pad_mask(ylens)).to(x).unsqueeze(-1) + x = F.interpolate( + x.transpose(1, 2).contiguous(), size=ylens.max(), mode="linear" + ) + out = self.model(x).transpose(1, 2).contiguous() + olens = ylens + return out * mask, olens + + def inference(self, x1, x2, mel_len1, mel_len2): + # x in (B, T, D) + x2 = F.interpolate( + x2.transpose(1, 2).contiguous(), size=mel_len2, mode="linear" + ) + if x1.shape[1] != 0: + x1 = F.interpolate( + x1.transpose(1, 2).contiguous(), size=mel_len1, mode="linear" + ) + x = torch.concat([x1, x2], dim=2) + else: + x = x2 + out = self.model(x).transpose(1, 2).contiguous() + return out, mel_len1 + mel_len2 diff --git a/almeval/models/stepaudio/cosyvoice/hifigan/f0_predictor.py b/almeval/models/stepaudio/cosyvoice/hifigan/f0_predictor.py new file mode 100644 index 0000000000000000000000000000000000000000..f465fcbecf9606e8627571aa4592289da50fec26 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/hifigan/f0_predictor.py @@ -0,0 +1,55 @@ +# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Kai Hu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import torch +import torch.nn as nn +from torch.nn.utils import weight_norm + + +class ConvRNNF0Predictor(nn.Module): + def __init__( + self, num_class: int = 1, in_channels: int = 80, cond_channels: int = 512 + ): + super().__init__() + + self.num_class = num_class + self.condnet = nn.Sequential( + weight_norm( + nn.Conv1d(in_channels, cond_channels, kernel_size=3, padding=1) + ), + nn.ELU(), + weight_norm( + nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1) + ), + nn.ELU(), + weight_norm( + nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1) + ), + nn.ELU(), + weight_norm( + nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1) + ), + nn.ELU(), + weight_norm( + nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1) + ), + nn.ELU(), + ) + self.classifier = nn.Linear( + in_features=cond_channels, out_features=self.num_class + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.condnet(x) + x = x.transpose(1, 2) + return torch.abs(self.classifier(x).squeeze(-1)) diff --git a/almeval/models/stepaudio/cosyvoice/hifigan/generator.py b/almeval/models/stepaudio/cosyvoice/hifigan/generator.py new file mode 100644 index 0000000000000000000000000000000000000000..4d02c03da59775503b85998127b295289d284287 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/hifigan/generator.py @@ -0,0 +1,566 @@ +# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Kai Hu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""HIFI-GAN""" + +import typing as tp +import time +import numpy as np +from scipy.signal import get_window +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn import Conv1d +from torch.nn import ConvTranspose1d +from torch.nn.utils import remove_weight_norm +from torch.nn.utils import weight_norm +from torch.distributions.uniform import Uniform + +from cosyvoice.transformer.activation import Snake +from cosyvoice.utils.common import get_padding +from cosyvoice.utils.common import init_weights + + +"""hifigan based generator implementation. + +This code is modified from https://github.com/jik876/hifi-gan + ,https://github.com/kan-bayashi/ParallelWaveGAN and + https://github.com/NVIDIA/BigVGAN + +""" + + +class ResBlock(torch.nn.Module): + """Residual block module in HiFiGAN/BigVGAN.""" + + def __init__( + self, + channels: int = 512, + kernel_size: int = 3, + dilations: tp.List[int] = [1, 3, 5], + ): + super(ResBlock, self).__init__() + self.convs1 = nn.ModuleList() + self.convs2 = nn.ModuleList() + + for dilation in dilations: + self.convs1.append( + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation, + padding=get_padding(kernel_size, dilation), + ) + ) + ) + self.convs2.append( + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1), + ) + ) + ) + self.convs1.apply(init_weights) + self.convs2.apply(init_weights) + self.activations1 = nn.ModuleList( + [Snake(channels, alpha_logscale=False) for _ in range(len(self.convs1))] + ) + self.activations2 = nn.ModuleList( + [Snake(channels, alpha_logscale=False) for _ in range(len(self.convs2))] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for idx in range(len(self.convs1)): + xt = self.activations1[idx](x) + xt = self.convs1[idx](xt) + xt = self.activations2[idx](xt) + xt = self.convs2[idx](xt) + x = xt + x + return x + + def remove_weight_norm(self): + for idx in range(len(self.convs1)): + remove_weight_norm(self.convs1[idx]) + remove_weight_norm(self.convs2[idx]) + + +class SineGen(torch.nn.Module): + """Definition of sine generator + SineGen(samp_rate, harmonic_num = 0, + sine_amp = 0.1, noise_std = 0.003, + voiced_threshold = 0, + flag_for_pulse=False) + samp_rate: sampling rate in Hz + harmonic_num: number of harmonic overtones (default 0) + sine_amp: amplitude of sine-wavefrom (default 0.1) + noise_std: std of Gaussian noise (default 0.003) + voiced_thoreshold: F0 threshold for U/V classification (default 0) + flag_for_pulse: this SinGen is used inside PulseGen (default False) + Note: when flag_for_pulse is True, the first time step of a voiced + segment is always sin(np.pi) or cos(0) + """ + + def __init__( + self, + samp_rate, + harmonic_num=0, + sine_amp=0.1, + noise_std=0.003, + voiced_threshold=0, + ): + super(SineGen, self).__init__() + self.sine_amp = sine_amp + self.noise_std = noise_std + self.harmonic_num = harmonic_num + self.sampling_rate = samp_rate + self.voiced_threshold = voiced_threshold + + def _f02uv(self, f0): + # generate uv signal + uv = (f0 > self.voiced_threshold).type(torch.float32) + return uv + + @torch.no_grad() + def forward(self, f0): + """ + :param f0: [B, 1, sample_len], Hz + :return: [B, 1, sample_len] + """ + + F_mat = torch.zeros((f0.size(0), self.harmonic_num + 1, f0.size(-1))).to( + f0.device + ) + for i in range(self.harmonic_num + 1): + F_mat[:, i : i + 1, :] = f0 * (i + 1) / self.sampling_rate + + theta_mat = 2 * np.pi * (torch.cumsum(F_mat, dim=-1) % 1) + u_dist = Uniform(low=-np.pi, high=np.pi) + phase_vec = u_dist.sample( + sample_shape=(f0.size(0), self.harmonic_num + 1, 1) + ).to(F_mat.device) + phase_vec[:, 0, :] = 0 + + # generate sine waveforms + sine_waves = self.sine_amp * torch.sin(theta_mat + phase_vec) + + # generate uv signal + uv = self._f02uv(f0) + + # noise: for unvoiced should be similar to sine_amp + # std = self.sine_amp/3 -> max value ~ self.sine_amp + # . for voiced regions is self.noise_std + noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3 + noise = noise_amp * torch.randn_like(sine_waves) + + # first: set the unvoiced part to 0 by uv + # then: additive noise + sine_waves = sine_waves * uv + noise + return sine_waves, uv, noise + + +class SourceModuleHnNSF(torch.nn.Module): + """SourceModule for hn-nsf + SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1, + add_noise_std=0.003, voiced_threshod=0) + sampling_rate: sampling_rate in Hz + harmonic_num: number of harmonic above F0 (default: 0) + sine_amp: amplitude of sine source signal (default: 0.1) + add_noise_std: std of additive Gaussian noise (default: 0.003) + note that amplitude of noise in unvoiced is decided + by sine_amp + voiced_threshold: threhold to set U/V given F0 (default: 0) + Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) + F0_sampled (batchsize, length, 1) + Sine_source (batchsize, length, 1) + noise_source (batchsize, length 1) + uv (batchsize, length, 1) + """ + + def __init__( + self, + sampling_rate, + upsample_scale, + harmonic_num=0, + sine_amp=0.1, + add_noise_std=0.003, + voiced_threshod=0, + ): + super(SourceModuleHnNSF, self).__init__() + + self.sine_amp = sine_amp + self.noise_std = add_noise_std + + # to produce sine waveforms + self.l_sin_gen = SineGen( + sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshod + ) + + # to merge source harmonics into a single excitation + self.l_linear = torch.nn.Linear(harmonic_num + 1, 1) + self.l_tanh = torch.nn.Tanh() + + def forward(self, x): + """ + Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) + F0_sampled (batchsize, length, 1) + Sine_source (batchsize, length, 1) + noise_source (batchsize, length 1) + """ + # source for harmonic branch + with torch.no_grad(): + sine_wavs, uv, _ = self.l_sin_gen(x.transpose(1, 2)) + sine_wavs = sine_wavs.transpose(1, 2) + uv = uv.transpose(1, 2) + sine_merge = self.l_tanh(self.l_linear(sine_wavs)) + + # source for noise branch, in the same shape as uv + noise = torch.randn_like(uv) * self.sine_amp / 3 + return sine_merge, noise, uv + + +class HiFTGenerator(nn.Module): + """ + HiFTNet Generator: Neural Source Filter + ISTFTNet + https://arxiv.org/abs/2309.09493 + """ + + def __init__( + self, + in_channels: int = 80, + base_channels: int = 512, + nb_harmonics: int = 8, + sampling_rate: int = 22050, + nsf_alpha: float = 0.1, + nsf_sigma: float = 0.003, + nsf_voiced_threshold: float = 10, + upsample_rates: tp.List[int] = [8, 8], + upsample_kernel_sizes: tp.List[int] = [16, 16], + istft_params: tp.Dict[str, int] = {"n_fft": 16, "hop_len": 4}, + resblock_kernel_sizes: tp.List[int] = [3, 7, 11], + resblock_dilation_sizes: tp.List[tp.List[int]] = [ + [1, 3, 5], + [1, 3, 5], + [1, 3, 5], + ], + source_resblock_kernel_sizes: tp.List[int] = [7, 11], + source_resblock_dilation_sizes: tp.List[tp.List[int]] = [[1, 3, 5], [1, 3, 5]], + lrelu_slope: float = 0.1, + audio_limit: float = 0.99, + f0_predictor: torch.nn.Module = None, + ): + super(HiFTGenerator, self).__init__() + + self.out_channels = 1 + self.nb_harmonics = nb_harmonics + self.sampling_rate = sampling_rate + self.istft_params = istft_params + self.lrelu_slope = lrelu_slope + self.audio_limit = audio_limit + + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_rates) + self.upsample_rates = upsample_rates + self.m_source = SourceModuleHnNSF( + sampling_rate=sampling_rate, + upsample_scale=np.prod(upsample_rates) * istft_params["hop_len"], + harmonic_num=nb_harmonics, + sine_amp=nsf_alpha, + add_noise_std=nsf_sigma, + voiced_threshod=nsf_voiced_threshold, + ) + self.f0_upsamp = torch.nn.Upsample( + scale_factor=np.prod(upsample_rates) * istft_params["hop_len"] + ) + + self.conv_pre = weight_norm(Conv1d(in_channels, base_channels, 7, 1, padding=3)) + + # Up + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): + self.ups.append( + weight_norm( + ConvTranspose1d( + base_channels // (2**i), + base_channels // (2 ** (i + 1)), + k, + u, + padding=(k - u) // 2, + ) + ) + ) + + # Down + self.source_downs = nn.ModuleList() + self.source_resblocks = nn.ModuleList() + downsample_rates = [1] + upsample_rates[::-1][:-1] + downsample_cum_rates = np.cumprod(downsample_rates) + for i, (u, k, d) in enumerate( + zip( + downsample_cum_rates[::-1], + source_resblock_kernel_sizes, + source_resblock_dilation_sizes, + ) + ): + if u == 1: + self.source_downs.append( + Conv1d( + istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), 1, 1 + ) + ) + else: + self.source_downs.append( + Conv1d( + istft_params["n_fft"] + 2, + base_channels // (2 ** (i + 1)), + u * 2, + u, + padding=(u // 2), + ) + ) + + self.source_resblocks.append( + ResBlock(base_channels // (2 ** (i + 1)), k, d) + ) + + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = base_channels // (2 ** (i + 1)) + for _, (k, d) in enumerate( + zip(resblock_kernel_sizes, resblock_dilation_sizes) + ): + self.resblocks.append(ResBlock(ch, k, d)) + + self.conv_post = weight_norm( + Conv1d(ch, istft_params["n_fft"] + 2, 7, 1, padding=3) + ) + self.ups.apply(init_weights) + self.conv_post.apply(init_weights) + self.reflection_pad = nn.ReflectionPad1d((1, 0)) + self.stft_window = torch.from_numpy( + get_window("hann", istft_params["n_fft"], fftbins=True).astype(np.float32) + ).cuda() + self.f0_predictor = f0_predictor + self.inference_buffers = {} + self.inference_graphs = {} + + def _f02source(self, f0: torch.Tensor) -> torch.Tensor: + f0 = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t + + har_source, _, _ = self.m_source(f0) + return har_source.transpose(1, 2) + + def _stft(self, x): + spec = torch.stft( + x, + self.istft_params["n_fft"], + self.istft_params["hop_len"], + self.istft_params["n_fft"], + window=self.stft_window, + return_complex=True, + ) + spec = torch.view_as_real(spec) # [B, F, TT, 2] + return spec[..., 0], spec[..., 1] + + def _istft(self, magnitude, phase): + magnitude = torch.clip(magnitude, max=1e2) + real = magnitude * torch.cos(phase) + img = magnitude * torch.sin(phase) + inverse_transform = torch.istft( + torch.complex(real, img), + self.istft_params["n_fft"], + self.istft_params["hop_len"], + self.istft_params["n_fft"], + window=self.stft_window, + ) + return inverse_transform + + def forward( + self, x: torch.Tensor, cache_source: torch.Tensor = torch.zeros(1, 1, 0) + ) -> torch.Tensor: + f0 = self.f0_predictor(x) + s = self._f02source(f0) + + # use cache_source to avoid glitch + if cache_source.shape[2] != 0: + s[:, :, : cache_source.shape[2]] = cache_source + + s_stft_real, s_stft_imag = self._stft(s.squeeze(1)) + s_stft = torch.cat([s_stft_real, s_stft_imag], dim=1) + + x = self.conv_pre(x) + for i in range(self.num_upsamples): + x = F.leaky_relu(x, self.lrelu_slope) + x = self.ups[i](x) + + if i == self.num_upsamples - 1: + x = self.reflection_pad(x) + + # fusion + si = self.source_downs[i](s_stft) + si = self.source_resblocks[i](si) + x = x + si + + xs = None + for j in range(self.num_kernels): + if xs is None: + xs = self.resblocks[i * self.num_kernels + j](x) + else: + xs += self.resblocks[i * self.num_kernels + j](x) + x = xs / self.num_kernels + + x = F.leaky_relu(x) + x = self.conv_post(x) + magnitude = torch.exp(x[:, : self.istft_params["n_fft"] // 2 + 1, :]) + phase = torch.sin( + x[:, self.istft_params["n_fft"] // 2 + 1 :, :] + ) # actually, sin is redundancy + + x = self._istft(magnitude, phase) + x = torch.clamp(x, -self.audio_limit, self.audio_limit) + return x, s + + def remove_weight_norm(self): + print("Removing weight norm...") + for l in self.ups: + remove_weight_norm(l) + for l in self.resblocks: + l.remove_weight_norm() + remove_weight_norm(self.conv_pre) + remove_weight_norm(self.conv_post) + self.source_module.remove_weight_norm() + for l in self.source_downs: + remove_weight_norm(l) + for l in self.source_resblocks: + l.remove_weight_norm() + + @torch.inference_mode() + def _inference_impl(self, mel: torch.Tensor, s_stft: torch.Tensor) -> torch.Tensor: + x = self.conv_pre(mel) + for i in range(self.num_upsamples): + x = F.leaky_relu(x, self.lrelu_slope) + x = self.ups[i](x) + + if i == self.num_upsamples - 1: + x = self.reflection_pad(x) + + # fusion + si = self.source_downs[i](s_stft) + si = self.source_resblocks[i](si) + x = x + si + + xs = None + for j in range(self.num_kernels): + if xs is None: + xs = self.resblocks[i * self.num_kernels + j](x) + else: + xs += self.resblocks[i * self.num_kernels + j](x) + x = xs / self.num_kernels + + x = F.leaky_relu(x) + x = self.conv_post(x) + magnitude = torch.exp(x[:, : self.istft_params["n_fft"] // 2 + 1, :]) + phase = torch.sin( + x[:, self.istft_params["n_fft"] // 2 + 1 :, :] + ) # actually, sin is redundancy + # print(f"mel: {mel.shape}, magnitude: {magnitude.shape}, phase: {phase.shape}") + return magnitude, phase + + @torch.inference_mode() + def inference( + self, mel: torch.Tensor, cache_source: torch.Tensor = torch.zeros(1, 1, 0) + ) -> torch.Tensor: + curr_seq_len = mel.shape[2] + f0 = self.f0_predictor(mel) + s = self._f02source(f0) + s_stft_real, s_stft_imag = self._stft(s.squeeze(1)) + s_stft = torch.cat([s_stft_real, s_stft_imag], dim=1) + + target_len = None + for seq_len in sorted(self.inference_buffers.keys()): + if curr_seq_len <= seq_len: + target_len = seq_len + break + + if target_len is not None: + buffer = self.inference_buffers[target_len] + + if curr_seq_len < target_len: + padded_mel = torch.zeros_like(buffer["mel"]) + padded_mel[:, :, :curr_seq_len] = mel + buffer["mel"].copy_(padded_mel) + padded_s_stft = torch.zeros_like(buffer["s_stft"]) + cur_s_stft_len = s_stft.shape[2] + padded_s_stft[:, :, :cur_s_stft_len] = s_stft + buffer["s_stft"].copy_(padded_s_stft) + + else: + buffer["mel"].copy_(mel) + buffer["s_stft"].copy_(s_stft) + cur_s_stft_len = s_stft.shape[2] + + self.inference_graphs[target_len].replay() + + magnitude, phase = ( + buffer["magnitude"][:, :, :cur_s_stft_len], + buffer["phase"][:, :, :cur_s_stft_len], + ) + else: + magnitude, phase = self._inference_impl(mel=mel, s_stft=s_stft) + + x = self._istft(magnitude, phase) + x = torch.clamp(x, -self.audio_limit, self.audio_limit) + return x, s + + @torch.inference_mode() + def capture_inference(self, seq_len_to_capture=[64, 128, 256, 512, 1024]): + start_time = time.time() + print( + f"capture inference for HiFTGenerator with seq_len_to_capture: {seq_len_to_capture}" + ) + for seq_len in seq_len_to_capture: + mel = torch.randn( + 1, 80, seq_len, device=torch.device("cuda"), dtype=torch.float32 + ) + f0 = self.f0_predictor(mel) + s = self._f02source(f0) + s_stft_real, s_stft_imag = self._stft(s.squeeze(1)) + s_stft = torch.cat([s_stft_real, s_stft_imag], dim=1) + + magnitude, phase = self._inference_impl(mel=mel, s_stft=s_stft) + torch.cuda.synchronize() + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + magnitude, phase = self._inference_impl(mel=mel, s_stft=s_stft) + inference_buffer = { + "mel": mel, + "s_stft": s_stft, + "magnitude": magnitude, + "phase": phase, + } + self.inference_buffers[seq_len] = inference_buffer + self.inference_graphs[seq_len] = g + + end_time = time.time() + print( + f"capture inference for HiFTGenerator with seq_len_to_capture: {seq_len_to_capture} takes {end_time - start_time} seconds" + ) diff --git a/almeval/models/stepaudio/cosyvoice/matcha/audio.py b/almeval/models/stepaudio/cosyvoice/matcha/audio.py new file mode 100644 index 0000000000000000000000000000000000000000..0a9b8db2a96e5c06ce04681ca373477dfa69fea2 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/matcha/audio.py @@ -0,0 +1,90 @@ +import numpy as np +import torch +import torch.utils.data +from librosa.filters import mel as librosa_mel_fn +from scipy.io.wavfile import read + +MAX_WAV_VALUE = 32768.0 + + +def load_wav(full_path): + sampling_rate, data = read(full_path) + return data, sampling_rate + + +def dynamic_range_compression(x, C=1, clip_val=1e-5): + return np.log(np.clip(x, a_min=clip_val, a_max=None) * C) + + +def dynamic_range_decompression(x, C=1): + return np.exp(x) / C + + +def dynamic_range_compression_torch(x, C=1, clip_val=1e-5): + return torch.log(torch.clamp(x, min=clip_val) * C) + + +def dynamic_range_decompression_torch(x, C=1): + return torch.exp(x) / C + + +def spectral_normalize_torch(magnitudes): + output = dynamic_range_compression_torch(magnitudes) + return output + + +def spectral_de_normalize_torch(magnitudes): + output = dynamic_range_decompression_torch(magnitudes) + return output + + +mel_basis = {} +hann_window = {} + + +def mel_spectrogram( + y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False +): + if torch.min(y) < -1.0: + print("min value is ", torch.min(y)) + if torch.max(y) > 1.0: + print("max value is ", torch.max(y)) + + global mel_basis, hann_window # pylint: disable=global-statement + if f"{str(fmax)}_{str(y.device)}" not in mel_basis: + mel = librosa_mel_fn( + sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax + ) + mel_basis[str(fmax) + "_" + str(y.device)] = ( + torch.from_numpy(mel).float().to(y.device) + ) + hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device) + + y = torch.nn.functional.pad( + y.unsqueeze(1), + (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), + mode="reflect", + ) + y = y.squeeze(1) + + spec = torch.view_as_real( + torch.stft( + y, + n_fft, + hop_length=hop_size, + win_length=win_size, + window=hann_window[str(y.device)], + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + ) + + spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)) + + spec = torch.matmul(mel_basis[str(fmax) + "_" + str(y.device)], spec) + spec = spectral_normalize_torch(spec) + + return spec diff --git a/almeval/models/stepaudio/cosyvoice/matcha/decoder.py b/almeval/models/stepaudio/cosyvoice/matcha/decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..d571dcca22a404a1272940cf1480820b58a10549 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/matcha/decoder.py @@ -0,0 +1,511 @@ +import math +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from conformer import ConformerBlock +from diffusers.models.activations import get_activation +from einops import pack, rearrange, repeat + +from cosyvoice.matcha.transformer import BasicTransformerBlock + + +class SinusoidalPosEmb(torch.nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + assert self.dim % 2 == 0, "SinusoidalPosEmb requires dim to be even" + + def forward(self, x, scale=1000): + if x.ndim < 1: + x = x.unsqueeze(0) + device = x.device + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb) + emb = scale * x.unsqueeze(1) * emb.unsqueeze(0) + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb + + +class MaskedGroupNorm(nn.GroupNorm): + """ + Masked verstion of the Group normalization. + + Based on: https://github.com/ptrblck/pytorch_misc/blob/20e8ea93bd458b88f921a87e2d4001a4eb753a02/batch_norm_manual.py + + Receives a N-dim tensor of sequence lengths per batch element + along with the regular input for masking. + + Check pytorch's GroupNorm implementation for argument details. + """ + + def __init__(self, num_groups, num_channels, eps=1e-5, affine=True): + super(MaskedGroupNorm, self).__init__(num_groups, num_channels, eps, affine) + + def forward(self, inp, mask=None): + assert ( + inp.shape[1] % self.num_groups == 0 + ), "Feature size not divisible by groups" + + # 计算有效长度 + seq_lengths = mask.sum(-1, keepdim=True) # [batch_size, 1] + + # 将输入reshape为groups + features_per_group = inp.shape[1] // self.num_groups + inp_r = inp.reshape( + inp.shape[0], self.num_groups, features_per_group, inp.shape[-1] + ) + mask_r = mask.unsqueeze(1) # [batch_size, 1, 1, length] + + # 计算masked mean和variance + masked_inp = inp_r * mask_r + n = seq_lengths * features_per_group # 每组的有效元素数量 + mean = masked_inp.sum([2, 3], keepdim=True) / (n.view(-1, 1, 1, 1) + 1e-5) + var = ((masked_inp - mean * mask_r) ** 2).sum([2, 3], keepdim=True) / ( + n.view(-1, 1, 1, 1) + 1e-5 + ) + + # 标准化 + inp_r = (inp_r - mean) / (torch.sqrt(var + self.eps)) + out = inp_r.reshape(inp.shape[0], self.num_channels, inp.shape[-1]) + + # 应用仿射变换 + if self.affine: + out = out * self.weight[None, :, None] + self.bias[None, :, None] + + return out + + +class Block1D(torch.nn.Module): + def __init__(self, dim, dim_out, groups=8): + super().__init__() + self.block = torch.nn.Sequential( + torch.nn.Conv1d(dim, dim_out, 3, padding=1), + torch.nn.GroupNorm(groups, dim_out), + # MaskedGroupNorm(groups, dim_out), + nn.Mish(), + ) + + def forward(self, x, mask): + output = self.block(x * mask) + return output * mask + return x * mask + + +class ResnetBlock1D(torch.nn.Module): + def __init__(self, dim, dim_out, time_emb_dim, groups=8): + super().__init__() + self.mlp = torch.nn.Sequential( + nn.Mish(), torch.nn.Linear(time_emb_dim, dim_out) + ) + + self.block1 = Block1D(dim, dim_out, groups=groups) + self.block2 = Block1D(dim_out, dim_out, groups=groups) + + self.res_conv = torch.nn.Conv1d(dim, dim_out, 1) + + def forward(self, x, mask, time_emb): + h = self.block1(x, mask) + h += self.mlp(time_emb).unsqueeze(-1) + h = self.block2(h, mask) + output = h + self.res_conv(x * mask) + return output + + +class Downsample1D(nn.Module): + def __init__(self, dim): + super().__init__() + self.conv = torch.nn.Conv1d(dim, dim, 3, 2, 1) + + def forward(self, x): + return self.conv(x) + + +class TimestepEmbedding(nn.Module): + def __init__( + self, + in_channels: int, + time_embed_dim: int, + act_fn: str = "silu", + out_dim: int = None, + post_act_fn: Optional[str] = None, + cond_proj_dim=None, + ): + super().__init__() + + self.linear_1 = nn.Linear(in_channels, time_embed_dim) + + if cond_proj_dim is not None: + self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False) + else: + self.cond_proj = None + + self.act = get_activation(act_fn) + + if out_dim is not None: + time_embed_dim_out = out_dim + else: + time_embed_dim_out = time_embed_dim + self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out) + + if post_act_fn is None: + self.post_act = None + else: + self.post_act = get_activation(post_act_fn) + + def forward(self, sample, condition=None): + if condition is not None: + sample = sample + self.cond_proj(condition) + sample = self.linear_1(sample) + + if self.act is not None: + sample = self.act(sample) + + sample = self.linear_2(sample) + + if self.post_act is not None: + sample = self.post_act(sample) + return sample + + +class Upsample1D(nn.Module): + """A 1D upsampling layer with an optional convolution. + + Parameters: + channels (`int`): + number of channels in the inputs and outputs. + use_conv (`bool`, default `False`): + option to use a convolution. + use_conv_transpose (`bool`, default `False`): + option to use a convolution transpose. + out_channels (`int`, optional): + number of output channels. Defaults to `channels`. + """ + + def __init__( + self, + channels, + use_conv=False, + use_conv_transpose=True, + out_channels=None, + name="conv", + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_conv_transpose = use_conv_transpose + self.name = name + + self.conv = None + if use_conv_transpose: + self.conv = nn.ConvTranspose1d(channels, self.out_channels, 4, 2, 1) + elif use_conv: + self.conv = nn.Conv1d(self.channels, self.out_channels, 3, padding=1) + + def forward(self, inputs): + assert inputs.shape[1] == self.channels + if self.use_conv_transpose: + return self.conv(inputs) + + outputs = F.interpolate(inputs, scale_factor=2.0, mode="nearest") + + if self.use_conv: + outputs = self.conv(outputs) + + return outputs + + +class ConformerWrapper(ConformerBlock): + def __init__( # pylint: disable=useless-super-delegation + self, + *, + dim, + dim_head=64, + heads=8, + ff_mult=4, + conv_expansion_factor=2, + conv_kernel_size=31, + attn_dropout=0, + ff_dropout=0, + conv_dropout=0, + conv_causal=False, + ): + super().__init__( + dim=dim, + dim_head=dim_head, + heads=heads, + ff_mult=ff_mult, + conv_expansion_factor=conv_expansion_factor, + conv_kernel_size=conv_kernel_size, + attn_dropout=attn_dropout, + ff_dropout=ff_dropout, + conv_dropout=conv_dropout, + conv_causal=conv_causal, + ) + + def forward( + self, + hidden_states, + attention_mask, + encoder_hidden_states=None, + encoder_attention_mask=None, + timestep=None, + ): + return super().forward(x=hidden_states, mask=attention_mask.bool()) + + +class Decoder(nn.Module): + def __init__( + self, + in_channels, + out_channels, + channels=(256, 256), + dropout=0.05, + attention_head_dim=64, + n_blocks=1, + num_mid_blocks=2, + num_heads=4, + act_fn="snake", + down_block_type="transformer", + mid_block_type="transformer", + up_block_type="transformer", + ): + super().__init__() + channels = tuple(channels) + self.in_channels = in_channels + self.out_channels = out_channels + + self.time_embeddings = SinusoidalPosEmb(in_channels) + time_embed_dim = channels[0] * 4 + self.time_mlp = TimestepEmbedding( + in_channels=in_channels, + time_embed_dim=time_embed_dim, + act_fn="silu", + ) + + self.down_blocks = nn.ModuleList([]) + self.mid_blocks = nn.ModuleList([]) + self.up_blocks = nn.ModuleList([]) + + output_channel = in_channels + for i in range(len(channels)): # pylint: disable=consider-using-enumerate + input_channel = output_channel + output_channel = channels[i] + is_last = i == len(channels) - 1 + resnet = ResnetBlock1D( + dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim + ) + transformer_blocks = nn.ModuleList( + [ + self.get_block( + down_block_type, + output_channel, + attention_head_dim, + num_heads, + dropout, + act_fn, + ) + for _ in range(n_blocks) + ] + ) + downsample = ( + Downsample1D(output_channel) + if not is_last + else nn.Conv1d(output_channel, output_channel, 3, padding=1) + ) + + self.down_blocks.append( + nn.ModuleList([resnet, transformer_blocks, downsample]) + ) + + for i in range(num_mid_blocks): + input_channel = channels[-1] + out_channels = channels[-1] + + resnet = ResnetBlock1D( + dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim + ) + + transformer_blocks = nn.ModuleList( + [ + self.get_block( + mid_block_type, + output_channel, + attention_head_dim, + num_heads, + dropout, + act_fn, + ) + for _ in range(n_blocks) + ] + ) + + self.mid_blocks.append(nn.ModuleList([resnet, transformer_blocks])) + + channels = channels[::-1] + (channels[0],) + for i in range(len(channels) - 1): + input_channel = channels[i] + output_channel = channels[i + 1] + is_last = i == len(channels) - 2 + + resnet = ResnetBlock1D( + dim=2 * input_channel, + dim_out=output_channel, + time_emb_dim=time_embed_dim, + ) + transformer_blocks = nn.ModuleList( + [ + self.get_block( + up_block_type, + output_channel, + attention_head_dim, + num_heads, + dropout, + act_fn, + ) + for _ in range(n_blocks) + ] + ) + upsample = ( + Upsample1D(output_channel, use_conv_transpose=True) + if not is_last + else nn.Conv1d(output_channel, output_channel, 3, padding=1) + ) + + self.up_blocks.append(nn.ModuleList([resnet, transformer_blocks, upsample])) + + self.final_block = Block1D(channels[-1], channels[-1]) + self.final_proj = nn.Conv1d(channels[-1], self.out_channels, 1) + + self.initialize_weights() + # nn.init.normal_(self.final_proj.weight) + + @staticmethod + def get_block(block_type, dim, attention_head_dim, num_heads, dropout, act_fn): + if block_type == "conformer": + block = ConformerWrapper( + dim=dim, + dim_head=attention_head_dim, + heads=num_heads, + ff_mult=1, + conv_expansion_factor=2, + ff_dropout=dropout, + attn_dropout=dropout, + conv_dropout=dropout, + conv_kernel_size=31, + ) + elif block_type == "transformer": + block = BasicTransformerBlock( + dim=dim, + num_attention_heads=num_heads, + attention_head_dim=attention_head_dim, + dropout=dropout, + activation_fn=act_fn, + ) + else: + raise ValueError(f"Unknown block type {block_type}") + + return block + + def initialize_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv1d): + nn.init.kaiming_normal_(m.weight, nonlinearity="relu") + + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + elif isinstance(m, nn.GroupNorm): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + elif isinstance(m, nn.Linear): + nn.init.kaiming_normal_(m.weight, nonlinearity="relu") + + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def forward(self, x, mask, mu, t, spks=None, cond=None): + """Forward pass of the UNet1DConditional model. + + Args: + x (torch.Tensor): shape (batch_size, in_channels, time) + mask (_type_): shape (batch_size, 1, time) + t (_type_): shape (batch_size) + spks (_type_, optional): shape: (batch_size, condition_channels). Defaults to None. + cond (_type_, optional): placeholder for future use. Defaults to None. + + Raises: + ValueError: _description_ + ValueError: _description_ + + Returns: + _type_: _description_ + """ + + t = self.time_embeddings(t) + t = self.time_mlp(t) + + x = pack([x, mu], "b * t")[0] + + if spks is not None: + spks = repeat(spks, "b c -> b c t", t=x.shape[-1]) + x = pack([x, spks], "b * t")[0] + + hiddens = [] + masks = [mask] + for resnet, transformer_blocks, downsample in self.down_blocks: + mask_down = masks[-1] + x = resnet(x, mask_down, t) + x = rearrange(x, "b c t -> b t c") + mask_down = rearrange(mask_down, "b 1 t -> b t") + for transformer_block in transformer_blocks: + x = transformer_block( + hidden_states=x, + attention_mask=mask_down, + timestep=t, + ) + x = rearrange(x, "b t c -> b c t") + mask_down = rearrange(mask_down, "b t -> b 1 t") + hiddens.append(x) # Save hidden states for skip connections + x = downsample(x * mask_down) + masks.append(mask_down[:, :, ::2]) + + masks = masks[:-1] + mask_mid = masks[-1] + + for resnet, transformer_blocks in self.mid_blocks: + x = resnet(x, mask_mid, t) + x = rearrange(x, "b c t -> b t c") + mask_mid = rearrange(mask_mid, "b 1 t -> b t") + for transformer_block in transformer_blocks: + x = transformer_block( + hidden_states=x, + attention_mask=mask_mid, + timestep=t, + ) + x = rearrange(x, "b t c -> b c t") + mask_mid = rearrange(mask_mid, "b t -> b 1 t") + + for resnet, transformer_blocks, upsample in self.up_blocks: + mask_up = masks.pop() + x = resnet(pack([x, hiddens.pop()], "b * t")[0], mask_up, t) + x = rearrange(x, "b c t -> b t c") + mask_up = rearrange(mask_up, "b 1 t -> b t") + for transformer_block in transformer_blocks: + x = transformer_block( + hidden_states=x, + attention_mask=mask_up, + timestep=t, + ) + x = rearrange(x, "b t c -> b c t") + mask_up = rearrange(mask_up, "b t -> b 1 t") + x = upsample(x * mask_up) + + x = self.final_block(x, mask_up) + output = self.final_proj(x * mask_up) + + return output * mask diff --git a/almeval/models/stepaudio/cosyvoice/matcha/flow_matching.py b/almeval/models/stepaudio/cosyvoice/matcha/flow_matching.py new file mode 100644 index 0000000000000000000000000000000000000000..678153c2eb345aa66ff5fbe643025bd2332315f8 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/matcha/flow_matching.py @@ -0,0 +1,141 @@ +from abc import ABC + +import torch +import torch.nn.functional as F + +from cosyvoice.matcha.decoder import Decoder + + +class BASECFM(torch.nn.Module, ABC): + def __init__( + self, + n_feats, + cfm_params, + n_spks=1, + spk_emb_dim=128, + ): + super().__init__() + self.n_feats = n_feats + self.n_spks = n_spks + self.spk_emb_dim = spk_emb_dim + self.solver = cfm_params.solver + if hasattr(cfm_params, "sigma_min"): + self.sigma_min = cfm_params.sigma_min + else: + self.sigma_min = 1e-4 + + self.estimator = None + + @torch.inference_mode() + def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None): + """Forward diffusion + + Args: + mu (torch.Tensor): output of encoder + shape: (batch_size, n_feats, mel_timesteps) + mask (torch.Tensor): output_mask + shape: (batch_size, 1, mel_timesteps) + n_timesteps (int): number of diffusion steps + temperature (float, optional): temperature for scaling noise. Defaults to 1.0. + spks (torch.Tensor, optional): speaker ids. Defaults to None. + shape: (batch_size, spk_emb_dim) + cond: Not used but kept for future purposes + + Returns: + sample: generated mel-spectrogram + shape: (batch_size, n_feats, mel_timesteps) + """ + z = torch.randn_like(mu) * temperature + t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device) + return self.solve_euler( + z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond + ) + + def solve_euler(self, x, t_span, mu, mask, spks, cond): + """ + Fixed euler solver for ODEs. + Args: + x (torch.Tensor): random noise + t_span (torch.Tensor): n_timesteps interpolated + shape: (n_timesteps + 1,) + mu (torch.Tensor): output of encoder + shape: (batch_size, n_feats, mel_timesteps) + mask (torch.Tensor): output_mask + shape: (batch_size, 1, mel_timesteps) + spks (torch.Tensor, optional): speaker ids. Defaults to None. + shape: (batch_size, spk_emb_dim) + cond: Not used but kept for future purposes + """ + t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0] + + # I am storing this because I can later plot it by putting a debugger here and saving it to a file + # Or in future might add like a return_all_steps flag + sol = [] + + for step in range(1, len(t_span)): + dphi_dt = self.estimator(x, mask, mu, t, spks, cond) + + x = x + dt * dphi_dt + t = t + dt + sol.append(x) + if step < len(t_span) - 1: + dt = t_span[step + 1] - t + + return sol[-1] + + def compute_loss(self, x1, mask, mu, spks=None, cond=None): + """Computes diffusion loss + + Args: + x1 (torch.Tensor): Target + shape: (batch_size, n_feats, mel_timesteps) + mask (torch.Tensor): target mask + shape: (batch_size, 1, mel_timesteps) + mu (torch.Tensor): output of encoder + shape: (batch_size, n_feats, mel_timesteps) + spks (torch.Tensor, optional): speaker embedding. Defaults to None. + shape: (batch_size, spk_emb_dim) + + Returns: + loss: conditional flow matching loss + y: conditional flow + shape: (batch_size, n_feats, mel_timesteps) + """ + b, _, t = mu.shape + + # random timestep + t = torch.rand([b, 1, 1], device=mu.device, dtype=mu.dtype) + # sample noise p(x_0) + z = torch.randn_like(x1) + + y = (1 - (1 - self.sigma_min) * t) * z + t * x1 + u = x1 - (1 - self.sigma_min) * z + + loss = F.mse_loss( + self.estimator(y, mask, mu, t.squeeze(), spks), u, reduction="sum" + ) / (torch.sum(mask) * u.shape[1]) + return loss, y + + +class CFM(BASECFM): + def __init__( + self, + in_channels, + out_channel, + cfm_params, + decoder_params, + n_spks=1, + spk_emb_dim=64, + ): + super().__init__( + n_feats=in_channels, + cfm_params=cfm_params, + n_spks=n_spks, + spk_emb_dim=spk_emb_dim, + ) + + in_channels = in_channels + (spk_emb_dim if n_spks > 1 else 0) + # Just change the architecture of the estimator here + self.estimator = Decoder( + in_channels=in_channels, out_channels=out_channel, **decoder_params + ) diff --git a/almeval/models/stepaudio/cosyvoice/matcha/transformer.py b/almeval/models/stepaudio/cosyvoice/matcha/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..12712dc79590d7eafa95f745d0be95c3b13bb0a4 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/matcha/transformer.py @@ -0,0 +1,443 @@ +from typing import Any, Dict, Optional + +import torch +import torch.nn as nn +from diffusers.models.attention import ( + GEGLU, + GELU, + AdaLayerNorm, + AdaLayerNormZero, + ApproximateGELU, +) +from diffusers.models.attention_processor import Attention +from diffusers.models.lora import LoRACompatibleLinear +from diffusers.utils.torch_utils import maybe_allow_in_graph + + +class SnakeBeta(nn.Module): + """ + A modified Snake function which uses separate parameters for the magnitude of the periodic components + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter that controls frequency + - beta - trainable parameter that controls magnitude + References: + - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: + https://arxiv.org/abs/2006.08195 + Examples: + >>> a1 = snakebeta(256) + >>> x = torch.randn(256) + >>> x = a1(x) + """ + + def __init__( + self, + in_features, + out_features, + alpha=1.0, + alpha_trainable=True, + alpha_logscale=True, + ): + """ + Initialization. + INPUT: + - in_features: shape of the input + - alpha - trainable parameter that controls frequency + - beta - trainable parameter that controls magnitude + alpha is initialized to 1 by default, higher values = higher-frequency. + beta is initialized to 1 by default, higher values = higher-magnitude. + alpha will be trained along with the rest of your model. + """ + super().__init__() + self.in_features = ( + out_features if isinstance(out_features, list) else [out_features] + ) + self.proj = LoRACompatibleLinear(in_features, out_features) + + # initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # log scale alphas initialized to zeros + self.alpha = nn.Parameter(torch.zeros(self.in_features) * alpha) + self.beta = nn.Parameter(torch.zeros(self.in_features) * alpha) + else: # linear scale alphas initialized to ones + self.alpha = nn.Parameter(torch.ones(self.in_features) * alpha) + self.beta = nn.Parameter(torch.ones(self.in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + self.beta.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + """ + Forward pass of the function. + Applies the function to the input elementwise. + SnakeBeta ∶= x + 1/b * sin^2 (xa) + """ + x = self.proj(x) + if self.alpha_logscale: + alpha = torch.exp(self.alpha) + beta = torch.exp(self.beta) + else: + alpha = self.alpha + beta = self.beta + + x = x + (1.0 / (beta + self.no_div_by_zero)) * torch.pow( + torch.sin(x * alpha), 2 + ) + + return x + + +class FeedForward(nn.Module): + r""" + A feed-forward layer. + + Parameters: + dim (`int`): The number of channels in the input. + dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`. + mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + final_dropout (`bool` *optional*, defaults to False): Apply a final dropout. + """ + + def __init__( + self, + dim: int, + dim_out: Optional[int] = None, + mult: int = 4, + dropout: float = 0.0, + activation_fn: str = "geglu", + final_dropout: bool = False, + ): + super().__init__() + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + if activation_fn == "gelu": + act_fn = GELU(dim, inner_dim) + if activation_fn == "gelu-approximate": + act_fn = GELU(dim, inner_dim, approximate="tanh") + elif activation_fn == "geglu": + act_fn = GEGLU(dim, inner_dim) + elif activation_fn == "geglu-approximate": + act_fn = ApproximateGELU(dim, inner_dim) + elif activation_fn == "snakebeta": + act_fn = SnakeBeta(dim, inner_dim) + + self.net = nn.ModuleList([]) + # project in + self.net.append(act_fn) + # project dropout + self.net.append(nn.Dropout(dropout)) + # project out + self.net.append(LoRACompatibleLinear(inner_dim, dim_out)) + # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout + if final_dropout: + self.net.append(nn.Dropout(dropout)) + + def forward(self, hidden_states): + for module in self.net: + hidden_states = module(hidden_states) + return hidden_states + + +@maybe_allow_in_graph +class BasicTransformerBlock(nn.Module): + r""" + A basic Transformer block. + + Parameters: + dim (`int`): The number of channels in the input and output. + num_attention_heads (`int`): The number of heads to use for multi-head attention. + attention_head_dim (`int`): The number of channels in each head. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. + only_cross_attention (`bool`, *optional*): + Whether to use only cross-attention layers. In this case two cross attention layers are used. + double_self_attention (`bool`, *optional*): + Whether to use two self-attention layers. In this case no cross attention layers are used. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + num_embeds_ada_norm (: + obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. + attention_bias (: + obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + dropout=0.0, + cross_attention_dim: Optional[int] = None, + activation_fn: str = "geglu", + num_embeds_ada_norm: Optional[int] = None, + attention_bias: bool = False, + only_cross_attention: bool = False, + double_self_attention: bool = False, + upcast_attention: bool = False, + norm_elementwise_affine: bool = True, + norm_type: str = "layer_norm", + final_dropout: bool = False, + ): + super().__init__() + self.only_cross_attention = only_cross_attention + + self.use_ada_layer_norm_zero = ( + num_embeds_ada_norm is not None + ) and norm_type == "ada_norm_zero" + self.use_ada_layer_norm = ( + num_embeds_ada_norm is not None + ) and norm_type == "ada_norm" + + if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: + raise ValueError( + f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" + f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." + ) + + # Define 3 blocks. Each block has its own normalization layer. + # 1. Self-Attn + if self.use_ada_layer_norm: + self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) + elif self.use_ada_layer_norm_zero: + self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm) + else: + self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=cross_attention_dim if only_cross_attention else None, + upcast_attention=upcast_attention, + ) + + # 2. Cross-Attn + if cross_attention_dim is not None or double_self_attention: + # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. + # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during + # the second cross attention block. + self.norm2 = ( + AdaLayerNorm(dim, num_embeds_ada_norm) + if self.use_ada_layer_norm + else nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) + ) + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=( + cross_attention_dim if not double_self_attention else None + ), + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + # scale_qk=False, # uncomment this to not to use flash attention + ) # is self-attn if encoder_hidden_states is none + else: + self.norm2 = None + self.attn2 = None + + # 3. Feed-forward + self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) + self.ff = FeedForward( + dim, + dropout=dropout, + activation_fn=activation_fn, + final_dropout=final_dropout, + ) + + # let chunk size default to None + self._chunk_size = None + self._chunk_dim = 0 + + def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int): + # Sets chunk feed-forward + self._chunk_size = chunk_size + self._chunk_dim = dim + + def forward_native( + self, + hidden_states: torch.FloatTensor, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + timestep: Optional[torch.LongTensor] = None, + cross_attention_kwargs: Dict[str, Any] = None, + class_labels: Optional[torch.LongTensor] = None, + ): + # Notice that normalization is always applied before the real computation in the following blocks. + # 1. Self-Attention + if self.use_ada_layer_norm: + norm_hidden_states = self.norm1(hidden_states, timestep) + elif self.use_ada_layer_norm_zero: + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( + hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + else: + norm_hidden_states = self.norm1(hidden_states) + + cross_attention_kwargs = ( + cross_attention_kwargs if cross_attention_kwargs is not None else {} + ) + + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=( + encoder_hidden_states if self.only_cross_attention else None + ), + attention_mask=( + encoder_attention_mask if self.only_cross_attention else attention_mask + ), + **cross_attention_kwargs, + ) + if self.use_ada_layer_norm_zero: + attn_output = gate_msa.unsqueeze(1) * attn_output + hidden_states = attn_output + hidden_states + + # 2. Cross-Attention + if self.attn2 is not None: + norm_hidden_states = ( + self.norm2(hidden_states, timestep) + if self.use_ada_layer_norm + else self.norm2(hidden_states) + ) + + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + **cross_attention_kwargs, + ) + hidden_states = attn_output + hidden_states + + # 3. Feed-forward + norm_hidden_states = self.norm3(hidden_states) + + if self.use_ada_layer_norm_zero: + norm_hidden_states = ( + norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ) + + if self._chunk_size is not None: + # "feed_forward_chunk_size" can be used to save memory + if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0: + raise ValueError( + f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." + ) + + num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size + ff_output = torch.cat( + [ + self.ff(hid_slice) + for hid_slice in norm_hidden_states.chunk( + num_chunks, dim=self._chunk_dim + ) + ], + dim=self._chunk_dim, + ) + else: + ff_output = self.ff(norm_hidden_states) + + if self.use_ada_layer_norm_zero: + ff_output = gate_mlp.unsqueeze(1) * ff_output + + hidden_states = ff_output + hidden_states + + return hidden_states + + def forward( + self, + hidden_states: torch.FloatTensor, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + timestep: Optional[torch.LongTensor] = None, + cross_attention_kwargs: Dict[str, Any] = None, + class_labels: Optional[torch.LongTensor] = None, + ): + # Notice that normalization is always applied before the real computation in the following blocks. + # 1. Self-Attention + if self.use_ada_layer_norm: + norm_hidden_states = self.norm1(hidden_states, timestep) + elif self.use_ada_layer_norm_zero: + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( + hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + else: + norm_hidden_states = self.norm1(hidden_states) + + cross_attention_kwargs = ( + cross_attention_kwargs if cross_attention_kwargs is not None else {} + ) + + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=( + encoder_hidden_states if self.only_cross_attention else None + ), + attention_mask=( + encoder_attention_mask if self.only_cross_attention else attention_mask + ), + **cross_attention_kwargs, + ) + if self.use_ada_layer_norm_zero: + attn_output = gate_msa.unsqueeze(1) * attn_output + hidden_states = attn_output + hidden_states + + # 2. Cross-Attention + if self.attn2 is not None: + norm_hidden_states = ( + self.norm2(hidden_states, timestep) + if self.use_ada_layer_norm + else self.norm2(hidden_states) + ) + + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + **cross_attention_kwargs, + ) + hidden_states = attn_output + hidden_states + + # 3. Feed-forward + norm_hidden_states = self.norm3(hidden_states) + + if self.use_ada_layer_norm_zero: + norm_hidden_states = ( + norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ) + + if self._chunk_size is not None: + # "feed_forward_chunk_size" can be used to save memory + if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0: + raise ValueError( + f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." + ) + + num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size + ff_output = torch.cat( + [ + self.ff(hid_slice) + for hid_slice in norm_hidden_states.chunk( + num_chunks, dim=self._chunk_dim + ) + ], + dim=self._chunk_dim, + ) + else: + ff_output = self.ff(norm_hidden_states) + + if self.use_ada_layer_norm_zero: + ff_output = gate_mlp.unsqueeze(1) * ff_output + + hidden_states = ff_output + hidden_states + + return hidden_states diff --git a/almeval/models/stepaudio/cosyvoice/transformer/__init__.py b/almeval/models/stepaudio/cosyvoice/transformer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/almeval/models/stepaudio/cosyvoice/transformer/activation.py b/almeval/models/stepaudio/cosyvoice/transformer/activation.py new file mode 100644 index 0000000000000000000000000000000000000000..507fff4ba2b72af3aaff87cc85f98cf3612434d7 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/transformer/activation.py @@ -0,0 +1,87 @@ +# Copyright (c) 2020 Johns Hopkins University (Shinji Watanabe) +# 2020 Northwestern Polytechnical University (Pengcheng Guo) +# 2020 Mobvoi Inc (Binbin Zhang) +# 2024 Alibaba Inc (Xiang Lyu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Swish() activation function for Conformer.""" + +import torch +from torch import nn, sin, pow +from torch.nn import Parameter + + +class Swish(torch.nn.Module): + """Construct an Swish object.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Return Swish activation function.""" + return x * torch.sigmoid(x) + + +# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license. +# LICENSE is in incl_licenses directory. +class Snake(nn.Module): + """ + Implementation of a sine-based periodic activation function + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter + References: + - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: + https://arxiv.org/abs/2006.08195 + Examples: + >>> a1 = snake(256) + >>> x = torch.randn(256) + >>> x = a1(x) + """ + + def __init__( + self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False + ): + """ + Initialization. + INPUT: + - in_features: shape of the input + - alpha: trainable parameter + alpha is initialized to 1 by default, higher values = higher-frequency. + alpha will be trained along with the rest of your model. + """ + super(Snake, self).__init__() + self.in_features = in_features + + # initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # log scale alphas initialized to zeros + self.alpha = Parameter(torch.zeros(in_features) * alpha) + else: # linear scale alphas initialized to ones + self.alpha = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + """ + Forward pass of the function. + Applies the function to the input elementwise. + Snake ∶= x + 1/a * sin^2 (xa) + """ + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] + if self.alpha_logscale: + alpha = torch.exp(alpha) + x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2) + + return x diff --git a/almeval/models/stepaudio/cosyvoice/transformer/attention.py b/almeval/models/stepaudio/cosyvoice/transformer/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..4b2a17b2471d5486d38ad36312ff2f41f70ed5c1 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/transformer/attention.py @@ -0,0 +1,322 @@ +# Copyright (c) 2019 Shigeki Karita +# 2020 Mobvoi Inc (Binbin Zhang) +# 2022 Xingchen Song (sxc19@mails.tsinghua.edu.cn) +# 2024 Alibaba Inc (Xiang Lyu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Multi-Head Attention layer definition.""" + +import math +from typing import Tuple + +import torch +from torch import nn + + +class MultiHeadedAttention(nn.Module): + """Multi-Head Attention layer. + + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + + """ + + def __init__( + self, n_head: int, n_feat: int, dropout_rate: float, key_bias: bool = True + ): + """Construct an MultiHeadedAttention object.""" + super().__init__() + assert n_feat % n_head == 0 + # We assume d_v always equals d_k + self.d_k = n_feat // n_head + self.h = n_head + self.linear_q = nn.Linear(n_feat, n_feat) + self.linear_k = nn.Linear(n_feat, n_feat, bias=key_bias) + self.linear_v = nn.Linear(n_feat, n_feat) + self.linear_out = nn.Linear(n_feat, n_feat) + self.dropout = nn.Dropout(p=dropout_rate) + + def forward_qkv( + self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Transform query, key and value. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + + Returns: + torch.Tensor: Transformed query tensor, size + (#batch, n_head, time1, d_k). + torch.Tensor: Transformed key tensor, size + (#batch, n_head, time2, d_k). + torch.Tensor: Transformed value tensor, size + (#batch, n_head, time2, d_k). + + """ + n_batch = query.size(0) + q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k) + k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k) + v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k) + q = q.transpose(1, 2) # (batch, head, time1, d_k) + k = k.transpose(1, 2) # (batch, head, time2, d_k) + v = v.transpose(1, 2) # (batch, head, time2, d_k) + + return q, k, v + + def forward_attention( + self, + value: torch.Tensor, + scores: torch.Tensor, + mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), + ) -> torch.Tensor: + """Compute attention context vector. + + Args: + value (torch.Tensor): Transformed value, size + (#batch, n_head, time2, d_k). + scores (torch.Tensor): Attention score, size + (#batch, n_head, time1, time2). + mask (torch.Tensor): Mask, size (#batch, 1, time2) or + (#batch, time1, time2), (0, 0, 0) means fake mask. + + Returns: + torch.Tensor: Transformed value (#batch, time1, d_model) + weighted by the attention score (#batch, time1, time2). + + """ + n_batch = value.size(0) + # NOTE(xcsong): When will `if mask.size(2) > 0` be True? + # 1. onnx(16/4) [WHY? Because we feed real cache & real mask for the + # 1st chunk to ease the onnx export.] + # 2. pytorch training + if mask.size(2) > 0: # time2 > 0 + mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) + # For last chunk, time2 might be larger than scores.size(-1) + mask = mask[:, :, :, : scores.size(-1)] # (batch, 1, *, time2) + scores = scores.masked_fill(mask, -float("inf")) + attn = torch.softmax(scores, dim=-1).masked_fill( + mask, 0.0 + ) # (batch, head, time1, time2) + # NOTE(xcsong): When will `if mask.size(2) > 0` be False? + # 1. onnx(16/-1, -1/-1, 16/0) + # 2. jit (16/-1, -1/-1, 16/0, 16/4) + else: + attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) + + p_attn = self.dropout(attn) + x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) + x = ( + x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) + ) # (batch, time1, d_model) + + return self.linear_out(x) # (batch, time1, d_model) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), + pos_emb: torch.Tensor = torch.empty(0), + cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute scaled dot product attention. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + 1.When applying cross attention between decoder and encoder, + the batch padding mask for input is in (#batch, 1, T) shape. + 2.When applying self attention of encoder, + the mask is in (#batch, T, T) shape. + 3.When applying self attention of decoder, + the mask is in (#batch, L, L) shape. + 4.If the different position in decoder see different block + of the encoder, such as Mocha, the passed in mask could be + in (#batch, L, T) shape. But there is no such case in current + CosyVoice. + cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2), + where `cache_t == chunk_size * num_decoding_left_chunks` + and `head * d_k == size` + + + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2) + where `cache_t == chunk_size * num_decoding_left_chunks` + and `head * d_k == size` + + """ + q, k, v = self.forward_qkv(query, key, value) + + # NOTE(xcsong): + # when export onnx model, for 1st chunk, we feed + # cache(1, head, 0, d_k * 2) (16/-1, -1/-1, 16/0 mode) + # or cache(1, head, real_cache_t, d_k * 2) (16/4 mode). + # In all modes, `if cache.size(0) > 0` will alwayse be `True` + # and we will always do splitting and + # concatnation(this will simplify onnx export). Note that + # it's OK to concat & split zero-shaped tensors(see code below). + # when export jit model, for 1st chunk, we always feed + # cache(0, 0, 0, 0) since jit supports dynamic if-branch. + # >>> a = torch.ones((1, 2, 0, 4)) + # >>> b = torch.ones((1, 2, 3, 4)) + # >>> c = torch.cat((a, b), dim=2) + # >>> torch.equal(b, c) # True + # >>> d = torch.split(a, 2, dim=-1) + # >>> torch.equal(d[0], d[1]) # True + if cache.size(0) > 0: + key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1) + k = torch.cat([key_cache, k], dim=2) + v = torch.cat([value_cache, v], dim=2) + # NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's + # non-trivial to calculate `next_cache_start` here. + new_cache = torch.cat((k, v), dim=-1) + + scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) + return self.forward_attention(v, scores, mask), new_cache + + +class RelPositionMultiHeadedAttention(MultiHeadedAttention): + """Multi-Head Attention layer with relative position encoding. + Paper: https://arxiv.org/abs/1901.02860 + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + """ + + def __init__( + self, n_head: int, n_feat: int, dropout_rate: float, key_bias: bool = True + ): + """Construct an RelPositionMultiHeadedAttention object.""" + super().__init__(n_head, n_feat, dropout_rate, key_bias) + # linear transformation for positional encoding + self.linear_pos = nn.Linear(n_feat, n_feat, bias=False) + # these two learnable bias are used in matrix c and matrix d + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k)) + self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k)) + torch.nn.init.xavier_uniform_(self.pos_bias_u) + torch.nn.init.xavier_uniform_(self.pos_bias_v) + + def rel_shift(self, x: torch.Tensor) -> torch.Tensor: + """Compute relative positional encoding. + + Args: + x (torch.Tensor): Input tensor (batch, head, time1, 2*time1-1). + time1 means the length of query vector. + + Returns: + torch.Tensor: Output tensor. + + """ + zero_pad = torch.zeros( + (x.size()[0], x.size()[1], x.size()[2], 1), device=x.device, dtype=x.dtype + ) + x_padded = torch.cat([zero_pad, x], dim=-1) + + x_padded = x_padded.view(x.size()[0], x.size()[1], x.size(3) + 1, x.size(2)) + x = x_padded[:, :, 1:].view_as(x)[ + :, :, :, : x.size(-1) // 2 + 1 + ] # only keep the positions from 0 to time2 + return x + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), + pos_emb: torch.Tensor = torch.empty(0), + cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute 'Scaled Dot Product Attention' with rel. positional encoding. + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2), (0, 0, 0) means fake mask. + pos_emb (torch.Tensor): Positional embedding tensor + (#batch, time2, size). + cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2), + where `cache_t == chunk_size * num_decoding_left_chunks` + and `head * d_k == size` + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2) + where `cache_t == chunk_size * num_decoding_left_chunks` + and `head * d_k == size` + """ + q, k, v = self.forward_qkv(query, key, value) + q = q.transpose(1, 2) # (batch, time1, head, d_k) + + # NOTE(xcsong): + # when export onnx model, for 1st chunk, we feed + # cache(1, head, 0, d_k * 2) (16/-1, -1/-1, 16/0 mode) + # or cache(1, head, real_cache_t, d_k * 2) (16/4 mode). + # In all modes, `if cache.size(0) > 0` will alwayse be `True` + # and we will always do splitting and + # concatnation(this will simplify onnx export). Note that + # it's OK to concat & split zero-shaped tensors(see code below). + # when export jit model, for 1st chunk, we always feed + # cache(0, 0, 0, 0) since jit supports dynamic if-branch. + # >>> a = torch.ones((1, 2, 0, 4)) + # >>> b = torch.ones((1, 2, 3, 4)) + # >>> c = torch.cat((a, b), dim=2) + # >>> torch.equal(b, c) # True + # >>> d = torch.split(a, 2, dim=-1) + # >>> torch.equal(d[0], d[1]) # True + if cache.size(0) > 0: + key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1) + k = torch.cat([key_cache, k], dim=2) + v = torch.cat([value_cache, v], dim=2) + # NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's + # non-trivial to calculate `next_cache_start` here. + new_cache = torch.cat((k, v), dim=-1) + + n_batch_pos = pos_emb.size(0) + p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k) + p = p.transpose(1, 2) # (batch, head, time1, d_k) + + # (batch, head, time1, d_k) + q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) + # (batch, head, time1, d_k) + q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) + + # compute attention score + # first compute matrix a and matrix c + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + # (batch, head, time1, time2) + matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) + + # compute matrix b and matrix d + # (batch, head, time1, time2) + matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) + # NOTE(Xiang Lyu): Keep rel_shift since espnet rel_pos_emb is used + if matrix_ac.shape != matrix_bd.shape: + matrix_bd = self.rel_shift(matrix_bd) + + scores = (matrix_ac + matrix_bd) / math.sqrt( + self.d_k + ) # (batch, head, time1, time2) + + return self.forward_attention(v, scores, mask), new_cache diff --git a/almeval/models/stepaudio/cosyvoice/transformer/convolution.py b/almeval/models/stepaudio/cosyvoice/transformer/convolution.py new file mode 100644 index 0000000000000000000000000000000000000000..ef3dfc2a88146db06ea69e897d6aab9c28895bdb --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/transformer/convolution.py @@ -0,0 +1,147 @@ +# Copyright (c) 2020 Mobvoi Inc. (authors: Binbin Zhang, Di Wu) +# 2024 Alibaba Inc (Xiang Lyu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Modified from ESPnet(https://github.com/espnet/espnet) +"""ConvolutionModule definition.""" + +from typing import Tuple + +import torch +from torch import nn + + +class ConvolutionModule(nn.Module): + """ConvolutionModule in Conformer model.""" + + def __init__( + self, + channels: int, + kernel_size: int = 15, + activation: nn.Module = nn.ReLU(), + norm: str = "batch_norm", + causal: bool = False, + bias: bool = True, + ): + """Construct an ConvolutionModule object. + Args: + channels (int): The number of channels of conv layers. + kernel_size (int): Kernel size of conv layers. + causal (int): Whether use causal convolution or not + """ + super().__init__() + + self.pointwise_conv1 = nn.Conv1d( + channels, + 2 * channels, + kernel_size=1, + stride=1, + padding=0, + bias=bias, + ) + # self.lorder is used to distinguish if it's a causal convolution, + # if self.lorder > 0: it's a causal convolution, the input will be + # padded with self.lorder frames on the left in forward. + # else: it's a symmetrical convolution + if causal: + padding = 0 + self.lorder = kernel_size - 1 + else: + # kernel_size should be an odd number for none causal convolution + assert (kernel_size - 1) % 2 == 0 + padding = (kernel_size - 1) // 2 + self.lorder = 0 + self.depthwise_conv = nn.Conv1d( + channels, + channels, + kernel_size, + stride=1, + padding=padding, + groups=channels, + bias=bias, + ) + + assert norm in ["batch_norm", "layer_norm"] + if norm == "batch_norm": + self.use_layer_norm = False + self.norm = nn.BatchNorm1d(channels) + else: + self.use_layer_norm = True + self.norm = nn.LayerNorm(channels) + + self.pointwise_conv2 = nn.Conv1d( + channels, + channels, + kernel_size=1, + stride=1, + padding=0, + bias=bias, + ) + self.activation = activation + + def forward( + self, + x: torch.Tensor, + mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), + cache: torch.Tensor = torch.zeros((0, 0, 0)), + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute convolution module. + Args: + x (torch.Tensor): Input tensor (#batch, time, channels). + mask_pad (torch.Tensor): used for batch padding (#batch, 1, time), + (0, 0, 0) means fake mask. + cache (torch.Tensor): left context cache, it is only + used in causal convolution (#batch, channels, cache_t), + (0, 0, 0) meas fake cache. + Returns: + torch.Tensor: Output tensor (#batch, time, channels). + """ + # exchange the temporal dimension and the feature dimension + x = x.transpose(1, 2) # (#batch, channels, time) + + # mask batch padding + if mask_pad.size(2) > 0: # time > 0 + x.masked_fill_(~mask_pad, 0.0) + + if self.lorder > 0: + if cache.size(2) == 0: # cache_t == 0 + x = nn.functional.pad(x, (self.lorder, 0), "constant", 0.0) + else: + assert cache.size(0) == x.size(0) # equal batch + assert cache.size(1) == x.size(1) # equal channel + x = torch.cat((cache, x), dim=2) + assert x.size(2) > self.lorder + new_cache = x[:, :, -self.lorder :] + else: + # It's better we just return None if no cache is required, + # However, for JIT export, here we just fake one tensor instead of + # None. + new_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device) + + # GLU mechanism + x = self.pointwise_conv1(x) # (batch, 2*channel, dim) + x = nn.functional.glu(x, dim=1) # (batch, channel, dim) + + # 1D Depthwise Conv + x = self.depthwise_conv(x) + if self.use_layer_norm: + x = x.transpose(1, 2) + x = self.activation(self.norm(x)) + if self.use_layer_norm: + x = x.transpose(1, 2) + x = self.pointwise_conv2(x) + # mask batch padding + if mask_pad.size(2) > 0: # time > 0 + x.masked_fill_(~mask_pad, 0.0) + + return x.transpose(1, 2), new_cache diff --git a/almeval/models/stepaudio/cosyvoice/transformer/decoder.py b/almeval/models/stepaudio/cosyvoice/transformer/decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..db1e30ce3ded1724828a50d12f798ee41df27ef0 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/transformer/decoder.py @@ -0,0 +1,418 @@ +# Copyright (c) 2021 Mobvoi Inc. (authors: Binbin Zhang, Di Wu) +# 2024 Alibaba Inc (Xiang Lyu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Modified from ESPnet(https://github.com/espnet/espnet) +"""Decoder definition.""" +from typing import Tuple, List, Optional + +import torch +import torch.utils.checkpoint as ckpt +import logging + +from cosyvoice.transformer.decoder_layer import DecoderLayer +from cosyvoice.transformer.positionwise_feed_forward import ( + PositionwiseFeedForward, +) +from cosyvoice.utils.class_utils import ( + COSYVOICE_EMB_CLASSES, + COSYVOICE_ATTENTION_CLASSES, + COSYVOICE_ACTIVATION_CLASSES, +) +from cosyvoice.utils.mask import subsequent_mask, make_pad_mask + + +class TransformerDecoder(torch.nn.Module): + """Base class of Transfomer decoder module. + Args: + vocab_size: output dim + encoder_output_size: dimension of attention + attention_heads: the number of heads of multi head attention + linear_units: the hidden units number of position-wise feedforward + num_blocks: the number of decoder blocks + dropout_rate: dropout rate + self_attention_dropout_rate: dropout rate for attention + input_layer: input layer type + use_output_layer: whether to use output layer + pos_enc_class: PositionalEncoding or ScaledPositionalEncoding + normalize_before: + True: use layer_norm before each sub-block of a layer. + False: use layer_norm after each sub-block of a layer. + src_attention: if false, encoder-decoder cross attention is not + applied, such as CIF model + key_bias: whether use bias in attention.linear_k, False for whisper models. + gradient_checkpointing: rerunning a forward-pass segment for each + checkpointed segment during backward. + tie_word_embedding: Tie or clone module weights depending of whether we are + using TorchScript or not + """ + + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + normalize_before: bool = True, + src_attention: bool = True, + key_bias: bool = True, + activation_type: str = "relu", + gradient_checkpointing: bool = False, + tie_word_embedding: bool = False, + ): + super().__init__() + attention_dim = encoder_output_size + activation = COSYVOICE_ACTIVATION_CLASSES[activation_type]() + + self.embed = torch.nn.Sequential( + ( + torch.nn.Identity() + if input_layer == "no_pos" + else torch.nn.Embedding(vocab_size, attention_dim) + ), + COSYVOICE_EMB_CLASSES[input_layer](attention_dim, positional_dropout_rate), + ) + + self.normalize_before = normalize_before + self.after_norm = torch.nn.LayerNorm(attention_dim, eps=1e-5) + self.use_output_layer = use_output_layer + if use_output_layer: + self.output_layer = torch.nn.Linear(attention_dim, vocab_size) + else: + self.output_layer = torch.nn.Identity() + self.num_blocks = num_blocks + self.decoders = torch.nn.ModuleList( + [ + DecoderLayer( + attention_dim, + COSYVOICE_ATTENTION_CLASSES["selfattn"]( + attention_heads, + attention_dim, + self_attention_dropout_rate, + key_bias, + ), + ( + COSYVOICE_ATTENTION_CLASSES["selfattn"]( + attention_heads, + attention_dim, + src_attention_dropout_rate, + key_bias, + ) + if src_attention + else None + ), + PositionwiseFeedForward( + attention_dim, linear_units, dropout_rate, activation + ), + dropout_rate, + normalize_before, + ) + for _ in range(self.num_blocks) + ] + ) + + self.gradient_checkpointing = gradient_checkpointing + self.tie_word_embedding = tie_word_embedding + + def forward( + self, + memory: torch.Tensor, + memory_mask: torch.Tensor, + ys_in_pad: torch.Tensor, + ys_in_lens: torch.Tensor, + r_ys_in_pad: torch.Tensor = torch.empty(0), + reverse_weight: float = 0.0, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Forward decoder. + Args: + memory: encoded memory, float32 (batch, maxlen_in, feat) + memory_mask: encoder memory mask, (batch, 1, maxlen_in) + ys_in_pad: padded input token ids, int64 (batch, maxlen_out) + ys_in_lens: input lengths of this batch (batch) + r_ys_in_pad: not used in transformer decoder, in order to unify api + with bidirectional decoder + reverse_weight: not used in transformer decoder, in order to unify + api with bidirectional decode + Returns: + (tuple): tuple containing: + x: decoded token score before softmax (batch, maxlen_out, + vocab_size) if use_output_layer is True, + torch.tensor(0.0), in order to unify api with bidirectional decoder + olens: (batch, ) + NOTE(xcsong): + We pass the `__call__` method of the modules instead of `forward` to the + checkpointing API because `__call__` attaches all the hooks of the module. + https://discuss.pytorch.org/t/any-different-between-model-input-and-model-forward-input/3690/2 + """ + tgt = ys_in_pad + maxlen = tgt.size(1) + # tgt_mask: (B, 1, L) + tgt_mask = ~make_pad_mask(ys_in_lens, maxlen).unsqueeze(1) + tgt_mask = tgt_mask.to(tgt.device) + # m: (1, L, L) + m = subsequent_mask(tgt_mask.size(-1), device=tgt_mask.device).unsqueeze(0) + # tgt_mask: (B, L, L) + tgt_mask = tgt_mask & m + x, _ = self.embed(tgt) + if self.gradient_checkpointing and self.training: + x = self.forward_layers_checkpointed(x, tgt_mask, memory, memory_mask) + else: + x = self.forward_layers(x, tgt_mask, memory, memory_mask) + if self.normalize_before: + x = self.after_norm(x) + if self.use_output_layer: + x = self.output_layer(x) + olens = tgt_mask.sum(1) + return x, torch.tensor(0.0), olens + + def forward_layers( + self, + x: torch.Tensor, + tgt_mask: torch.Tensor, + memory: torch.Tensor, + memory_mask: torch.Tensor, + ) -> torch.Tensor: + for layer in self.decoders: + x, tgt_mask, memory, memory_mask = layer(x, tgt_mask, memory, memory_mask) + return x + + @torch.jit.unused + def forward_layers_checkpointed( + self, + x: torch.Tensor, + tgt_mask: torch.Tensor, + memory: torch.Tensor, + memory_mask: torch.Tensor, + ) -> torch.Tensor: + for layer in self.decoders: + x, tgt_mask, memory, memory_mask = ckpt.checkpoint( + layer.__call__, x, tgt_mask, memory, memory_mask + ) + return x + + def forward_one_step( + self, + memory: torch.Tensor, + memory_mask: torch.Tensor, + tgt: torch.Tensor, + tgt_mask: torch.Tensor, + cache: Optional[List[torch.Tensor]] = None, + ) -> Tuple[torch.Tensor, List[torch.Tensor]]: + """Forward one step. + This is only used for decoding. + Args: + memory: encoded memory, float32 (batch, maxlen_in, feat) + memory_mask: encoded memory mask, (batch, 1, maxlen_in) + tgt: input token ids, int64 (batch, maxlen_out) + tgt_mask: input token mask, (batch, maxlen_out) + dtype=torch.uint8 in PyTorch 1.2- + dtype=torch.bool in PyTorch 1.2+ (include 1.2) + cache: cached output list of (batch, max_time_out-1, size) + Returns: + y, cache: NN output value and cache per `self.decoders`. + y.shape` is (batch, maxlen_out, token) + """ + x, _ = self.embed(tgt) + new_cache = [] + for i, decoder in enumerate(self.decoders): + if cache is None: + c = None + else: + c = cache[i] + x, tgt_mask, memory, memory_mask = decoder( + x, tgt_mask, memory, memory_mask, cache=c + ) + new_cache.append(x) + if self.normalize_before: + y = self.after_norm(x[:, -1]) + else: + y = x[:, -1] + if self.use_output_layer: + y = torch.log_softmax(self.output_layer(y), dim=-1) + return y, new_cache + + def tie_or_clone_weights(self, jit_mode: bool = True): + """Tie or clone module weights (between word_emb and output_layer) + depending of whether we are using TorchScript or not""" + if not self.use_output_layer: + return + if jit_mode: + logging.info("clone emb.weight to output.weight") + self.output_layer.weight = torch.nn.Parameter(self.embed[0].weight.clone()) + else: + logging.info("tie emb.weight with output.weight") + self.output_layer.weight = self.embed[0].weight + + if getattr(self.output_layer, "bias", None) is not None: + self.output_layer.bias.data = torch.nn.functional.pad( + self.output_layer.bias.data, + ( + 0, + self.output_layer.weight.shape[0] - self.output_layer.bias.shape[0], + ), + "constant", + 0, + ) + + +class BiTransformerDecoder(torch.nn.Module): + """Base class of Transfomer decoder module. + Args: + vocab_size: output dim + encoder_output_size: dimension of attention + attention_heads: the number of heads of multi head attention + linear_units: the hidden units number of position-wise feedforward + num_blocks: the number of decoder blocks + r_num_blocks: the number of right to left decoder blocks + dropout_rate: dropout rate + self_attention_dropout_rate: dropout rate for attention + input_layer: input layer type + use_output_layer: whether to use output layer + pos_enc_class: PositionalEncoding or ScaledPositionalEncoding + normalize_before: + True: use layer_norm before each sub-block of a layer. + False: use layer_norm after each sub-block of a layer. + key_bias: whether use bias in attention.linear_k, False for whisper models. + """ + + def __init__( + self, + vocab_size: int, + encoder_output_size: int, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + r_num_blocks: int = 0, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + self_attention_dropout_rate: float = 0.0, + src_attention_dropout_rate: float = 0.0, + input_layer: str = "embed", + use_output_layer: bool = True, + normalize_before: bool = True, + key_bias: bool = True, + gradient_checkpointing: bool = False, + tie_word_embedding: bool = False, + ): + + super().__init__() + self.tie_word_embedding = tie_word_embedding + self.left_decoder = TransformerDecoder( + vocab_size, + encoder_output_size, + attention_heads, + linear_units, + num_blocks, + dropout_rate, + positional_dropout_rate, + self_attention_dropout_rate, + src_attention_dropout_rate, + input_layer, + use_output_layer, + normalize_before, + key_bias=key_bias, + gradient_checkpointing=gradient_checkpointing, + tie_word_embedding=tie_word_embedding, + ) + + self.right_decoder = TransformerDecoder( + vocab_size, + encoder_output_size, + attention_heads, + linear_units, + r_num_blocks, + dropout_rate, + positional_dropout_rate, + self_attention_dropout_rate, + src_attention_dropout_rate, + input_layer, + use_output_layer, + normalize_before, + key_bias=key_bias, + gradient_checkpointing=gradient_checkpointing, + tie_word_embedding=tie_word_embedding, + ) + + def forward( + self, + memory: torch.Tensor, + memory_mask: torch.Tensor, + ys_in_pad: torch.Tensor, + ys_in_lens: torch.Tensor, + r_ys_in_pad: torch.Tensor, + reverse_weight: float = 0.0, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Forward decoder. + Args: + memory: encoded memory, float32 (batch, maxlen_in, feat) + memory_mask: encoder memory mask, (batch, 1, maxlen_in) + ys_in_pad: padded input token ids, int64 (batch, maxlen_out) + ys_in_lens: input lengths of this batch (batch) + r_ys_in_pad: padded input token ids, int64 (batch, maxlen_out), + used for right to left decoder + reverse_weight: used for right to left decoder + Returns: + (tuple): tuple containing: + x: decoded token score before softmax (batch, maxlen_out, + vocab_size) if use_output_layer is True, + r_x: x: decoded token score (right to left decoder) + before softmax (batch, maxlen_out, vocab_size) + if use_output_layer is True, + olens: (batch, ) + """ + l_x, _, olens = self.left_decoder(memory, memory_mask, ys_in_pad, ys_in_lens) + r_x = torch.tensor(0.0) + if reverse_weight > 0.0: + r_x, _, olens = self.right_decoder( + memory, memory_mask, r_ys_in_pad, ys_in_lens + ) + return l_x, r_x, olens + + def forward_one_step( + self, + memory: torch.Tensor, + memory_mask: torch.Tensor, + tgt: torch.Tensor, + tgt_mask: torch.Tensor, + cache: Optional[List[torch.Tensor]] = None, + ) -> Tuple[torch.Tensor, List[torch.Tensor]]: + """Forward one step. + This is only used for decoding. + Args: + memory: encoded memory, float32 (batch, maxlen_in, feat) + memory_mask: encoded memory mask, (batch, 1, maxlen_in) + tgt: input token ids, int64 (batch, maxlen_out) + tgt_mask: input token mask, (batch, maxlen_out) + dtype=torch.uint8 in PyTorch 1.2- + dtype=torch.bool in PyTorch 1.2+ (include 1.2) + cache: cached output list of (batch, max_time_out-1, size) + Returns: + y, cache: NN output value and cache per `self.decoders`. + y.shape` is (batch, maxlen_out, token) + """ + return self.left_decoder.forward_one_step( + memory, memory_mask, tgt, tgt_mask, cache + ) + + def tie_or_clone_weights(self, jit_mode: bool = True): + """Tie or clone module weights (between word_emb and output_layer) + depending of whether we are using TorchScript or not""" + self.left_decoder.tie_or_clone_weights(jit_mode) + self.right_decoder.tie_or_clone_weights(jit_mode) diff --git a/almeval/models/stepaudio/cosyvoice/transformer/decoder_layer.py b/almeval/models/stepaudio/cosyvoice/transformer/decoder_layer.py new file mode 100644 index 0000000000000000000000000000000000000000..6f1420732591b96ba2f09335e96dee99eae71003 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/transformer/decoder_layer.py @@ -0,0 +1,132 @@ +# Copyright (c) 2019 Shigeki Karita +# 2020 Mobvoi Inc (Binbin Zhang) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Decoder self-attention layer definition.""" +from typing import Optional, Tuple + +import torch +from torch import nn + + +class DecoderLayer(nn.Module): + """Single decoder layer module. + + Args: + size (int): Input dimension. + self_attn (torch.nn.Module): Self-attention module instance. + `MultiHeadedAttention` instance can be used as the argument. + src_attn (torch.nn.Module): Inter-attention module instance. + `MultiHeadedAttention` instance can be used as the argument. + If `None` is passed, Inter-attention is not used, such as + CIF, GPT, and other decoder only model. + feed_forward (torch.nn.Module): Feed-forward module instance. + `PositionwiseFeedForward` instance can be used as the argument. + dropout_rate (float): Dropout rate. + normalize_before (bool): + True: use layer_norm before each sub-block. + False: to use layer_norm after each sub-block. + """ + + def __init__( + self, + size: int, + self_attn: nn.Module, + src_attn: Optional[nn.Module], + feed_forward: nn.Module, + dropout_rate: float, + normalize_before: bool = True, + ): + """Construct an DecoderLayer object.""" + super().__init__() + self.size = size + self.self_attn = self_attn + self.src_attn = src_attn + self.feed_forward = feed_forward + self.norm1 = nn.LayerNorm(size, eps=1e-5) + self.norm2 = nn.LayerNorm(size, eps=1e-5) + self.norm3 = nn.LayerNorm(size, eps=1e-5) + self.dropout = nn.Dropout(dropout_rate) + self.normalize_before = normalize_before + + def forward( + self, + tgt: torch.Tensor, + tgt_mask: torch.Tensor, + memory: torch.Tensor, + memory_mask: torch.Tensor, + cache: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute decoded features. + + Args: + tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size). + tgt_mask (torch.Tensor): Mask for input tensor + (#batch, maxlen_out). + memory (torch.Tensor): Encoded memory + (#batch, maxlen_in, size). + memory_mask (torch.Tensor): Encoded memory mask + (#batch, maxlen_in). + cache (torch.Tensor): cached tensors. + (#batch, maxlen_out - 1, size). + + Returns: + torch.Tensor: Output tensor (#batch, maxlen_out, size). + torch.Tensor: Mask for output tensor (#batch, maxlen_out). + torch.Tensor: Encoded memory (#batch, maxlen_in, size). + torch.Tensor: Encoded memory mask (#batch, maxlen_in). + + """ + residual = tgt + if self.normalize_before: + tgt = self.norm1(tgt) + + if cache is None: + tgt_q = tgt + tgt_q_mask = tgt_mask + else: + # compute only the last frame query keeping dim: max_time_out -> 1 + assert cache.shape == ( + tgt.shape[0], + tgt.shape[1] - 1, + self.size, + ), "{cache.shape} == {(tgt.shape[0], tgt.shape[1] - 1, self.size)}" + tgt_q = tgt[:, -1:, :] + residual = residual[:, -1:, :] + tgt_q_mask = tgt_mask[:, -1:, :] + + x = residual + self.dropout(self.self_attn(tgt_q, tgt, tgt, tgt_q_mask)[0]) + if not self.normalize_before: + x = self.norm1(x) + + if self.src_attn is not None: + residual = x + if self.normalize_before: + x = self.norm2(x) + x = residual + self.dropout( + self.src_attn(x, memory, memory, memory_mask)[0] + ) + if not self.normalize_before: + x = self.norm2(x) + + residual = x + if self.normalize_before: + x = self.norm3(x) + x = residual + self.dropout(self.feed_forward(x)) + if not self.normalize_before: + x = self.norm3(x) + + if cache is not None: + x = torch.cat([cache, x], dim=1) + + return x, tgt_mask, memory, memory_mask diff --git a/almeval/models/stepaudio/cosyvoice/transformer/embedding.py b/almeval/models/stepaudio/cosyvoice/transformer/embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..7d4f6e1af263da2b9144648208e3f2a659220d16 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/transformer/embedding.py @@ -0,0 +1,293 @@ +# Copyright (c) 2020 Mobvoi Inc. (authors: Binbin Zhang, Di Wu) +# 2024 Alibaba Inc (Xiang Lyu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Modified from ESPnet(https://github.com/espnet/espnet) +"""Positonal Encoding Module.""" + +import math +from typing import Tuple, Union + +import torch +import torch.nn.functional as F +import numpy as np + + +class PositionalEncoding(torch.nn.Module): + """Positional encoding. + + :param int d_model: embedding dim + :param float dropout_rate: dropout rate + :param int max_len: maximum input length + + PE(pos, 2i) = sin(pos/(10000^(2i/dmodel))) + PE(pos, 2i+1) = cos(pos/(10000^(2i/dmodel))) + """ + + def __init__( + self, + d_model: int, + dropout_rate: float, + max_len: int = 5000, + reverse: bool = False, + ): + """Construct an PositionalEncoding object.""" + super().__init__() + self.d_model = d_model + self.xscale = math.sqrt(self.d_model) + self.dropout = torch.nn.Dropout(p=dropout_rate) + self.max_len = max_len + + self.pe = torch.zeros(self.max_len, self.d_model) + position = torch.arange(0, self.max_len, dtype=torch.float32).unsqueeze(1) + div_term = torch.exp( + torch.arange(0, self.d_model, 2, dtype=torch.float32) + * -(math.log(10000.0) / self.d_model) + ) + self.pe[:, 0::2] = torch.sin(position * div_term) + self.pe[:, 1::2] = torch.cos(position * div_term) + self.pe = self.pe.unsqueeze(0) + + def forward( + self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0 + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Add positional encoding. + + Args: + x (torch.Tensor): Input. Its shape is (batch, time, ...) + offset (int, torch.tensor): position offset + + Returns: + torch.Tensor: Encoded tensor. Its shape is (batch, time, ...) + torch.Tensor: for compatibility to RelPositionalEncoding + """ + + self.pe = self.pe.to(x.device) + pos_emb = self.position_encoding(offset, x.size(1), False) + x = x * self.xscale + pos_emb + return self.dropout(x), self.dropout(pos_emb) + + def position_encoding( + self, offset: Union[int, torch.Tensor], size: int, apply_dropout: bool = True + ) -> torch.Tensor: + """For getting encoding in a streaming fashion + + Attention!!!!! + we apply dropout only once at the whole utterance level in a none + streaming way, but will call this function several times with + increasing input size in a streaming scenario, so the dropout will + be applied several times. + + Args: + offset (int or torch.tensor): start offset + size (int): required size of position encoding + + Returns: + torch.Tensor: Corresponding encoding + """ + # How to subscript a Union type: + # https://github.com/pytorch/pytorch/issues/69434 + if isinstance(offset, int): + assert offset + size <= self.max_len + pos_emb = self.pe[:, offset : offset + size] + elif isinstance(offset, torch.Tensor) and offset.dim() == 0: # scalar + assert offset + size <= self.max_len + pos_emb = self.pe[:, offset : offset + size] + else: # for batched streaming decoding on GPU + assert torch.max(offset) + size <= self.max_len + index = offset.unsqueeze(1) + torch.arange(0, size).to( + offset.device + ) # B X T + flag = index > 0 + # remove negative offset + index = index * flag + pos_emb = F.embedding(index, self.pe[0]) # B X T X d_model + + if apply_dropout: + pos_emb = self.dropout(pos_emb) + return pos_emb + + +class RelPositionalEncoding(PositionalEncoding): + """Relative positional encoding module. + See : Appendix B in https://arxiv.org/abs/1901.02860 + Args: + d_model (int): Embedding dimension. + dropout_rate (float): Dropout rate. + max_len (int): Maximum input length. + """ + + def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000): + """Initialize class.""" + super().__init__(d_model, dropout_rate, max_len, reverse=True) + + def forward( + self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0 + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute positional encoding. + Args: + x (torch.Tensor): Input tensor (batch, time, `*`). + Returns: + torch.Tensor: Encoded tensor (batch, time, `*`). + torch.Tensor: Positional embedding tensor (1, time, `*`). + """ + self.pe = self.pe.to(x.device) + x = x * self.xscale + pos_emb = self.position_encoding(offset, x.size(1), False) + return self.dropout(x), self.dropout(pos_emb) + + +class WhisperPositionalEncoding(PositionalEncoding): + """Sinusoids position encoding used in openai-whisper.encoder""" + + def __init__(self, d_model: int, dropout_rate: float, max_len: int = 1500): + super().__init__(d_model, dropout_rate, max_len) + self.xscale = 1.0 + log_timescale_increment = np.log(10000) / (d_model // 2 - 1) + inv_timescales = torch.exp( + -log_timescale_increment * torch.arange(d_model // 2) + ) + scaled_time = ( + torch.arange(max_len)[:, np.newaxis] * inv_timescales[np.newaxis, :] + ) + pe = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1) + delattr(self, "pe") + self.register_buffer("pe", pe.unsqueeze(0)) + + +class LearnablePositionalEncoding(PositionalEncoding): + """Learnable position encoding used in openai-whisper.decoder""" + + def __init__(self, d_model: int, dropout_rate: float, max_len: int = 448): + super().__init__(d_model, dropout_rate, max_len) + # NOTE(xcsong): overwrite self.pe & self.xscale + self.pe = torch.nn.Parameter(torch.empty(1, max_len, d_model)) + self.xscale = 1.0 + + +class NoPositionalEncoding(torch.nn.Module): + """No position encoding""" + + def __init__(self, d_model: int, dropout_rate: float): + super().__init__() + self.d_model = d_model + self.dropout = torch.nn.Dropout(p=dropout_rate) + + def forward( + self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0 + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Just return zero vector for interface compatibility""" + pos_emb = torch.zeros(1, x.size(1), self.d_model).to(x.device) + return self.dropout(x), pos_emb + + def position_encoding( + self, offset: Union[int, torch.Tensor], size: int + ) -> torch.Tensor: + return torch.zeros(1, size, self.d_model) + + +class EspnetRelPositionalEncoding(torch.nn.Module): + """Relative positional encoding module (new implementation). + + Details can be found in https://github.com/espnet/espnet/pull/2816. + + See : Appendix B in https://arxiv.org/abs/1901.02860 + + Args: + d_model (int): Embedding dimension. + dropout_rate (float): Dropout rate. + max_len (int): Maximum input length. + + """ + + def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000): + """Construct an PositionalEncoding object.""" + super(EspnetRelPositionalEncoding, self).__init__() + self.d_model = d_model + self.xscale = math.sqrt(self.d_model) + self.dropout = torch.nn.Dropout(p=dropout_rate) + self.pe = None + self.extend_pe(torch.tensor(0.0).expand(1, max_len)) + + def extend_pe(self, x: torch.Tensor): + """Reset the positional encodings.""" + if self.pe is not None: + # self.pe contains both positive and negative parts + # the length of self.pe is 2 * input_len - 1 + if self.pe.size(1) >= x.size(1) * 2 - 1: + if self.pe.dtype != x.dtype or self.pe.device != x.device: + self.pe = self.pe.to(dtype=x.dtype, device=x.device) + return + # Suppose `i` means to the position of query vecotr and `j` means the + # position of key vector. We use position relative positions when keys + # are to the left (i>j) and negative relative positions otherwise (i Tuple[torch.Tensor, torch.Tensor]: + """Add positional encoding. + + Args: + x (torch.Tensor): Input tensor (batch, time, `*`). + + Returns: + torch.Tensor: Encoded tensor (batch, time, `*`). + + """ + self.extend_pe(x) + x = x * self.xscale + pos_emb = self.position_encoding(size=x.size(1), offset=offset) + return self.dropout(x), self.dropout(pos_emb) + + def position_encoding( + self, offset: Union[int, torch.Tensor], size: int + ) -> torch.Tensor: + """For getting encoding in a streaming fashion + + Attention!!!!! + we apply dropout only once at the whole utterance level in a none + streaming way, but will call this function several times with + increasing input size in a streaming scenario, so the dropout will + be applied several times. + + Args: + offset (int or torch.tensor): start offset + size (int): required size of position encoding + + Returns: + torch.Tensor: Corresponding encoding + """ + pos_emb = self.pe[ + :, + self.pe.size(1) // 2 - size + 1 : self.pe.size(1) // 2 + size, + ] + return pos_emb diff --git a/almeval/models/stepaudio/cosyvoice/transformer/encoder.py b/almeval/models/stepaudio/cosyvoice/transformer/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..76d00cff324bb4c7e0da1de9ad3d4e8417205fe6 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/transformer/encoder.py @@ -0,0 +1,633 @@ +# Copyright (c) 2021 Mobvoi Inc (Binbin Zhang, Di Wu) +# 2022 Xingchen Song (sxc19@mails.tsinghua.edu.cn) +# 2024 Alibaba Inc (Xiang Lyu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Modified from ESPnet(https://github.com/espnet/espnet) +"""Encoder definition.""" +from typing import Tuple +import time + +import torch +import torch.utils.checkpoint as ckpt +import torch.nn.functional as F + +from cosyvoice.transformer.convolution import ConvolutionModule +from cosyvoice.transformer.encoder_layer import ( + TransformerEncoderLayer, +) +from cosyvoice.transformer.encoder_layer import ( + ConformerEncoderLayer, +) +from cosyvoice.transformer.positionwise_feed_forward import ( + PositionwiseFeedForward, +) +from cosyvoice.utils.class_utils import ( + COSYVOICE_EMB_CLASSES, + COSYVOICE_SUBSAMPLE_CLASSES, + COSYVOICE_ATTENTION_CLASSES, + COSYVOICE_ACTIVATION_CLASSES, +) +from cosyvoice.utils.mask import make_pad_mask +from cosyvoice.utils.mask import add_optional_chunk_mask + + +class BaseEncoder(torch.nn.Module): + + def __init__( + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.0, + input_layer: str = "conv2d", + pos_enc_layer_type: str = "abs_pos", + normalize_before: bool = True, + static_chunk_size: int = 0, + use_dynamic_chunk: bool = False, + global_cmvn: torch.nn.Module = None, + use_dynamic_left_chunk: bool = False, + gradient_checkpointing: bool = False, + ): + """ + Args: + input_size (int): input dim + output_size (int): dimension of attention + attention_heads (int): the number of heads of multi head attention + linear_units (int): the hidden units number of position-wise feed + forward + num_blocks (int): the number of decoder blocks + dropout_rate (float): dropout rate + attention_dropout_rate (float): dropout rate in attention + positional_dropout_rate (float): dropout rate after adding + positional encoding + input_layer (str): input layer type. + optional [linear, conv2d, conv2d6, conv2d8] + pos_enc_layer_type (str): Encoder positional encoding layer type. + opitonal [abs_pos, scaled_abs_pos, rel_pos, no_pos] + normalize_before (bool): + True: use layer_norm before each sub-block of a layer. + False: use layer_norm after each sub-block of a layer. + static_chunk_size (int): chunk size for static chunk training and + decoding + use_dynamic_chunk (bool): whether use dynamic chunk size for + training or not, You can only use fixed chunk(chunk_size > 0) + or dyanmic chunk size(use_dynamic_chunk = True) + global_cmvn (Optional[torch.nn.Module]): Optional GlobalCMVN module + use_dynamic_left_chunk (bool): whether use dynamic left chunk in + dynamic chunk training + key_bias: whether use bias in attention.linear_k, False for whisper models. + gradient_checkpointing: rerunning a forward-pass segment for each + checkpointed segment during backward. + """ + super().__init__() + self._output_size = output_size + + self.global_cmvn = global_cmvn + self.embed = COSYVOICE_SUBSAMPLE_CLASSES[input_layer]( + input_size, + output_size, + dropout_rate, + COSYVOICE_EMB_CLASSES[pos_enc_layer_type]( + output_size, positional_dropout_rate + ), + ) + + self.normalize_before = normalize_before + self.after_norm = torch.nn.LayerNorm(output_size, eps=1e-5) + self.static_chunk_size = static_chunk_size + self.use_dynamic_chunk = use_dynamic_chunk + self.use_dynamic_left_chunk = use_dynamic_left_chunk + self.gradient_checkpointing = gradient_checkpointing + + def output_size(self) -> int: + return self._output_size + + def forward( + self, + xs: torch.Tensor, + xs_lens: torch.Tensor, + decoding_chunk_size: int = 0, + num_decoding_left_chunks: int = -1, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Embed positions in tensor. + + Args: + xs: padded input tensor (B, T, D) + xs_lens: input length (B) + decoding_chunk_size: decoding chunk size for dynamic chunk + 0: default for training, use random dynamic chunk. + <0: for decoding, use full chunk. + >0: for decoding, use fixed chunk size as set. + num_decoding_left_chunks: number of left chunks, this is for decoding, + the chunk size is decoding_chunk_size. + >=0: use num_decoding_left_chunks + <0: use all left chunks + Returns: + encoder output tensor xs, and subsampled masks + xs: padded output tensor (B, T' ~= T/subsample_rate, D) + masks: torch.Tensor batch padding mask after subsample + (B, 1, T' ~= T/subsample_rate) + NOTE(xcsong): + We pass the `__call__` method of the modules instead of `forward` to the + checkpointing API because `__call__` attaches all the hooks of the module. + https://discuss.pytorch.org/t/any-different-between-model-input-and-model-forward-input/3690/2 + """ + T = xs.size(1) + masks = ~make_pad_mask(xs_lens, T).unsqueeze(1) # (B, 1, T) + if self.global_cmvn is not None: + xs = self.global_cmvn(xs) + xs, pos_emb, masks = self.embed(xs, masks) + mask_pad = masks # (B, 1, T/subsample_rate) + chunk_masks = add_optional_chunk_mask( + xs, + masks, + self.use_dynamic_chunk, + self.use_dynamic_left_chunk, + decoding_chunk_size, + self.static_chunk_size, + num_decoding_left_chunks, + ) + print(f"chunk_masks shape: {chunk_masks.shape}") + if self.gradient_checkpointing and self.training: + xs = self.forward_layers_checkpointed(xs, chunk_masks, pos_emb, mask_pad) + else: + xs = self.forward_layers(xs, chunk_masks, pos_emb, mask_pad) + if self.normalize_before: + xs = self.after_norm(xs) + # Here we assume the mask is not changed in encoder layers, so just + # return the masks before encoder layers, and the masks will be used + # for cross attention with decoder later + return xs, masks + + def forward_layers( + self, + xs: torch.Tensor, + chunk_masks: torch.Tensor, + pos_emb: torch.Tensor, + mask_pad: torch.Tensor, + ) -> torch.Tensor: + for layer in self.encoders: + xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad) + return xs + + @torch.jit.unused + def forward_layers_checkpointed( + self, + xs: torch.Tensor, + chunk_masks: torch.Tensor, + pos_emb: torch.Tensor, + mask_pad: torch.Tensor, + ) -> torch.Tensor: + for layer in self.encoders: + xs, chunk_masks, _, _ = ckpt.checkpoint( + layer.__call__, xs, chunk_masks, pos_emb, mask_pad + ) + return xs + + @torch.jit.export + def forward_chunk( + self, + xs: torch.Tensor, + offset: int, + required_cache_size: int, + att_cache: torch.Tensor = torch.zeros(0, 0, 0, 0), + cnn_cache: torch.Tensor = torch.zeros(0, 0, 0, 0), + att_mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ Forward just one chunk + + Args: + xs (torch.Tensor): chunk input, with shape (b=1, time, mel-dim), + where `time == (chunk_size - 1) * subsample_rate + \ + subsample.right_context + 1` + offset (int): current offset in encoder output time stamp + required_cache_size (int): cache size required for next chunk + compuation + >=0: actual cache size + <0: means all history cache is required + att_cache (torch.Tensor): cache tensor for KEY & VALUE in + transformer/conformer attention, with shape + (elayers, head, cache_t1, d_k * 2), where + `head * d_k == hidden-dim` and + `cache_t1 == chunk_size * num_decoding_left_chunks`. + cnn_cache (torch.Tensor): cache tensor for cnn_module in conformer, + (elayers, b=1, hidden-dim, cache_t2), where + `cache_t2 == cnn.lorder - 1` + + Returns: + torch.Tensor: output of current input xs, + with shape (b=1, chunk_size, hidden-dim). + torch.Tensor: new attention cache required for next chunk, with + dynamic shape (elayers, head, ?, d_k * 2) + depending on required_cache_size. + torch.Tensor: new conformer cnn cache required for next chunk, with + same shape as the original cnn_cache. + + """ + assert xs.size(0) == 1 + # tmp_masks is just for interface compatibility + tmp_masks = torch.ones(1, xs.size(1), device=xs.device, dtype=torch.bool) + tmp_masks = tmp_masks.unsqueeze(1) + if self.global_cmvn is not None: + xs = self.global_cmvn(xs) + # NOTE(xcsong): Before embed, shape(xs) is (b=1, time, mel-dim) + xs, pos_emb, _ = self.embed(xs, tmp_masks, offset) + # NOTE(xcsong): After embed, shape(xs) is (b=1, chunk_size, hidden-dim) + elayers, cache_t1 = att_cache.size(0), att_cache.size(2) + chunk_size = xs.size(1) + attention_key_size = cache_t1 + chunk_size + pos_emb = self.embed.position_encoding( + offset=offset - cache_t1, size=attention_key_size + ) + if required_cache_size < 0: + next_cache_start = 0 + elif required_cache_size == 0: + next_cache_start = attention_key_size + else: + next_cache_start = max(attention_key_size - required_cache_size, 0) + r_att_cache = [] + r_cnn_cache = [] + for i, layer in enumerate(self.encoders): + # NOTE(xcsong): Before layer.forward + # shape(att_cache[i:i + 1]) is (1, head, cache_t1, d_k * 2), + # shape(cnn_cache[i]) is (b=1, hidden-dim, cache_t2) + xs, _, new_att_cache, new_cnn_cache = layer( + xs, + att_mask, + pos_emb, + att_cache=att_cache[i : i + 1] if elayers > 0 else att_cache, + cnn_cache=cnn_cache[i] if cnn_cache.size(0) > 0 else cnn_cache, + ) + # NOTE(xcsong): After layer.forward + # shape(new_att_cache) is (1, head, attention_key_size, d_k * 2), + # shape(new_cnn_cache) is (b=1, hidden-dim, cache_t2) + r_att_cache.append(new_att_cache[:, :, next_cache_start:, :]) + r_cnn_cache.append(new_cnn_cache.unsqueeze(0)) + if self.normalize_before: + xs = self.after_norm(xs) + + # NOTE(xcsong): shape(r_att_cache) is (elayers, head, ?, d_k * 2), + # ? may be larger than cache_t1, it depends on required_cache_size + r_att_cache = torch.cat(r_att_cache, dim=0) + # NOTE(xcsong): shape(r_cnn_cache) is (e, b=1, hidden-dim, cache_t2) + r_cnn_cache = torch.cat(r_cnn_cache, dim=0) + + return (xs, r_att_cache, r_cnn_cache) + + @torch.jit.unused + def forward_chunk_by_chunk( + self, + xs: torch.Tensor, + decoding_chunk_size: int, + num_decoding_left_chunks: int = -1, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward input chunk by chunk with chunk_size like a streaming + fashion + + Here we should pay special attention to computation cache in the + streaming style forward chunk by chunk. Three things should be taken + into account for computation in the current network: + 1. transformer/conformer encoder layers output cache + 2. convolution in conformer + 3. convolution in subsampling + + However, we don't implement subsampling cache for: + 1. We can control subsampling module to output the right result by + overlapping input instead of cache left context, even though it + wastes some computation, but subsampling only takes a very + small fraction of computation in the whole model. + 2. Typically, there are several covolution layers with subsampling + in subsampling module, it is tricky and complicated to do cache + with different convolution layers with different subsampling + rate. + 3. Currently, nn.Sequential is used to stack all the convolution + layers in subsampling, we need to rewrite it to make it work + with cache, which is not preferred. + Args: + xs (torch.Tensor): (1, max_len, dim) + chunk_size (int): decoding chunk size + """ + assert decoding_chunk_size > 0 + # The model is trained by static or dynamic chunk + assert self.static_chunk_size > 0 or self.use_dynamic_chunk + subsampling = self.embed.subsampling_rate + context = self.embed.right_context + 1 # Add current frame + stride = subsampling * decoding_chunk_size + decoding_window = (decoding_chunk_size - 1) * subsampling + context + num_frames = xs.size(1) + att_cache: torch.Tensor = torch.zeros((0, 0, 0, 0), device=xs.device) + cnn_cache: torch.Tensor = torch.zeros((0, 0, 0, 0), device=xs.device) + outputs = [] + offset = 0 + required_cache_size = decoding_chunk_size * num_decoding_left_chunks + + # Feed forward overlap input step by step + for cur in range(0, num_frames - context + 1, stride): + end = min(cur + decoding_window, num_frames) + chunk_xs = xs[:, cur:end, :] + (y, att_cache, cnn_cache) = self.forward_chunk( + chunk_xs, offset, required_cache_size, att_cache, cnn_cache + ) + outputs.append(y) + offset += y.size(1) + ys = torch.cat(outputs, 1) + masks = torch.ones((1, 1, ys.size(1)), device=ys.device, dtype=torch.bool) + return ys, masks + + +class TransformerEncoder(BaseEncoder): + """Transformer encoder module.""" + + def __init__( + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.0, + input_layer: str = "conv2d", + pos_enc_layer_type: str = "abs_pos", + normalize_before: bool = True, + static_chunk_size: int = 0, + use_dynamic_chunk: bool = False, + global_cmvn: torch.nn.Module = None, + use_dynamic_left_chunk: bool = False, + key_bias: bool = True, + selfattention_layer_type: str = "selfattn", + activation_type: str = "relu", + gradient_checkpointing: bool = False, + ): + """Construct TransformerEncoder + + See Encoder for the meaning of each parameter. + """ + super().__init__( + input_size, + output_size, + attention_heads, + linear_units, + num_blocks, + dropout_rate, + positional_dropout_rate, + attention_dropout_rate, + input_layer, + pos_enc_layer_type, + normalize_before, + static_chunk_size, + use_dynamic_chunk, + global_cmvn, + use_dynamic_left_chunk, + gradient_checkpointing, + ) + activation = COSYVOICE_ACTIVATION_CLASSES[activation_type]() + self.encoders = torch.nn.ModuleList( + [ + TransformerEncoderLayer( + output_size, + COSYVOICE_ATTENTION_CLASSES[selfattention_layer_type]( + attention_heads, output_size, attention_dropout_rate, key_bias + ), + PositionwiseFeedForward( + output_size, linear_units, dropout_rate, activation + ), + dropout_rate, + normalize_before, + ) + for _ in range(num_blocks) + ] + ) + + +class ConformerEncoder(BaseEncoder): + """Conformer encoder module.""" + + def __init__( + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.0, + input_layer: str = "conv2d", + pos_enc_layer_type: str = "rel_pos", + normalize_before: bool = True, + static_chunk_size: int = 0, + use_dynamic_chunk: bool = False, + global_cmvn: torch.nn.Module = None, + use_dynamic_left_chunk: bool = False, + positionwise_conv_kernel_size: int = 1, + macaron_style: bool = True, + selfattention_layer_type: str = "rel_selfattn", + activation_type: str = "swish", + use_cnn_module: bool = True, + cnn_module_kernel: int = 15, + causal: bool = False, + cnn_module_norm: str = "batch_norm", + key_bias: bool = True, + gradient_checkpointing: bool = False, + ): + """Construct ConformerEncoder + + Args: + input_size to use_dynamic_chunk, see in BaseEncoder + positionwise_conv_kernel_size (int): Kernel size of positionwise + conv1d layer. + macaron_style (bool): Whether to use macaron style for + positionwise layer. + selfattention_layer_type (str): Encoder attention layer type, + the parameter has no effect now, it's just for configure + compatibility. + activation_type (str): Encoder activation function type. + use_cnn_module (bool): Whether to use convolution module. + cnn_module_kernel (int): Kernel size of convolution module. + causal (bool): whether to use causal convolution or not. + key_bias: whether use bias in attention.linear_k, False for whisper models. + """ + super().__init__( + input_size, + output_size, + attention_heads, + linear_units, + num_blocks, + dropout_rate, + positional_dropout_rate, + attention_dropout_rate, + input_layer, + pos_enc_layer_type, + normalize_before, + static_chunk_size, + use_dynamic_chunk, + global_cmvn, + use_dynamic_left_chunk, + gradient_checkpointing, + ) + activation = COSYVOICE_ACTIVATION_CLASSES[activation_type]() + + # self-attention module definition + encoder_selfattn_layer_args = ( + attention_heads, + output_size, + attention_dropout_rate, + key_bias, + ) + # feed-forward module definition + positionwise_layer_args = ( + output_size, + linear_units, + dropout_rate, + activation, + ) + # convolution module definition + convolution_layer_args = ( + output_size, + cnn_module_kernel, + activation, + cnn_module_norm, + causal, + ) + + self.encoders = torch.nn.ModuleList( + [ + ConformerEncoderLayer( + output_size, + COSYVOICE_ATTENTION_CLASSES[selfattention_layer_type]( + *encoder_selfattn_layer_args + ), + PositionwiseFeedForward(*positionwise_layer_args), + ( + PositionwiseFeedForward(*positionwise_layer_args) + if macaron_style + else None + ), + ( + ConvolutionModule(*convolution_layer_args) + if use_cnn_module + else None + ), + dropout_rate, + normalize_before, + ) + for _ in range(num_blocks) + ] + ) + self.inference_buffers = {} + self.inference_graphs = {} + + @torch.inference_mode() + def capture_inference(self, seq_len_to_capture=[128, 256, 512, 1024]): + device = next(self.parameters()).device + start_time = time.time() + print( + f"Start capture_inference for ConformerEncoder, seq_len_to_capture: {seq_len_to_capture}" + ) + + for seq_len in seq_len_to_capture: + xs = torch.randn( + 1, seq_len, self._output_size, device=device, dtype=torch.bfloat16 + ) + xs_lens = torch.tensor([seq_len], device=device, dtype=torch.int32) + decoding_chunk_size = 0 + num_decoding_left_chunks = -1 + + T = xs.size(1) + masks = ~make_pad_mask(xs_lens, T).unsqueeze(1) # (B, 1, T) + if self.global_cmvn is not None: + xs = self.global_cmvn(xs) + xs, pos_emb, masks = self.embed(xs, masks) + mask_pad = masks # (B, 1, T/subsample_rate) + chunk_masks = add_optional_chunk_mask( + xs, + masks, + self.use_dynamic_chunk, + self.use_dynamic_left_chunk, + decoding_chunk_size, + self.static_chunk_size, + num_decoding_left_chunks, + ) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out = self.forward_layers(xs, chunk_masks, pos_emb, mask_pad) + + self.inference_graphs[seq_len] = g + self.inference_buffers[seq_len] = { + "xs": xs, + "chunk_masks": chunk_masks, + "pos_emb": pos_emb, + "mask_pad": mask_pad, + "out": out, + } + end_time = time.time() + print( + f"Finish capture_inference for ConformerEncoder, time elapsed: {end_time - start_time}" + ) + + @torch.inference_mode() + def inference(self, xs: torch.Tensor, xs_lens: torch.Tensor): + curr_seq_len = xs.shape[1] + target_len = None + + for seq_len in sorted(self.inference_graphs.keys()): + if seq_len >= curr_seq_len: + target_len = seq_len + break + + if target_len is not None: + xs = F.pad(xs, (0, 0, 0, target_len - curr_seq_len), "constant", 0) + + decoding_chunk_size = 0 + num_decoding_left_chunks = -1 + + T = xs.size(1) + masks = ~make_pad_mask(xs_lens, T).unsqueeze(1) # (B, 1, T) + if self.global_cmvn is not None: + xs = self.global_cmvn(xs) + xs, pos_emb, masks = self.embed(xs, masks) + mask_pad = masks # (B, 1, T/subsample_rate) + chunk_masks = add_optional_chunk_mask( + xs, + masks, + self.use_dynamic_chunk, + self.use_dynamic_left_chunk, + decoding_chunk_size, + self.static_chunk_size, + num_decoding_left_chunks, + ) + + if target_len is not None: + buffer = self.inference_buffers[target_len] + buffer["xs"].copy_(xs) + buffer["chunk_masks"].copy_(chunk_masks) + buffer["pos_emb"].copy_(pos_emb) + buffer["mask_pad"].copy_(mask_pad) + + self.inference_graphs[target_len].replay() + + out = buffer["out"][:, :curr_seq_len, :] + else: + out = self.forward_layers(xs, chunk_masks, pos_emb, mask_pad) + + if self.normalize_before: + out = self.after_norm(out) + return out, masks diff --git a/almeval/models/stepaudio/cosyvoice/transformer/encoder_layer.py b/almeval/models/stepaudio/cosyvoice/transformer/encoder_layer.py new file mode 100644 index 0000000000000000000000000000000000000000..59045afa6ff78e3973737e47b6734989d2b823b9 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/transformer/encoder_layer.py @@ -0,0 +1,237 @@ +# Copyright (c) 2021 Mobvoi Inc (Binbin Zhang, Di Wu) +# 2022 Xingchen Song (sxc19@mails.tsinghua.edu.cn) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Modified from ESPnet(https://github.com/espnet/espnet) +"""Encoder self-attention layer definition.""" + +from typing import Optional, Tuple + +import torch +from torch import nn + + +class TransformerEncoderLayer(nn.Module): + """Encoder layer module. + + Args: + size (int): Input dimension. + self_attn (torch.nn.Module): Self-attention module instance. + `MultiHeadedAttention` or `RelPositionMultiHeadedAttention` + instance can be used as the argument. + feed_forward (torch.nn.Module): Feed-forward module instance. + `PositionwiseFeedForward`, instance can be used as the argument. + dropout_rate (float): Dropout rate. + normalize_before (bool): + True: use layer_norm before each sub-block. + False: to use layer_norm after each sub-block. + """ + + def __init__( + self, + size: int, + self_attn: torch.nn.Module, + feed_forward: torch.nn.Module, + dropout_rate: float, + normalize_before: bool = True, + ): + """Construct an EncoderLayer object.""" + super().__init__() + self.self_attn = self_attn + self.feed_forward = feed_forward + self.norm1 = nn.LayerNorm(size, eps=1e-5) + self.norm2 = nn.LayerNorm(size, eps=1e-5) + self.dropout = nn.Dropout(dropout_rate) + self.size = size + self.normalize_before = normalize_before + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor, + pos_emb: torch.Tensor, + mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), + att_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), + cnn_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute encoded features. + + Args: + x (torch.Tensor): (#batch, time, size) + mask (torch.Tensor): Mask tensor for the input (#batch, time，time), + (0, 0, 0) means fake mask. + pos_emb (torch.Tensor): just for interface compatibility + to ConformerEncoderLayer + mask_pad (torch.Tensor): does not used in transformer layer, + just for unified api with conformer. + att_cache (torch.Tensor): Cache tensor of the KEY & VALUE + (#batch=1, head, cache_t1, d_k * 2), head * d_k == size. + cnn_cache (torch.Tensor): Convolution cache in conformer layer + (#batch=1, size, cache_t2), not used here, it's for interface + compatibility to ConformerEncoderLayer. + Returns: + torch.Tensor: Output tensor (#batch, time, size). + torch.Tensor: Mask tensor (#batch, time, time). + torch.Tensor: att_cache tensor, + (#batch=1, head, cache_t1 + time, d_k * 2). + torch.Tensor: cnn_cahce tensor (#batch=1, size, cache_t2). + + """ + residual = x + if self.normalize_before: + x = self.norm1(x) + x_att, new_att_cache = self.self_attn( + x, x, x, mask, pos_emb=pos_emb, cache=att_cache + ) + x = residual + self.dropout(x_att) + if not self.normalize_before: + x = self.norm1(x) + + residual = x + if self.normalize_before: + x = self.norm2(x) + x = residual + self.dropout(self.feed_forward(x)) + if not self.normalize_before: + x = self.norm2(x) + + fake_cnn_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device) + return x, mask, new_att_cache, fake_cnn_cache + + +class ConformerEncoderLayer(nn.Module): + """Encoder layer module. + Args: + size (int): Input dimension. + self_attn (torch.nn.Module): Self-attention module instance. + `MultiHeadedAttention` or `RelPositionMultiHeadedAttention` + instance can be used as the argument. + feed_forward (torch.nn.Module): Feed-forward module instance. + `PositionwiseFeedForward` instance can be used as the argument. + feed_forward_macaron (torch.nn.Module): Additional feed-forward module + instance. + `PositionwiseFeedForward` instance can be used as the argument. + conv_module (torch.nn.Module): Convolution module instance. + `ConvlutionModule` instance can be used as the argument. + dropout_rate (float): Dropout rate. + normalize_before (bool): + True: use layer_norm before each sub-block. + False: use layer_norm after each sub-block. + """ + + def __init__( + self, + size: int, + self_attn: torch.nn.Module, + feed_forward: Optional[nn.Module] = None, + feed_forward_macaron: Optional[nn.Module] = None, + conv_module: Optional[nn.Module] = None, + dropout_rate: float = 0.1, + normalize_before: bool = True, + ): + """Construct an EncoderLayer object.""" + super().__init__() + self.self_attn = self_attn + self.feed_forward = feed_forward + self.feed_forward_macaron = feed_forward_macaron + self.conv_module = conv_module + self.norm_ff = nn.LayerNorm(size, eps=1e-5) # for the FNN module + self.norm_mha = nn.LayerNorm(size, eps=1e-5) # for the MHA module + if feed_forward_macaron is not None: + self.norm_ff_macaron = nn.LayerNorm(size, eps=1e-5) + self.ff_scale = 0.5 + else: + self.ff_scale = 1.0 + if self.conv_module is not None: + self.norm_conv = nn.LayerNorm(size, eps=1e-5) # for the CNN module + self.norm_final = nn.LayerNorm( + size, eps=1e-5 + ) # for the final output of the block + self.dropout = nn.Dropout(dropout_rate) + self.size = size + self.normalize_before = normalize_before + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor, + pos_emb: torch.Tensor, + mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), + att_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), + cnn_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute encoded features. + + Args: + x (torch.Tensor): (#batch, time, size) + mask (torch.Tensor): Mask tensor for the input (#batch, time，time), + (0, 0, 0) means fake mask. + pos_emb (torch.Tensor): positional encoding, must not be None + for ConformerEncoderLayer. + mask_pad (torch.Tensor): batch padding mask used for conv module. + (#batch, 1，time), (0, 0, 0) means fake mask. + att_cache (torch.Tensor): Cache tensor of the KEY & VALUE + (#batch=1, head, cache_t1, d_k * 2), head * d_k == size. + cnn_cache (torch.Tensor): Convolution cache in conformer layer + (#batch=1, size, cache_t2) + Returns: + torch.Tensor: Output tensor (#batch, time, size). + torch.Tensor: Mask tensor (#batch, time, time). + torch.Tensor: att_cache tensor, + (#batch=1, head, cache_t1 + time, d_k * 2). + torch.Tensor: cnn_cahce tensor (#batch, size, cache_t2). + """ + + # whether to use macaron style + if self.feed_forward_macaron is not None: + residual = x + if self.normalize_before: + x = self.norm_ff_macaron(x) + x = residual + self.ff_scale * self.dropout(self.feed_forward_macaron(x)) + if not self.normalize_before: + x = self.norm_ff_macaron(x) + + # multi-headed self-attention module + residual = x + if self.normalize_before: + x = self.norm_mha(x) + x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb, att_cache) + x = residual + self.dropout(x_att) + if not self.normalize_before: + x = self.norm_mha(x) + + # convolution module + # Fake new cnn cache here, and then change it in conv_module + new_cnn_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device) + if self.conv_module is not None: + residual = x + if self.normalize_before: + x = self.norm_conv(x) + x, new_cnn_cache = self.conv_module(x, mask_pad, cnn_cache) + x = residual + self.dropout(x) + + if not self.normalize_before: + x = self.norm_conv(x) + + # feed forward module + residual = x + if self.normalize_before: + x = self.norm_ff(x) + + x = residual + self.ff_scale * self.dropout(self.feed_forward(x)) + if not self.normalize_before: + x = self.norm_ff(x) + + if self.conv_module is not None: + x = self.norm_final(x) + + return x, mask, new_att_cache, new_cnn_cache diff --git a/almeval/models/stepaudio/cosyvoice/transformer/label_smoothing_loss.py b/almeval/models/stepaudio/cosyvoice/transformer/label_smoothing_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..47574a25fd4369132bfc444cae1b0e6b85a89933 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/transformer/label_smoothing_loss.py @@ -0,0 +1,98 @@ +# Copyright (c) 2019 Shigeki Karita +# 2020 Mobvoi Inc (Binbin Zhang) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Label smoothing module.""" + +import torch +from torch import nn + + +class LabelSmoothingLoss(nn.Module): + """Label-smoothing loss. + + In a standard CE loss, the label's data distribution is: + [0,1,2] -> + [ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + ] + + In the smoothing version CE Loss,some probabilities + are taken from the true label prob (1.0) and are divided + among other labels. + + e.g. + smoothing=0.1 + [0,1,2] -> + [ + [0.9, 0.05, 0.05], + [0.05, 0.9, 0.05], + [0.05, 0.05, 0.9], + ] + + Args: + size (int): the number of class + padding_idx (int): padding class id which will be ignored for loss + smoothing (float): smoothing rate (0.0 means the conventional CE) + normalize_length (bool): + normalize loss by sequence length if True + normalize loss by batch size if False + """ + + def __init__( + self, + size: int, + padding_idx: int, + smoothing: float, + normalize_length: bool = False, + ): + """Construct an LabelSmoothingLoss object.""" + super(LabelSmoothingLoss, self).__init__() + self.criterion = nn.KLDivLoss(reduction="none") + self.padding_idx = padding_idx + self.confidence = 1.0 - smoothing + self.smoothing = smoothing + self.size = size + self.normalize_length = normalize_length + + def forward(self, x: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """Compute loss between x and target. + + The model outputs and data labels tensors are flatten to + (batch*seqlen, class) shape and a mask is applied to the + padding part which should not be calculated for loss. + + Args: + x (torch.Tensor): prediction (batch, seqlen, class) + target (torch.Tensor): + target signal masked with self.padding_id (batch, seqlen) + Returns: + loss (torch.Tensor) : The KL loss, scalar float value + """ + assert x.size(2) == self.size + batch_size = x.size(0) + x = x.view(-1, self.size) + target = target.view(-1) + # use zeros_like instead of torch.no_grad() for true_dist, + # since no_grad() can not be exported by JIT + true_dist = torch.zeros_like(x) + true_dist.fill_(self.smoothing / (self.size - 1)) + ignore = target == self.padding_idx # (B,) + total = len(target) - ignore.sum().item() + target = target.masked_fill(ignore, 0) # avoid -1 index + true_dist.scatter_(1, target.unsqueeze(1), self.confidence) + kl = self.criterion(torch.log_softmax(x, dim=1), true_dist) + denom = total if self.normalize_length else batch_size + return kl.masked_fill(ignore.unsqueeze(1), 0).sum() / denom diff --git a/almeval/models/stepaudio/cosyvoice/transformer/positionwise_feed_forward.py b/almeval/models/stepaudio/cosyvoice/transformer/positionwise_feed_forward.py new file mode 100644 index 0000000000000000000000000000000000000000..3e60fb55fd9149745821f505c98a6fb53cfc3d9f --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/transformer/positionwise_feed_forward.py @@ -0,0 +1,116 @@ +# Copyright (c) 2019 Shigeki Karita +# 2020 Mobvoi Inc (Binbin Zhang) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Positionwise feed forward layer definition.""" + +import torch + + +class PositionwiseFeedForward(torch.nn.Module): + """Positionwise feed forward layer. + + FeedForward are appied on each position of the sequence. + The output dim is same with the input dim. + + Args: + idim (int): Input dimenstion. + hidden_units (int): The number of hidden units. + dropout_rate (float): Dropout rate. + activation (torch.nn.Module): Activation function + """ + + def __init__( + self, + idim: int, + hidden_units: int, + dropout_rate: float, + activation: torch.nn.Module = torch.nn.ReLU(), + ): + """Construct a PositionwiseFeedForward object.""" + super(PositionwiseFeedForward, self).__init__() + self.w_1 = torch.nn.Linear(idim, hidden_units) + self.activation = activation + self.dropout = torch.nn.Dropout(dropout_rate) + self.w_2 = torch.nn.Linear(hidden_units, idim) + + def forward(self, xs: torch.Tensor) -> torch.Tensor: + """Forward function. + + Args: + xs: input tensor (B, L, D) + Returns: + output tensor, (B, L, D) + """ + return self.w_2(self.dropout(self.activation(self.w_1(xs)))) + + +class MoEFFNLayer(torch.nn.Module): + """ + Mixture of expert with Positionwise feed forward layer + See also figure 1 in https://arxiv.org/pdf/2305.15663.pdf + The output dim is same with the input dim. + + Modified from https://github.com/Lightning-AI/lit-gpt/pull/823 + https://github.com/mistralai/mistral-src/blob/b46d6/moe_one_file_ref.py#L203-L219 + Args: + n_expert: number of expert. + n_expert_per_token: The actual number of experts used for each frame + idim (int): Input dimenstion. + hidden_units (int): The number of hidden units. + dropout_rate (float): Dropout rate. + activation (torch.nn.Module): Activation function + """ + + def __init__( + self, + n_expert: int, + n_expert_per_token: int, + idim: int, + hidden_units: int, + dropout_rate: float, + activation: torch.nn.Module = torch.nn.ReLU(), + ): + super(MoEFFNLayer, self).__init__() + self.gate = torch.nn.Linear(idim, n_expert, bias=False) + self.experts = torch.nn.ModuleList( + PositionwiseFeedForward(idim, hidden_units, dropout_rate, activation) + for _ in range(n_expert) + ) + self.n_expert_per_token = n_expert_per_token + + def forward(self, xs: torch.Tensor) -> torch.Tensor: + """Foward function. + Args: + xs: input tensor (B, L, D) + Returns: + output tensor, (B, L, D) + + """ + B, L, D = xs.size() # batch size, sequence length, embedding dimension (idim) + xs = xs.view(-1, D) # (B*L, D) + router = self.gate(xs) # (B*L, n_expert) + logits, indices = torch.topk( + router, self.n_expert_per_token + ) # probs:(B*L, n_expert), indices: (B*L, n_expert) + weights = torch.nn.functional.softmax(logits, dim=1, dtype=torch.float).to( + dtype=xs.dtype + ) # (B*L, n_expert_per_token) + output = torch.zeros_like(xs) # (B*L, D) + for i, expert in enumerate(self.experts): + mask = indices == i + batch_idx, ith_expert = torch.where(mask) + output[batch_idx] += weights[batch_idx, ith_expert, None] * expert( + xs[batch_idx] + ) + return output.view(B, L, D) diff --git a/almeval/models/stepaudio/cosyvoice/transformer/subsampling.py b/almeval/models/stepaudio/cosyvoice/transformer/subsampling.py new file mode 100644 index 0000000000000000000000000000000000000000..01052c78c05e066c2305382b0e4aee8597254c00 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/transformer/subsampling.py @@ -0,0 +1,391 @@ +# Copyright (c) 2021 Mobvoi Inc (Binbin Zhang, Di Wu) +# 2024 Alibaba Inc (Xiang Lyu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Modified from ESPnet(https://github.com/espnet/espnet) +"""Subsampling layer definition.""" + +from typing import Tuple, Union + +import torch + + +class BaseSubsampling(torch.nn.Module): + + def __init__(self): + super().__init__() + self.right_context = 0 + self.subsampling_rate = 1 + + def position_encoding( + self, offset: Union[int, torch.Tensor], size: int + ) -> torch.Tensor: + return self.pos_enc.position_encoding(offset, size) + + +class EmbedinigNoSubsampling(BaseSubsampling): + """Embedding input without subsampling""" + + def __init__( + self, idim: int, odim: int, dropout_rate: float, pos_enc_class: torch.nn.Module + ): + super().__init__() + self.embed = torch.nn.Embedding(idim, odim) + self.pos_enc = pos_enc_class + + def forward( + self, + x: torch.Tensor, + x_mask: torch.Tensor, + offset: Union[int, torch.Tensor] = 0, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Input x. + + Args: + x (torch.Tensor): Input tensor (#batch, time, idim). + x_mask (torch.Tensor): Input mask (#batch, 1, time). + + Returns: + torch.Tensor: linear input tensor (#batch, time', odim), + where time' = time . + torch.Tensor: linear input mask (#batch, 1, time'), + where time' = time . + + """ + x = self.embed(x) + x, pos_emb = self.pos_enc(x, offset) + return x, pos_emb, x_mask + + +class LinearNoSubsampling(BaseSubsampling): + """Linear transform the input without subsampling + + Args: + idim (int): Input dimension. + odim (int): Output dimension. + dropout_rate (float): Dropout rate. + + """ + + def __init__( + self, idim: int, odim: int, dropout_rate: float, pos_enc_class: torch.nn.Module + ): + """Construct an linear object.""" + super().__init__() + self.out = torch.nn.Sequential( + torch.nn.Linear(idim, odim), + torch.nn.LayerNorm(odim, eps=1e-5), + torch.nn.Dropout(dropout_rate), + ) + self.pos_enc = pos_enc_class + self.right_context = 0 + self.subsampling_rate = 1 + + def forward( + self, + x: torch.Tensor, + x_mask: torch.Tensor, + offset: Union[int, torch.Tensor] = 0, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Input x. + + Args: + x (torch.Tensor): Input tensor (#batch, time, idim). + x_mask (torch.Tensor): Input mask (#batch, 1, time). + + Returns: + torch.Tensor: linear input tensor (#batch, time', odim), + where time' = time . + torch.Tensor: linear input mask (#batch, 1, time'), + where time' = time . + + """ + x = self.out(x) + x, pos_emb = self.pos_enc(x, offset) + return x, pos_emb, x_mask + + +class Conv1dSubsampling2(BaseSubsampling): + """Convolutional 1D subsampling (to 1/2 length). + It is designed for Whisper, ref: + https://github.com/openai/whisper/blob/main/whisper/model.py + + Args: + idim (int): Input dimension. + odim (int): Output dimension. + dropout_rate (float): Dropout rate. + + """ + + def __init__( + self, idim: int, odim: int, dropout_rate: float, pos_enc_class: torch.nn.Module + ): + """Construct an Conv1dSubsampling2 object.""" + super().__init__() + self.conv = torch.nn.Sequential( + torch.nn.Conv1d(idim, odim, kernel_size=3, padding=1), + torch.nn.GELU(), + torch.nn.Conv1d(odim, odim, kernel_size=3, stride=2, padding=1), + torch.nn.GELU(), + ) + self.pos_enc = pos_enc_class + # The right context for every conv layer is computed by: + # (kernel_size - 1) * frame_rate_of_this_layer + self.subsampling_rate = 2 + # 4 = (3 - 1) * 1 + (3 - 1) * 1 + self.right_context = 4 + + def forward( + self, + x: torch.Tensor, + x_mask: torch.Tensor, + offset: Union[int, torch.Tensor] = 0, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Subsample x. + + Args: + x (torch.Tensor): Input tensor (#batch, time, idim). + x_mask (torch.Tensor): Input mask (#batch, 1, time). + + Returns: + torch.Tensor: Subsampled tensor (#batch, time', odim), + where time' = time // 2. + torch.Tensor: Subsampled mask (#batch, 1, time'), + where time' = time // 2. + torch.Tensor: positional encoding + + """ + time = x.size(1) + x = x.transpose(1, 2) # (b, f, t) + x = self.conv(x) + x = x.transpose(1, 2) # (b, t, f) + x, pos_emb = self.pos_enc(x, offset) + return x, pos_emb, x_mask[:, :, (time + 1) % 2 :: 2] + + +class Conv2dSubsampling4(BaseSubsampling): + """Convolutional 2D subsampling (to 1/4 length). + + Args: + idim (int): Input dimension. + odim (int): Output dimension. + dropout_rate (float): Dropout rate. + + """ + + def __init__( + self, idim: int, odim: int, dropout_rate: float, pos_enc_class: torch.nn.Module + ): + """Construct an Conv2dSubsampling4 object.""" + super().__init__() + self.conv = torch.nn.Sequential( + torch.nn.Conv2d(1, odim, 3, 2), + torch.nn.ReLU(), + torch.nn.Conv2d(odim, odim, 3, 2), + torch.nn.ReLU(), + ) + self.out = torch.nn.Sequential( + torch.nn.Linear(odim * (((idim - 1) // 2 - 1) // 2), odim) + ) + self.pos_enc = pos_enc_class + # The right context for every conv layer is computed by: + # (kernel_size - 1) * frame_rate_of_this_layer + self.subsampling_rate = 4 + # 6 = (3 - 1) * 1 + (3 - 1) * 2 + self.right_context = 6 + + def forward( + self, + x: torch.Tensor, + x_mask: torch.Tensor, + offset: Union[int, torch.Tensor] = 0, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Subsample x. + + Args: + x (torch.Tensor): Input tensor (#batch, time, idim). + x_mask (torch.Tensor): Input mask (#batch, 1, time). + + Returns: + torch.Tensor: Subsampled tensor (#batch, time', odim), + where time' = time // 4. + torch.Tensor: Subsampled mask (#batch, 1, time'), + where time' = time // 4. + torch.Tensor: positional encoding + + """ + x = x.unsqueeze(1) # (b, c=1, t, f) + x = self.conv(x) + b, c, t, f = x.size() + x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f)) + x, pos_emb = self.pos_enc(x, offset) + return x, pos_emb, x_mask[:, :, 2::2][:, :, 2::2] + + +class Conv2dSubsampling6(BaseSubsampling): + """Convolutional 2D subsampling (to 1/6 length). + Args: + idim (int): Input dimension. + odim (int): Output dimension. + dropout_rate (float): Dropout rate. + pos_enc (torch.nn.Module): Custom position encoding layer. + """ + + def __init__( + self, idim: int, odim: int, dropout_rate: float, pos_enc_class: torch.nn.Module + ): + """Construct an Conv2dSubsampling6 object.""" + super().__init__() + self.conv = torch.nn.Sequential( + torch.nn.Conv2d(1, odim, 3, 2), + torch.nn.ReLU(), + torch.nn.Conv2d(odim, odim, 5, 3), + torch.nn.ReLU(), + ) + self.linear = torch.nn.Linear(odim * (((idim - 1) // 2 - 2) // 3), odim) + self.pos_enc = pos_enc_class + # 10 = (3 - 1) * 1 + (5 - 1) * 2 + self.subsampling_rate = 6 + self.right_context = 10 + + def forward( + self, + x: torch.Tensor, + x_mask: torch.Tensor, + offset: Union[int, torch.Tensor] = 0, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Subsample x. + Args: + x (torch.Tensor): Input tensor (#batch, time, idim). + x_mask (torch.Tensor): Input mask (#batch, 1, time). + + Returns: + torch.Tensor: Subsampled tensor (#batch, time', odim), + where time' = time // 6. + torch.Tensor: Subsampled mask (#batch, 1, time'), + where time' = time // 6. + torch.Tensor: positional encoding + """ + x = x.unsqueeze(1) # (b, c, t, f) + x = self.conv(x) + b, c, t, f = x.size() + x = self.linear(x.transpose(1, 2).contiguous().view(b, t, c * f)) + x, pos_emb = self.pos_enc(x, offset) + return x, pos_emb, x_mask[:, :, 2::2][:, :, 4::3] + + +class Conv2dSubsampling8(BaseSubsampling): + """Convolutional 2D subsampling (to 1/8 length). + + Args: + idim (int): Input dimension. + odim (int): Output dimension. + dropout_rate (float): Dropout rate. + + """ + + def __init__( + self, idim: int, odim: int, dropout_rate: float, pos_enc_class: torch.nn.Module + ): + """Construct an Conv2dSubsampling8 object.""" + super().__init__() + self.conv = torch.nn.Sequential( + torch.nn.Conv2d(1, odim, 3, 2), + torch.nn.ReLU(), + torch.nn.Conv2d(odim, odim, 3, 2), + torch.nn.ReLU(), + torch.nn.Conv2d(odim, odim, 3, 2), + torch.nn.ReLU(), + ) + self.linear = torch.nn.Linear( + odim * ((((idim - 1) // 2 - 1) // 2 - 1) // 2), odim + ) + self.pos_enc = pos_enc_class + self.subsampling_rate = 8 + # 14 = (3 - 1) * 1 + (3 - 1) * 2 + (3 - 1) * 4 + self.right_context = 14 + + def forward( + self, + x: torch.Tensor, + x_mask: torch.Tensor, + offset: Union[int, torch.Tensor] = 0, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Subsample x. + + Args: + x (torch.Tensor): Input tensor (#batch, time, idim). + x_mask (torch.Tensor): Input mask (#batch, 1, time). + + Returns: + torch.Tensor: Subsampled tensor (#batch, time', odim), + where time' = time // 8. + torch.Tensor: Subsampled mask (#batch, 1, time'), + where time' = time // 8. + torch.Tensor: positional encoding + """ + x = x.unsqueeze(1) # (b, c, t, f) + x = self.conv(x) + b, c, t, f = x.size() + x = self.linear(x.transpose(1, 2).contiguous().view(b, t, c * f)) + x, pos_emb = self.pos_enc(x, offset) + return x, pos_emb, x_mask[:, :, 2::2][:, :, 2::2][:, :, 2::2] + + +class LegacyLinearNoSubsampling(BaseSubsampling): + """Linear transform the input without subsampling + + Args: + idim (int): Input dimension. + odim (int): Output dimension. + dropout_rate (float): Dropout rate. + + """ + + def __init__( + self, idim: int, odim: int, dropout_rate: float, pos_enc_class: torch.nn.Module + ): + """Construct an linear object.""" + super().__init__() + self.out = torch.nn.Sequential( + torch.nn.Linear(idim, odim), + torch.nn.LayerNorm(odim, eps=1e-5), + torch.nn.Dropout(dropout_rate), + torch.nn.ReLU(), + ) + self.pos_enc = pos_enc_class + self.right_context = 0 + self.subsampling_rate = 1 + + def forward( + self, + x: torch.Tensor, + x_mask: torch.Tensor, + offset: Union[int, torch.Tensor] = 0, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Input x. + + Args: + x (torch.Tensor): Input tensor (#batch, time, idim). + x_mask (torch.Tensor): Input mask (#batch, 1, time). + + Returns: + torch.Tensor: linear input tensor (#batch, time', odim), + where time' = time . + torch.Tensor: linear input mask (#batch, 1, time'), + where time' = time . + + """ + x = self.out(x) + x, pos_emb = self.pos_enc(x, offset) + return x, pos_emb, x_mask diff --git a/almeval/models/stepaudio/cosyvoice/utils/__init__.py b/almeval/models/stepaudio/cosyvoice/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/almeval/models/stepaudio/cosyvoice/utils/audio.py b/almeval/models/stepaudio/cosyvoice/utils/audio.py new file mode 100644 index 0000000000000000000000000000000000000000..c837df56e74d52c7ef64a33a49851fb521352630 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/utils/audio.py @@ -0,0 +1,90 @@ +import numpy as np +import torch +import torch.utils.data +from librosa.filters import mel as librosa_mel_fn +from scipy.io.wavfile import read + +MAX_WAV_VALUE = 32768.0 + + +def load_wav(full_path): + sampling_rate, data = read(full_path) + return data, sampling_rate + + +def dynamic_range_compression(x, C=1, clip_val=1e-5): + return np.log(np.clip(x, a_min=clip_val, a_max=None) * C) + + +def dynamic_range_decompression(x, C=1): + return np.exp(x) / C + + +def dynamic_range_compression_torch(x, C=1, clip_val=1e-5): + return torch.log(torch.clamp(x, min=clip_val) * C) + + +def dynamic_range_decompression_torch(x, C=1): + return torch.exp(x) / C + + +def spectral_normalize_torch(magnitudes): + output = dynamic_range_compression_torch(magnitudes) + return output + + +def spectral_de_normalize_torch(magnitudes): + output = dynamic_range_decompression_torch(magnitudes) + return output + + +mel_basis = {} +hann_window = {} + + +def mel_spectrogram( + y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False +): + # if torch.min(y) < -1.0: + # print("min value is ", torch.min(y)) + # if torch.max(y) > 1.0: + # print("max value is ", torch.max(y)) + + global mel_basis, hann_window # pylint: disable=global-statement + if f"{str(fmax)}_{str(y.device)}" not in mel_basis: + mel = librosa_mel_fn( + sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax + ) + mel_basis[str(fmax) + "_" + str(y.device)] = ( + torch.from_numpy(mel).float().to(y.device) + ) + hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device) + + y = torch.nn.functional.pad( + y.unsqueeze(1), + (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), + mode="reflect", + ) + y = y.squeeze(1) + + spec = torch.view_as_real( + torch.stft( + y, + n_fft, + hop_length=hop_size, + win_length=win_size, + window=hann_window[str(y.device)], + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + ) + + spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)) + + spec = torch.matmul(mel_basis[str(fmax) + "_" + str(y.device)], spec) + spec = spectral_normalize_torch(spec) + + return spec diff --git a/almeval/models/stepaudio/cosyvoice/utils/class_utils.py b/almeval/models/stepaudio/cosyvoice/utils/class_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..52d7a959a9097fc21fb53f5ac43e4c1d15d953c3 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/utils/class_utils.py @@ -0,0 +1,78 @@ +# Copyright [2023-11-28] +# 2024 Alibaba Inc (authors: Xiang Lyu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import torch + +from cosyvoice.transformer.activation import Swish +from cosyvoice.transformer.subsampling import ( + LinearNoSubsampling, + EmbedinigNoSubsampling, + Conv1dSubsampling2, + Conv2dSubsampling4, + Conv2dSubsampling6, + Conv2dSubsampling8, +) +from cosyvoice.transformer.embedding import ( + PositionalEncoding, + RelPositionalEncoding, + WhisperPositionalEncoding, + LearnablePositionalEncoding, + NoPositionalEncoding, +) +from cosyvoice.transformer.attention import ( + MultiHeadedAttention, + RelPositionMultiHeadedAttention, +) +from cosyvoice.transformer.embedding import ( + EspnetRelPositionalEncoding, +) +from cosyvoice.transformer.subsampling import ( + LegacyLinearNoSubsampling, +) + + +COSYVOICE_ACTIVATION_CLASSES = { + "hardtanh": torch.nn.Hardtanh, + "tanh": torch.nn.Tanh, + "relu": torch.nn.ReLU, + "selu": torch.nn.SELU, + "swish": getattr(torch.nn, "SiLU", Swish), + "gelu": torch.nn.GELU, +} + +COSYVOICE_SUBSAMPLE_CLASSES = { + "linear": LinearNoSubsampling, + "linear_legacy": LegacyLinearNoSubsampling, + "embed": EmbedinigNoSubsampling, + "conv1d2": Conv1dSubsampling2, + "conv2d": Conv2dSubsampling4, + "conv2d6": Conv2dSubsampling6, + "conv2d8": Conv2dSubsampling8, + "paraformer_dummy": torch.nn.Identity, +} + +COSYVOICE_EMB_CLASSES = { + "embed": PositionalEncoding, + "abs_pos": PositionalEncoding, + "rel_pos": RelPositionalEncoding, + "rel_pos_espnet": EspnetRelPositionalEncoding, + "no_pos": NoPositionalEncoding, + "abs_pos_whisper": WhisperPositionalEncoding, + "embed_learnable_pe": LearnablePositionalEncoding, +} + +COSYVOICE_ATTENTION_CLASSES = { + "selfattn": MultiHeadedAttention, + "rel_selfattn": RelPositionMultiHeadedAttention, +} diff --git a/almeval/models/stepaudio/cosyvoice/utils/common.py b/almeval/models/stepaudio/cosyvoice/utils/common.py new file mode 100644 index 0000000000000000000000000000000000000000..e9611b6aa8ea5a12f50357d9b78af0e5e969ab85 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/utils/common.py @@ -0,0 +1,169 @@ +# Copyright (c) 2020 Mobvoi Inc (Binbin Zhang) +# 2024 Alibaba Inc (authors: Xiang Lyu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Modified from ESPnet(https://github.com/espnet/espnet) +"""Unility functions for Transformer.""" + +import random +from typing import List + +import numpy as np +import torch + +IGNORE_ID = -1 + + +def pad_list(xs: List[torch.Tensor], pad_value: int): + """Perform padding for the list of tensors. + + Args: + xs (List): List of Tensors [(T_1, `*`), (T_2, `*`), ..., (T_B, `*`)]. + pad_value (float): Value for padding. + + Returns: + Tensor: Padded tensor (B, Tmax, `*`). + + Examples: + >>> x = [torch.ones(4), torch.ones(2), torch.ones(1)] + >>> x + [tensor([1., 1., 1., 1.]), tensor([1., 1.]), tensor([1.])] + >>> pad_list(x, 0) + tensor([[1., 1., 1., 1.], + [1., 1., 0., 0.], + [1., 0., 0., 0.]]) + + """ + max_len = max([len(item) for item in xs]) + batchs = len(xs) + ndim = xs[0].ndim + if ndim == 1: + pad_res = torch.zeros(batchs, max_len, dtype=xs[0].dtype, device=xs[0].device) + elif ndim == 2: + pad_res = torch.zeros( + batchs, max_len, xs[0].shape[1], dtype=xs[0].dtype, device=xs[0].device + ) + elif ndim == 3: + pad_res = torch.zeros( + batchs, + max_len, + xs[0].shape[1], + xs[0].shape[2], + dtype=xs[0].dtype, + device=xs[0].device, + ) + else: + raise ValueError(f"Unsupported ndim: {ndim}") + pad_res.fill_(pad_value) + for i in range(batchs): + pad_res[i, : len(xs[i])] = xs[i] + return pad_res + + +def th_accuracy( + pad_outputs: torch.Tensor, pad_targets: torch.Tensor, ignore_label: int +) -> torch.Tensor: + """Calculate accuracy. + + Args: + pad_outputs (Tensor): Prediction tensors (B * Lmax, D). + pad_targets (LongTensor): Target label tensors (B, Lmax). + ignore_label (int): Ignore label id. + + Returns: + torch.Tensor: Accuracy value (0.0 - 1.0). + + """ + pad_pred = pad_outputs.view( + pad_targets.size(0), pad_targets.size(1), pad_outputs.size(1) + ).argmax(2) + mask = pad_targets != ignore_label + numerator = torch.sum( + pad_pred.masked_select(mask) == pad_targets.masked_select(mask) + ) + denominator = torch.sum(mask) + return (numerator / denominator).detach() + + +def get_padding(kernel_size, dilation=1): + return int((kernel_size * dilation - dilation) / 2) + + +def init_weights(m, mean=0.0, std=0.01): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + m.weight.data.normal_(mean, std) + + +# Repetition Aware Sampling in VALL-E 2 +def ras_sampling( + weighted_scores, + decoded_tokens, + sampling, + top_p=0.8, + top_k=25, + win_size=10, + tau_r=0.1, +): + top_ids = nucleus_sampling(weighted_scores, top_p=top_p, top_k=top_k) + rep_num = ( + (torch.tensor(decoded_tokens[-win_size:]).to(weighted_scores.device) == top_ids) + .sum() + .item() + ) + if rep_num >= win_size * tau_r: + top_ids = random_sampling(weighted_scores, decoded_tokens, sampling) + return top_ids + + +def nucleus_sampling(weighted_scores, top_p=0.8, top_k=25): + prob, indices = [], [] + cum_prob = 0.0 + sorted_value, sorted_idx = weighted_scores.softmax(dim=0).sort( + descending=True, stable=True + ) + for i in range(len(sorted_idx)): + # sampling both top-p and numbers. + if cum_prob < top_p and len(prob) < top_k: + cum_prob += sorted_value[i] + prob.append(sorted_value[i]) + indices.append(sorted_idx[i]) + else: + break + prob = torch.tensor(prob).to(weighted_scores) + indices = torch.tensor(indices, dtype=torch.long).to(weighted_scores.device) + top_ids = indices[prob.multinomial(1, replacement=True)] + return top_ids + + +def random_sampling(weighted_scores, decoded_tokens, sampling): + top_ids = weighted_scores.softmax(dim=0).multinomial(1, replacement=True) + return top_ids + + +def fade_in_out(fade_in_mel, fade_out_mel, window): + device = fade_in_mel.device + fade_in_mel, fade_out_mel = fade_in_mel.cpu(), fade_out_mel.cpu() + mel_overlap_len = int(window.shape[0] / 2) + fade_in_mel[..., :mel_overlap_len] = ( + fade_in_mel[..., :mel_overlap_len] * window[:mel_overlap_len] + + fade_out_mel[..., -mel_overlap_len:] * window[mel_overlap_len:] + ) + return fade_in_mel.to(device) + + +def set_all_random_seed(seed): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) diff --git a/almeval/models/stepaudio/cosyvoice/utils/executor.py b/almeval/models/stepaudio/cosyvoice/utils/executor.py new file mode 100644 index 0000000000000000000000000000000000000000..d79e9350f27804e972c5e642b58dc6a450cadb7c --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/utils/executor.py @@ -0,0 +1,151 @@ +# Copyright (c) 2020 Mobvoi Inc (Binbin Zhang) +# 2024 Alibaba Inc (authors: Xiang Lyu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +from contextlib import nullcontext +import os + +import torch +import torch.distributed as dist + +from cosyvoice.utils.train_utils import ( + update_parameter_and_lr, + log_per_step, + log_per_save, + batch_forward, + batch_backward, + save_model, + cosyvoice_join, +) + + +class Executor: + + def __init__(self): + self.step = 0 + self.epoch = 0 + self.rank = int(os.environ.get("RANK", 0)) + self.device = torch.device("cuda:{}".format(self.rank)) + + def train_one_epoc( + self, + model, + optimizer, + scheduler, + train_data_loader, + cv_data_loader, + writer, + info_dict, + group_join, + ): + """Train one epoch""" + + lr = optimizer.param_groups[0]["lr"] + logging.info( + "Epoch {} TRAIN info lr {} rank {}".format(self.epoch, lr, self.rank) + ) + logging.info( + "using accumulate grad, new batch size is {} times" + " larger than before".format(info_dict["accum_grad"]) + ) + # A context manager to be used in conjunction with an instance of + # torch.nn.parallel.DistributedDataParallel to be able to train + # with uneven inputs across participating processes. + model.train() + model_context = ( + model.join if info_dict["train_engine"] == "torch_ddp" else nullcontext + ) + with model_context(): + for batch_idx, batch_dict in enumerate(train_data_loader): + info_dict["tag"] = "TRAIN" + info_dict["step"] = self.step + info_dict["epoch"] = self.epoch + info_dict["batch_idx"] = batch_idx + if cosyvoice_join(group_join, info_dict): + break + + # Disable gradient synchronizations across DDP processes. + # Within this context, gradients will be accumulated on module + # variables, which will later be synchronized. + if ( + info_dict["train_engine"] == "torch_ddp" + and (batch_idx + 1) % info_dict["accum_grad"] != 0 + ): + context = model.no_sync + # Used for single gpu training and DDP gradient synchronization + # processes. + else: + context = nullcontext + + with context(): + info_dict = batch_forward(model, batch_dict, info_dict) + info_dict = batch_backward(model, info_dict) + + info_dict = update_parameter_and_lr( + model, optimizer, scheduler, info_dict + ) + log_per_step(writer, info_dict) + # NOTE specify save_per_step in cosyvoice.yaml if you want to enable step save + if ( + info_dict["save_per_step"] > 0 + and (self.step + 1) % info_dict["save_per_step"] == 0 + and (batch_idx + 1) % info_dict["accum_grad"] == 0 + ): + dist.barrier() + self.cv( + model, cv_data_loader, writer, info_dict, on_batch_end=False + ) + model.train() + if (batch_idx + 1) % info_dict["accum_grad"] == 0: + self.step += 1 + dist.barrier() + self.cv(model, cv_data_loader, writer, info_dict, on_batch_end=True) + + @torch.inference_mode() + def cv(self, model, cv_data_loader, writer, info_dict, on_batch_end=True): + """Cross validation on""" + logging.info( + "Epoch {} Step {} on_batch_end {} CV rank {}".format( + self.epoch, self.step + 1, on_batch_end, self.rank + ) + ) + model.eval() + total_num_utts, total_loss_dict = 0, {} # avoid division by 0 + for batch_idx, batch_dict in enumerate(cv_data_loader): + info_dict["tag"] = "CV" + info_dict["step"] = self.step + info_dict["epoch"] = self.epoch + info_dict["batch_idx"] = batch_idx + + num_utts = len(batch_dict["utts"]) + total_num_utts += num_utts + + info_dict = batch_forward(model, batch_dict, info_dict) + + for k, v in info_dict["loss_dict"].items(): + if k not in total_loss_dict: + total_loss_dict[k] = [] + total_loss_dict[k].append(v.item() * num_utts) + log_per_step(None, info_dict) + for k, v in total_loss_dict.items(): + total_loss_dict[k] = sum(v) / total_num_utts + info_dict["loss_dict"] = total_loss_dict + log_per_save(writer, info_dict) + model_name = ( + "epoch_{}_whole".format(self.epoch) + if on_batch_end + else "epoch_{}_step_{}".format(self.epoch, self.step + 1) + ) + save_model(model, model_name, info_dict) diff --git a/almeval/models/stepaudio/cosyvoice/utils/file_utils.py b/almeval/models/stepaudio/cosyvoice/utils/file_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..40cfa24528c99cae008e0d9b0e9515c4abbcb7db --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/utils/file_utils.py @@ -0,0 +1,49 @@ +# Copyright (c) 2021 Mobvoi Inc. (authors: Binbin Zhang) +# 2024 Alibaba Inc (authors: Xiang Lyu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import torchaudio +import logging + +logging.getLogger("matplotlib").setLevel(logging.WARNING) +logging.basicConfig(level=logging.DEBUG, format="%(asctime)s %(levelname)s %(message)s") + + +def read_lists(list_file): + lists = [] + with open(list_file, "r", encoding="utf8") as fin: + for line in fin: + lists.append(line.strip()) + return lists + + +def read_json_lists(list_file): + lists = read_lists(list_file) + results = {} + for fn in lists: + with open(fn, "r", encoding="utf8") as fin: + results.update(json.load(fin)) + return results + + +def load_wav(wav, target_sr): + speech, sample_rate = torchaudio.load(wav) + speech = speech.mean(dim=0, keepdim=True) + if sample_rate != target_sr: + # assert sample_rate > target_sr, 'wav sample rate {} must be greater than {}'.format(sample_rate, target_sr) + speech = torchaudio.transforms.Resample( + orig_freq=sample_rate, new_freq=target_sr + )(speech) + return speech diff --git a/almeval/models/stepaudio/cosyvoice/utils/frontend_utils.py b/almeval/models/stepaudio/cosyvoice/utils/frontend_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9e536f23a280fccd80b34494f2654cc035e1b74a --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/utils/frontend_utils.py @@ -0,0 +1,142 @@ +# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Zhihao Du) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re + +chinese_char_pattern = re.compile(r"[\u4e00-\u9fff]+") + + +# whether contain chinese character +def contains_chinese(text): + return bool(chinese_char_pattern.search(text)) + + +# replace special symbol +def replace_corner_mark(text): + text = text.replace("²", "平方") + text = text.replace("³", "立方") + return text + + +# remove meaningless symbol +def remove_bracket(text): + text = text.replace("（", "").replace("）", "") + text = text.replace("【", "").replace("】", "") + text = text.replace("`", "").replace("`", "") + text = text.replace("——", " ") + return text + + +# spell Arabic numerals +def spell_out_number(text: str, inflect_parser): + new_text = [] + st = None + for i, c in enumerate(text): + if not c.isdigit(): + if st is not None: + num_str = inflect_parser.number_to_words(text[st:i]) + new_text.append(num_str) + st = None + new_text.append(c) + else: + if st is None: + st = i + if st is not None and st < len(text): + num_str = inflect_parser.number_to_words(text[st:]) + new_text.append(num_str) + return "".join(new_text) + + +# split paragrah logic： +# 1. per sentence max len token_max_n, min len token_min_n, merge if last sentence len less than merge_len +# 2. cal sentence len according to lang +# 3. split sentence according to puncatation +def split_paragraph( + text: str, + tokenize, + lang="zh", + token_max_n=80, + token_min_n=60, + merge_len=20, + comma_split=False, +): + def calc_utt_length(_text: str): + if lang == "zh": + return len(_text) + else: + return len(tokenize(_text)) + + def should_merge(_text: str): + if lang == "zh": + return len(_text) < merge_len + else: + return len(tokenize(_text)) < merge_len + + if lang == "zh": + pounc = ["。", "？", "！", "；", "：", "、", ".", "?", "!", ";"] + else: + pounc = [".", "?", "!", ";", ":"] + if comma_split: + pounc.extend(["，", ","]) + + if text[-1] not in pounc: + if lang == "zh": + text += "。" + else: + text += "." + + st = 0 + utts = [] + for i, c in enumerate(text): + if c in pounc: + if len(text[st:i]) > 0: + utts.append(text[st:i] + c) + if i + 1 < len(text) and text[i + 1] in ['"', "”"]: + tmp = utts.pop(-1) + utts.append(tmp + text[i + 1]) + st = i + 2 + else: + st = i + 1 + + final_utts = [] + cur_utt = "" + for utt in utts: + if ( + calc_utt_length(cur_utt + utt) > token_max_n + and calc_utt_length(cur_utt) > token_min_n + ): + final_utts.append(cur_utt) + cur_utt = "" + cur_utt = cur_utt + utt + if len(cur_utt) > 0: + if should_merge(cur_utt) and len(final_utts) != 0: + final_utts[-1] = final_utts[-1] + cur_utt + else: + final_utts.append(cur_utt) + + return final_utts + + +# remove blank between chinese character +def replace_blank(text: str): + out_str = [] + for i, c in enumerate(text): + if c == " ": + if (text[i + 1].isascii() and text[i + 1] != " ") and ( + text[i - 1].isascii() and text[i - 1] != " " + ): + out_str.append(c) + else: + out_str.append(c) + return "".join(out_str) diff --git a/almeval/models/stepaudio/cosyvoice/utils/mask.py b/almeval/models/stepaudio/cosyvoice/utils/mask.py new file mode 100644 index 0000000000000000000000000000000000000000..be8bba4779acd4491088c690b8ac9665f92a5283 --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/utils/mask.py @@ -0,0 +1,226 @@ +# Copyright (c) 2019 Shigeki Karita +# 2020 Mobvoi Inc (Binbin Zhang) +# 2024 Alibaba Inc (authors: Xiang Lyu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +''' +def subsequent_mask( + size: int, + device: torch.device = torch.device("cpu"), +) -> torch.Tensor: + """Create mask for subsequent steps (size, size). + + This mask is used only in decoder which works in an auto-regressive mode. + This means the current step could only do attention with its left steps. + + In encoder, fully attention is used when streaming is not necessary and + the sequence is not long. In this case, no attention mask is needed. + + When streaming is need, chunk-based attention is used in encoder. See + subsequent_chunk_mask for the chunk-based attention mask. + + Args: + size (int): size of mask + str device (str): "cpu" or "cuda" or torch.Tensor.device + dtype (torch.device): result dtype + + Returns: + torch.Tensor: mask + + Examples: + >>> subsequent_mask(3) + [[1, 0, 0], + [1, 1, 0], + [1, 1, 1]] + """ + ret = torch.ones(size, size, device=device, dtype=torch.bool) + return torch.tril(ret) +''' + + +def subsequent_mask( + size: int, + device: torch.device = torch.device("cpu"), +) -> torch.Tensor: + """Create mask for subsequent steps (size, size). + + This mask is used only in decoder which works in an auto-regressive mode. + This means the current step could only do attention with its left steps. + + In encoder, fully attention is used when streaming is not necessary and + the sequence is not long. In this case, no attention mask is needed. + + When streaming is need, chunk-based attention is used in encoder. See + subsequent_chunk_mask for the chunk-based attention mask. + + Args: + size (int): size of mask + str device (str): "cpu" or "cuda" or torch.Tensor.device + dtype (torch.device): result dtype + + Returns: + torch.Tensor: mask + + Examples: + >>> subsequent_mask(3) + [[1, 0, 0], + [1, 1, 0], + [1, 1, 1]] + """ + arange = torch.arange(size, device=device) + mask = arange.expand(size, size) + arange = arange.unsqueeze(-1) + mask = mask <= arange + return mask + + +def subsequent_chunk_mask( + size: int, + chunk_size: int, + num_left_chunks: int = -1, + device: torch.device = torch.device("cpu"), +) -> torch.Tensor: + """Create mask for subsequent steps (size, size) with chunk size, + this is for streaming encoder + + Args: + size (int): size of mask + chunk_size (int): size of chunk + num_left_chunks (int): number of left chunks + <0: use full chunk + >=0: use num_left_chunks + device (torch.device): "cpu" or "cuda" or torch.Tensor.device + + Returns: + torch.Tensor: mask + + Examples: + >>> subsequent_chunk_mask(4, 2) + [[1, 1, 0, 0], + [1, 1, 0, 0], + [1, 1, 1, 1], + [1, 1, 1, 1]] + """ + ret = torch.zeros(size, size, device=device, dtype=torch.bool) + for i in range(size): + if num_left_chunks < 0: + start = 0 + else: + start = max((i // chunk_size - num_left_chunks) * chunk_size, 0) + ending = min((i // chunk_size + 1) * chunk_size, size) + ret[i, start:ending] = True + return ret + + +def add_optional_chunk_mask( + xs: torch.Tensor, + masks: torch.Tensor, + use_dynamic_chunk: bool, + use_dynamic_left_chunk: bool, + decoding_chunk_size: int, + static_chunk_size: int, + num_decoding_left_chunks: int, + enable_full_context: bool = True, +): + """Apply optional mask for encoder. + + Args: + xs (torch.Tensor): padded input, (B, L, D), L for max length + mask (torch.Tensor): mask for xs, (B, 1, L) + use_dynamic_chunk (bool): whether to use dynamic chunk or not + use_dynamic_left_chunk (bool): whether to use dynamic left chunk for + training. + decoding_chunk_size (int): decoding chunk size for dynamic chunk, it's + 0: default for training, use random dynamic chunk. + <0: for decoding, use full chunk. + >0: for decoding, use fixed chunk size as set. + static_chunk_size (int): chunk size for static chunk training/decoding + if it's greater than 0, if use_dynamic_chunk is true, + this parameter will be ignored + num_decoding_left_chunks: number of left chunks, this is for decoding, + the chunk size is decoding_chunk_size. + >=0: use num_decoding_left_chunks + <0: use all left chunks + enable_full_context (bool): + True: chunk size is either [1, 25] or full context(max_len) + False: chunk size ~ U[1, 25] + + Returns: + torch.Tensor: chunk mask of the input xs. + """ + # Whether to use chunk mask or not + if use_dynamic_chunk: + max_len = xs.size(1) + if decoding_chunk_size < 0: + chunk_size = max_len + num_left_chunks = -1 + elif decoding_chunk_size > 0: + chunk_size = decoding_chunk_size + num_left_chunks = num_decoding_left_chunks + else: + # chunk size is either [1, 25] or full context(max_len). + # Since we use 4 times subsampling and allow up to 1s(100 frames) + # delay, the maximum frame is 100 / 4 = 25. + chunk_size = torch.randint(1, max_len, (1,)).item() + num_left_chunks = -1 + if chunk_size > max_len // 2 and enable_full_context: + chunk_size = max_len + else: + chunk_size = chunk_size % 25 + 1 + if use_dynamic_left_chunk: + max_left_chunks = (max_len - 1) // chunk_size + num_left_chunks = torch.randint(0, max_left_chunks, (1,)).item() + chunk_masks = subsequent_chunk_mask( + xs.size(1), chunk_size, num_left_chunks, xs.device + ) # (L, L) + chunk_masks = chunk_masks.unsqueeze(0) # (1, L, L) + chunk_masks = masks & chunk_masks # (B, L, L) + elif static_chunk_size > 0: + num_left_chunks = num_decoding_left_chunks + chunk_masks = subsequent_chunk_mask( + xs.size(1), static_chunk_size, num_left_chunks, xs.device + ) # (L, L) + chunk_masks = chunk_masks.unsqueeze(0) # (1, L, L) + chunk_masks = masks & chunk_masks # (B, L, L) + else: + chunk_masks = masks + return chunk_masks + + +def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor: + """Make mask tensor containing indices of padded part. + + See description of make_non_pad_mask. + + Args: + lengths (torch.Tensor): Batch of lengths (B,). + Returns: + torch.Tensor: Mask tensor containing indices of padded part. + + Examples: + >>> lengths = [5, 3, 2] + >>> make_pad_mask(lengths) + masks = [[0, 0, 0, 0 ,0], + [0, 0, 0, 1, 1], + [0, 0, 1, 1, 1]] + """ + batch_size = lengths.size(0) + max_len = max_len if max_len > 0 else lengths.max().item() + seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device) + seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len) + seq_length_expand = lengths.unsqueeze(-1) + mask = seq_range_expand >= seq_length_expand + return mask diff --git a/almeval/models/stepaudio/cosyvoice/utils/scheduler.py b/almeval/models/stepaudio/cosyvoice/utils/scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..8235d212bcdfc59eacd0b111e14bb0d87831d3ae --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/utils/scheduler.py @@ -0,0 +1,761 @@ +# Copyright (c) 2020 Mobvoi Inc (Binbin Zhang) +# 2022 Ximalaya Inc (Yuguang Yang) +# 2024 Alibaba Inc (authors: Xiang Lyu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Modified from ESPnet(https://github.com/espnet/espnet) +# NeMo(https://github.com/NVIDIA/NeMo) + +from typing import Union + +import math +import warnings +import torch +from torch.optim.lr_scheduler import _LRScheduler + + +class WarmupLR(_LRScheduler): + """The WarmupLR scheduler + + This scheduler is almost same as NoamLR Scheduler except for following + difference: + + NoamLR: + lr = optimizer.lr * model_size ** -0.5 + * min(step ** -0.5, step * warmup_step ** -1.5) + WarmupLR: + lr = optimizer.lr * warmup_step ** 0.5 + * min(step ** -0.5, step * warmup_step ** -1.5) + + Note that the maximum lr equals to optimizer.lr in this scheduler. + + """ + + def __init__( + self, + optimizer: torch.optim.Optimizer, + warmup_steps: Union[int, float] = 25000, + last_epoch: int = -1, + ): + self.warmup_steps = warmup_steps + + # __init__() must be invoked before setting field + # because step() is also invoked in __init__() + super().__init__(optimizer, last_epoch) + + def __repr__(self): + return f"{self.__class__.__name__}(warmup_steps={self.warmup_steps})" + + def get_lr(self): + step_num = self.last_epoch + 1 + if self.warmup_steps == 0: + return [lr * step_num**-0.5 for lr in self.base_lrs] + else: + return [ + lr + * self.warmup_steps**0.5 + * min(step_num**-0.5, step_num * self.warmup_steps**-1.5) + for lr in self.base_lrs + ] + + def set_step(self, step: int): + self.last_epoch = step + + +class WarmupPolicy(_LRScheduler): + """Adds warmup kwargs and warmup logic to lr policy. + All arguments should be passed as kwargs for clarity, + Args: + warmup_steps: Number of training steps in warmup stage + warmup_ratio: Ratio of warmup steps to total steps + max_steps: Total number of steps while training or `None` for + infinite training + """ + + def __init__( + self, + optimizer, + *, + warmup_steps=None, + warmup_ratio=None, + max_steps=None, + min_lr=0.0, + last_epoch=-1, + ): + assert not ( + warmup_steps is not None and warmup_ratio is not None + ), "Either use particular number of step or ratio" + assert ( + warmup_ratio is None or max_steps is not None + ), "If there is a ratio, there should be a total steps" + + # It is necessary to assign all attributes *before* __init__, + # as class is wrapped by an inner class. + self.max_steps = max_steps + if warmup_steps is not None: + self.warmup_steps = warmup_steps + elif warmup_ratio is not None: + self.warmup_steps = int(warmup_ratio * max_steps) + else: + self.warmup_steps = 0 + + self.min_lr = min_lr + super().__init__(optimizer, last_epoch) + + def get_lr(self): + if not self._get_lr_called_within_step: + warnings.warn( + "To get the last learning rate computed " + "by the scheduler, please use `get_last_lr()`.", + UserWarning, + stacklevel=2, + ) + + step = self.last_epoch + + if step <= self.warmup_steps and self.warmup_steps > 0: + return self._get_warmup_lr(step) + + if step > self.max_steps: + return [self.min_lr for _ in self.base_lrs] + + return self._get_lr(step) + + def _get_warmup_lr(self, step): + lr_val = (step + 1) / (self.warmup_steps + 1) + return [initial_lr * lr_val for initial_lr in self.base_lrs] + + def _get_lr(self, step): + """Simple const lr policy""" + return self.base_lrs + + +class SquareRootConstantPolicy(_LRScheduler): + """Adds warmup kwargs and warmup logic to lr policy. + All arguments should be passed as kwargs for clarity, + Args: + warmup_steps: Number of training steps in warmup stage + warmup_ratio: Ratio of warmup steps to total steps + max_steps: Total number of steps while training or `None` for + infinite training + """ + + def __init__( + self, + optimizer, + *, + constant_steps=None, + constant_ratio=None, + max_steps=None, + min_lr=0.0, + last_epoch=-1, + ): + assert not ( + constant_steps is not None and constant_ratio is not None + ), "Either use particular number of step or ratio" + assert ( + constant_ratio is None or max_steps is not None + ), "If there is a ratio, there should be a total steps" + + # It is necessary to assign all attributes *before* __init__, + # as class is wrapped by an inner class. + self.max_steps = max_steps + if constant_steps is not None: + self.constant_steps = constant_steps + elif constant_ratio is not None: + self.constant_steps = int(constant_ratio * max_steps) + else: + self.constant_steps = 0 + + self.constant_lr = 1 / (constant_steps**0.5) + self.min_lr = min_lr + super().__init__(optimizer, last_epoch) + + def get_lr(self): + if not self._get_lr_called_within_step: + warnings.warn( + "To get the last learning rate computed " + "by the scheduler, please use `get_last_lr()`.", + UserWarning, + stacklevel=2, + ) + + step = self.last_epoch + + if step <= self.constant_steps: + return [self.constant_lr for _ in self.base_lrs] + + if step > self.max_steps: + return [self.min_lr for _ in self.base_lrs] + + return self._get_lr(step) + + def _get_lr(self, step): + """Simple const lr policy""" + return self.base_lrs + + +class WarmupHoldPolicy(WarmupPolicy): + """Variant of WarmupPolicy which maintains high + learning rate for a defined number of steps. + All arguments should be passed as kwargs for clarity, + Args: + warmup_steps: Number of training steps in warmup stage + warmup_ratio: Ratio of warmup steps to total steps + hold_steps: Number of training steps to + hold the learning rate after warm up + hold_ratio: Ratio of hold steps to total steps + max_steps: Total number of steps while training or `None` for + infinite training + """ + + def __init__( + self, + optimizer, + *, + warmup_steps=None, + warmup_ratio=None, + hold_steps=None, + hold_ratio=None, + max_steps=None, + min_lr=0.0, + last_epoch=-1, + ): + assert not ( + hold_steps is not None and hold_ratio is not None + ), "Either use particular number of step or ratio" + assert ( + hold_ratio is None or max_steps is not None + ), "If there is a ratio, there should be a total steps" + + self.min_lr = min_lr + self._last_warmup_lr = 0.0 + + # Necessary to duplicate as class attributes are hidden in inner class + self.max_steps = max_steps + if warmup_steps is not None: + self.warmup_steps = warmup_steps + elif warmup_ratio is not None: + self.warmup_steps = int(warmup_ratio * max_steps) + else: + self.warmup_steps = 0 + + if hold_steps is not None: + self.hold_steps = hold_steps + self.warmup_steps + elif hold_ratio is not None: + self.hold_steps = int(hold_ratio * max_steps) + self.warmup_steps + else: + self.hold_steps = 0 + + super().__init__( + optimizer, + warmup_steps=warmup_steps, + warmup_ratio=warmup_ratio, + max_steps=max_steps, + last_epoch=last_epoch, + min_lr=min_lr, + ) + + def get_lr(self): + if not self._get_lr_called_within_step: + warnings.warn( + "To get the last learning rate computed by the scheduler," + " " + "please use `get_last_lr()`.", + UserWarning, + stacklevel=2, + ) + + step = self.last_epoch + + # Warmup phase + if step <= self.warmup_steps and self.warmup_steps > 0: + return self._get_warmup_lr(step) + + # Hold phase + if (step >= self.warmup_steps) and (step < self.hold_steps): + return self.base_lrs + + if step > self.max_steps: + return [self.min_lr for _ in self.base_lrs] + + return self._get_lr(step) + + +class WarmupAnnealHoldPolicy(_LRScheduler): + """Adds warmup kwargs and warmup logic to lr policy. + All arguments should be passed as kwargs for clarity, + Args: + warmup_steps: Number of training steps in warmup stage + warmup_ratio: Ratio of warmup steps to total steps + max_steps: Total number of steps while training or `None` for + infinite training + min_lr: Minimum lr to hold the learning rate after decay at. + constant_steps: Number of steps to keep lr constant at. + constant_ratio: Ratio of steps to keep lr constant. + """ + + def __init__( + self, + optimizer, + *, + warmup_steps=None, + warmup_ratio=None, + constant_steps=None, + constant_ratio=None, + max_steps=None, + min_lr=0.0, + last_epoch=-1, + ): + assert not ( + warmup_steps is not None and warmup_ratio is not None + ), "Either use particular number of step or ratio" + assert not ( + constant_steps is not None and constant_ratio is not None + ), "Either use constant_steps or constant_ratio" + assert ( + warmup_ratio is None or max_steps is not None + ), "If there is a ratio, there should be a total steps" + + # It is necessary to assign all attributes *before* __init__, + # as class is wrapped by an inner class. + self.max_steps = max_steps + + if warmup_steps is not None: + self.warmup_steps = warmup_steps + elif warmup_ratio is not None: + self.warmup_steps = int(warmup_ratio * max_steps) + else: + self.warmup_steps = 0 + + if constant_steps is not None: + self.constant_steps = constant_steps + elif constant_ratio is not None: + self.constant_steps = int(constant_ratio * max_steps) + else: + self.constant_steps = 0 + + self.decay_steps = max_steps - (self.constant_steps + self.warmup_steps) + + self.min_lr = min_lr + super().__init__(optimizer, last_epoch) + + def get_lr(self): + if not self._get_lr_called_within_step: + warnings.warn( + "To get the last learning rate computed " + "by the scheduler, please use `get_last_lr()`.", + UserWarning, + stacklevel=2, + ) + + step = self.last_epoch + + # Warmup steps + if self.warmup_steps > 0 and step <= self.warmup_steps: + return self._get_warmup_lr(step) + + # Constant steps after warmup and decay + if ( + self.constant_steps > 0 + and (self.warmup_steps + self.decay_steps) < step <= self.max_steps + ): + return self._get_constant_lr(step) + + # Min lr after max steps of updates + if step > self.max_steps: + return [self.min_lr for _ in self.base_lrs] + + return self._get_lr(step) + + def _get_warmup_lr(self, step): + lr_val = (step + 1) / (self.warmup_steps + 1) + return [initial_lr * lr_val for initial_lr in self.base_lrs] + + def _get_constant_lr(self, step): + return [self.min_lr for _ in self.base_lrs] + + def _get_lr(self, step): + """Simple const lr policy""" + return self.base_lrs + + +def _squareroot_annealing(initial_lr, step, max_steps, min_lr): + mult = ((max_steps - step) / max_steps) ** 0.5 + out_lr = initial_lr * mult + out_lr = max(out_lr, min_lr) + return out_lr + + +def _square_annealing(initial_lr, step, max_steps, min_lr): + mult = ((max_steps - step) / max_steps) ** 2 + out_lr = initial_lr * mult + out_lr = max(out_lr, min_lr) + return out_lr + + +def _cosine_annealing(initial_lr, step, max_steps, min_lr): + mult = 0.5 * (1 + math.cos(math.pi * step / max_steps)) + out_lr = (initial_lr - min_lr) * mult + min_lr + return out_lr + + +def _linear_warmup_with_cosine_annealing( + max_lr, warmup_steps, step, decay_steps, min_lr +): + assert max_lr > min_lr + # Use linear warmup for the initial part. + if warmup_steps > 0 and step <= warmup_steps: + return max_lr * float(step) / float(warmup_steps) + + # For any steps larger than `decay_steps`, use `min_lr`. + if step > warmup_steps + decay_steps: + return min_lr + + # If we are done with the warmup period, use the decay style. + num_steps_ = step - warmup_steps + decay_steps_ = decay_steps + decay_ratio = float(num_steps_) / float(decay_steps_) + assert decay_ratio >= 0.0 + assert decay_ratio <= 1.0 + delta_lr = max_lr - min_lr + + coeff = 0.5 * (math.cos(math.pi * decay_ratio) + 1.0) + + return min_lr + coeff * delta_lr + + +def _poly_decay(initial_lr, step, decay_steps, power, min_lr, cycle): + if cycle: + multiplier = 1.0 if step == 0 else math.ceil(step / decay_steps) + decay_steps *= multiplier + else: + step = min(step, decay_steps) + p = step / decay_steps + lr = (initial_lr - min_lr) * math.pow(1.0 - p, power) + lr += min_lr + return lr + + +def _noam_hold_annealing( + initial_lr, step, warmup_steps, hold_steps, decay_rate, min_lr +): + # hold_steps = total number of steps + # to hold the LR, not the warmup + hold steps. + T_warmup_decay = max(1, warmup_steps**decay_rate) + T_hold_decay = max(1, (step - hold_steps) ** decay_rate) + lr = (initial_lr * T_warmup_decay) / T_hold_decay + lr = max(lr, min_lr) + return lr + + +class SquareAnnealing(WarmupPolicy): + + def __init__(self, optimizer, *, max_steps, min_lr=1e-5, last_epoch=-1, **kwargs): + super().__init__( + optimizer=optimizer, + max_steps=max_steps, + last_epoch=last_epoch, + min_lr=min_lr, + **kwargs, + ) + + def _get_lr(self, step): + new_lrs = [ + _square_annealing( + initial_lr=initial_lr, + step=step - self.warmup_steps, + max_steps=self.max_steps - self.warmup_steps, + min_lr=self.min_lr, + ) + for initial_lr in self.base_lrs + ] + return new_lrs + + +class SquareRootAnnealing(WarmupPolicy): + + def __init__(self, optimizer, *, max_steps, min_lr=0, last_epoch=-1, **kwargs): + super().__init__( + optimizer=optimizer, + max_steps=max_steps, + last_epoch=last_epoch, + min_lr=min_lr, + **kwargs, + ) + + def _get_lr(self, step): + new_lrs = [ + _squareroot_annealing( + initial_lr=initial_lr, + step=step, + max_steps=self.max_steps, + min_lr=self.min_lr, + ) + for initial_lr in self.base_lrs + ] + return new_lrs + + +class CosineAnnealing(WarmupAnnealHoldPolicy): + + def __init__(self, optimizer, *, max_steps, min_lr=0, last_epoch=-1, **kwargs): + super().__init__( + optimizer=optimizer, + max_steps=max_steps, + last_epoch=last_epoch, + min_lr=min_lr, + **kwargs, + ) + + def _get_lr(self, step): + for initial_lr in self.base_lrs: + if initial_lr < self.min_lr: + raise ValueError( + f"{self} received an initial learning rate " + f"that was lower than the minimum learning rate." + ) + + if self.constant_steps is None or self.constant_steps == 0: + new_lrs = [ + _cosine_annealing( + initial_lr=initial_lr, + step=step - self.warmup_steps, + max_steps=self.max_steps - self.warmup_steps, + min_lr=self.min_lr, + ) + for initial_lr in self.base_lrs + ] + else: + new_lrs = self._get_linear_warmup_with_cosine_annealing_lr(step) + return new_lrs + + def _get_warmup_lr(self, step): + if self.constant_steps is None or self.constant_steps == 0: + return super()._get_warmup_lr(step) + else: + # Use linear warmup for the initial part. + return self._get_linear_warmup_with_cosine_annealing_lr(step) + + def _get_constant_lr(self, step): + # Only called when `constant_steps` > 0. + return self._get_linear_warmup_with_cosine_annealing_lr(step) + + def _get_linear_warmup_with_cosine_annealing_lr(self, step): + # Cosine Schedule for Megatron LM, + # slightly different warmup schedule + constant LR at the end. + new_lrs = [ + _linear_warmup_with_cosine_annealing( + max_lr=self.base_lrs[0], + warmup_steps=self.warmup_steps, + step=step, + decay_steps=self.decay_steps, + min_lr=self.min_lr, + ) + for _ in self.base_lrs + ] + return new_lrs + + +class NoamAnnealing(_LRScheduler): + + def __init__( + self, + optimizer, + *, + d_model, + warmup_steps=None, + warmup_ratio=None, + max_steps=None, + min_lr=0.0, + last_epoch=-1, + ): + self._normalize = d_model ** (-0.5) + assert not ( + warmup_steps is not None and warmup_ratio is not None + ), "Either use particular number of step or ratio" + assert ( + warmup_ratio is None or max_steps is not None + ), "If there is a ratio, there should be a total steps" + + # It is necessary to assign all attributes *before* __init__, + # as class is wrapped by an inner class. + self.max_steps = max_steps + if warmup_steps is not None: + self.warmup_steps = warmup_steps + elif warmup_ratio is not None: + self.warmup_steps = int(warmup_ratio * max_steps) + else: + self.warmup_steps = 0 + + self.min_lr = min_lr + super().__init__(optimizer, last_epoch) + + def get_lr(self): + if not self._get_lr_called_within_step: + warnings.warn( + "To get the last learning rate computed " + "by the scheduler, please use `get_last_lr()`.", + UserWarning, + stacklevel=2, + ) + + step = max(1, self.last_epoch) + + for initial_lr in self.base_lrs: + if initial_lr < self.min_lr: + raise ValueError( + f"{self} received an initial learning rate " + f"that was lower than the minimum learning rate." + ) + + new_lrs = [ + self._noam_annealing(initial_lr=initial_lr, step=step) + for initial_lr in self.base_lrs + ] + return new_lrs + + def _noam_annealing(self, initial_lr, step): + if self.warmup_steps > 0: + mult = self._normalize * min( + step ** (-0.5), step * (self.warmup_steps ** (-1.5)) + ) + else: + mult = self._normalize * step ** (-0.5) + + out_lr = initial_lr * mult + if step > self.warmup_steps: + out_lr = max(out_lr, self.min_lr) + return out_lr + + +class NoamHoldAnnealing(WarmupHoldPolicy): + + def __init__( + self, + optimizer, + *, + max_steps, + decay_rate=0.5, + min_lr=0.0, + last_epoch=-1, + **kwargs, + ): + """ + From Nemo: + Implementation of the Noam Hold Annealing policy + from the SqueezeFormer paper. + + Unlike NoamAnnealing, the peak learning rate + can be explicitly set for this scheduler. + The schedule first performs linear warmup, + then holds the peak LR, then decays with some schedule for + the remainder of the steps. + Therefore the min-lr is still dependent + on the hyper parameters selected. + + It's schedule is determined by three factors- + + Warmup Steps: Initial stage, where linear warmup + occurs uptil the peak LR is reached. Unlike NoamAnnealing, + the peak LR is explicitly stated here instead of a scaling factor. + + Hold Steps: Intermediate stage, where the peak LR + is maintained for some number of steps. In this region, + the high peak LR allows the model to converge faster + if training is stable. However the high LR + may also cause instability during training. + Should usually be a significant fraction of training + steps (around 30-40% of the entire training steps). + + Decay Steps: Final stage, where the LR rapidly decays + with some scaling rate (set by decay rate). + To attain Noam decay, use 0.5, + for Squeezeformer recommended decay, use 1.0. + The fast decay after prolonged high LR during + hold phase allows for rapid convergence. + + References: + - [Squeezeformer: + An Efficient Transformer for Automatic Speech Recognition] + (https://arxiv.org/abs/2206.00888) + + Args: + optimizer: Pytorch compatible Optimizer object. + warmup_steps: Number of training steps in warmup stage + warmup_ratio: Ratio of warmup steps to total steps + hold_steps: Number of training steps to + hold the learning rate after warm up + hold_ratio: Ratio of hold steps to total steps + max_steps: Total number of steps while training or `None` for + infinite training + decay_rate: Float value describing the polynomial decay + after the hold period. Default value + of 0.5 corresponds to Noam decay. + min_lr: Minimum learning rate. + """ + self.decay_rate = decay_rate + super().__init__( + optimizer=optimizer, + max_steps=max_steps, + last_epoch=last_epoch, + min_lr=min_lr, + **kwargs, + ) + + def _get_lr(self, step): + if self.warmup_steps is None or self.warmup_steps == 0: + raise ValueError("Noam scheduler cannot be used without warmup steps") + + if self.hold_steps > 0: + hold_steps = self.hold_steps - self.warmup_steps + else: + hold_steps = 0 + + new_lrs = [ + _noam_hold_annealing( + initial_lr, + step=step, + warmup_steps=self.warmup_steps, + hold_steps=hold_steps, + decay_rate=self.decay_rate, + min_lr=self.min_lr, + ) + for initial_lr in self.base_lrs + ] + return new_lrs + + def set_step(self, step: int): + self.last_epoch = step + + +class ConstantLR(_LRScheduler): + """The ConstantLR scheduler + + This scheduler keeps a constant lr + + """ + + def __init__( + self, + optimizer: torch.optim.Optimizer, + ): + # __init__() must be invoked before setting field + # because step() is also invoked in __init__() + super().__init__(optimizer) + + def get_lr(self): + return self.base_lrs + + def set_step(self, step: int): + self.last_epoch = step diff --git a/almeval/models/stepaudio/cosyvoice/utils/train_utils.py b/almeval/models/stepaudio/cosyvoice/utils/train_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..fe9150c67ecc692266bcc31c6c5a35bf9bc84dac --- /dev/null +++ b/almeval/models/stepaudio/cosyvoice/utils/train_utils.py @@ -0,0 +1,350 @@ +# Copyright (c) 2021 Mobvoi Inc. (authors: Binbin Zhang) +# 2023 Horizon Inc. (authors: Xingchen Song) +# 2024 Alibaba Inc (authors: Xiang Lyu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from contextlib import nullcontext +import logging +import os +import torch +import json +import re +import datetime +import yaml + +import deepspeed +import torch.optim as optim +import torch.distributed as dist + +from torch.utils.tensorboard import SummaryWriter +from torch.utils.data import DataLoader +from torch.nn.utils import clip_grad_norm_ + +from deepspeed.runtime.zero.stage_1_and_2 import ( + estimate_zero2_model_states_mem_needs_all_live, +) + +from cosyvoice.dataset.dataset import Dataset +from cosyvoice.utils.scheduler import ( + WarmupLR, + NoamHoldAnnealing, + ConstantLR, +) + + +def init_distributed(args): + world_size = int(os.environ.get("WORLD_SIZE", 1)) + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + rank = int(os.environ.get("RANK", 0)) + logging.info( + "training on multiple gpus, this gpu {}".format(local_rank) + + ", rank {}, world_size {}".format(rank, world_size) + ) + if args.train_engine == "torch_ddp": + torch.cuda.set_device(local_rank) + dist.init_process_group(args.dist_backend) + else: + deepspeed.init_distributed(dist_backend=args.dist_backend) + return world_size, local_rank, rank + + +def init_dataset_and_dataloader(args, configs): + train_dataset = Dataset( + args.train_data, + data_pipeline=configs["data_pipeline"], + mode="train", + shuffle=True, + partition=True, + ) + cv_dataset = Dataset( + args.cv_data, + data_pipeline=configs["data_pipeline"], + mode="train", + shuffle=False, + partition=False, + ) + + # do not use persistent_workers=True, as whisper tokenizer opens tiktoken file each time when the for loop starts + train_data_loader = DataLoader( + train_dataset, + batch_size=None, + pin_memory=args.pin_memory, + num_workers=args.num_workers, + prefetch_factor=args.prefetch, + ) + cv_data_loader = DataLoader( + cv_dataset, + batch_size=None, + pin_memory=args.pin_memory, + num_workers=args.num_workers, + prefetch_factor=args.prefetch, + ) + return train_dataset, cv_dataset, train_data_loader, cv_data_loader + + +def check_modify_and_save_config(args, configs): + if args.train_engine == "torch_ddp": + configs["train_conf"]["dtype"] = "fp32" + else: + with open(args.deepspeed_config, "r") as fin: + ds_configs = json.load(fin) + if "fp16" in ds_configs and ds_configs["fp16"]["enabled"]: + configs["train_conf"]["dtype"] = "fp16" + elif "bf16" in ds_configs and ds_configs["bf16"]["enabled"]: + configs["train_conf"]["dtype"] = "bf16" + else: + configs["train_conf"]["dtype"] = "fp32" + assert ds_configs["train_micro_batch_size_per_gpu"] == 1 + # if use deepspeed, override ddp config + configs["train_conf"]["save_per_step"] = int( + configs["train_conf"]["save_per_step"] + * configs["train_conf"]["accum_grad"] + / ds_configs["gradient_accumulation_steps"] + ) + configs["train_conf"]["accum_grad"] = ds_configs["gradient_accumulation_steps"] + configs["train_conf"]["grad_clip"] = ds_configs["gradient_clipping"] + configs["train_conf"]["log_interval"] = ds_configs["steps_per_print"] + return configs + + +def wrap_cuda_model(args, model): + local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", 1)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + if args.train_engine == "torch_ddp": # native pytorch ddp + assert torch.cuda.is_available() + model.cuda() + model = torch.nn.parallel.DistributedDataParallel( + model, find_unused_parameters=True + ) + else: + if int(os.environ.get("RANK", 0)) == 0: + logging.info("Estimating model states memory needs (zero2)...") + estimate_zero2_model_states_mem_needs_all_live( + model, + num_gpus_per_node=local_world_size, + num_nodes=world_size // local_world_size, + ) + return model + + +def init_optimizer_and_scheduler(args, configs, model): + if configs["train_conf"]["optim"] == "adam": + optimizer = optim.Adam( + model.parameters(), **configs["train_conf"]["optim_conf"] + ) + elif configs["train_conf"]["optim"] == "adamw": + optimizer = optim.AdamW( + model.parameters(), **configs["train_conf"]["optim_conf"] + ) + else: + raise ValueError("unknown optimizer: " + configs["train_conf"]) + + if configs["train_conf"]["scheduler"] == "warmuplr": + scheduler_type = WarmupLR + scheduler = WarmupLR(optimizer, **configs["train_conf"]["scheduler_conf"]) + elif configs["train_conf"]["scheduler"] == "NoamHoldAnnealing": + scheduler_type = NoamHoldAnnealing + scheduler = NoamHoldAnnealing( + optimizer, **configs["train_conf"]["scheduler_conf"] + ) + elif configs["train_conf"]["scheduler"] == "constantlr": + scheduler_type = ConstantLR + scheduler = ConstantLR(optimizer) + else: + raise ValueError("unknown scheduler: " + configs["train_conf"]) + + # use deepspeed optimizer for speedup + if args.train_engine == "deepspeed": + + def scheduler(opt): + return scheduler_type(opt, **configs["train_conf"]["scheduler_conf"]) + + model, optimizer, _, scheduler = deepspeed.initialize( + args=args, + model=model, + optimizer=None, + lr_scheduler=scheduler, + model_parameters=model.parameters(), + ) + + return model, optimizer, scheduler + + +def init_summarywriter(args): + writer = None + if int(os.environ.get("RANK", 0)) == 0: + os.makedirs(args.model_dir, exist_ok=True) + writer = SummaryWriter(args.tensorboard_dir) + return writer + + +def save_model(model, model_name, info_dict): + rank = int(os.environ.get("RANK", 0)) + model_dir = info_dict["model_dir"] + save_model_path = os.path.join(model_dir, "{}.pt".format(model_name)) + + if info_dict["train_engine"] == "torch_ddp": + if rank == 0: + torch.save(model.module.state_dict(), save_model_path) + else: + with torch.no_grad(): + model.save_checkpoint( + save_dir=model_dir, tag=model_name, client_state=info_dict + ) + if rank == 0: + info_path = re.sub(".pt$", ".yaml", save_model_path) + info_dict["save_time"] = datetime.datetime.now().strftime("%d/%m/%Y %H:%M:%S") + with open(info_path, "w") as fout: + data = yaml.dump(info_dict) + fout.write(data) + logging.info( + "[Rank {}] Checkpoint: save to checkpoint {}".format(rank, save_model_path) + ) + + +def cosyvoice_join(group_join, info_dict): + world_size = int(os.environ.get("WORLD_SIZE", 1)) + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + rank = int(os.environ.get("RANK", 0)) + + if info_dict["batch_idx"] != 0: + # we try to join all rank in both ddp and deepspeed mode, in case different rank has different lr + try: + dist.monitored_barrier( + group=group_join, timeout=group_join.options._timeout + ) + return False + except RuntimeError as e: + logging.info( + "Detected uneven workload distribution: {}\n".format(e) + + "Break current worker to manually join all workers, " + + "world_size {}, current rank {}, current local_rank {}\n".format( + world_size, rank, local_rank + ) + ) + return True + else: + return False + + +def batch_forward(model, batch, info_dict): + device = int(os.environ.get("LOCAL_RANK", 0)) + + dtype = info_dict["dtype"] + if dtype == "fp16": + dtype = torch.float16 + elif dtype == "bf16": + dtype = torch.bfloat16 + else: # fp32 + dtype = torch.float32 + + if info_dict["train_engine"] == "torch_ddp": + autocast = nullcontext() + else: + autocast = torch.cuda.amp.autocast( + enabled=True, dtype=dtype, cache_enabled=False + ) + + with autocast: + info_dict["loss_dict"] = model(batch, device) + return info_dict + + +def batch_backward(model, info_dict): + if info_dict["train_engine"] == "deepspeed": + scaled_loss = model.backward(info_dict["loss_dict"]["loss"]) + else: + scaled_loss = info_dict["loss_dict"]["loss"] / info_dict["accum_grad"] + scaled_loss.backward() + + info_dict["loss_dict"]["loss"] = scaled_loss + return info_dict + + +def update_parameter_and_lr(model, optimizer, scheduler, info_dict): + grad_norm = 0.0 + if info_dict["train_engine"] == "deepspeed": + info_dict["is_gradient_accumulation_boundary"] = ( + model.is_gradient_accumulation_boundary() + ) + model.step() + grad_norm = model.get_global_grad_norm() + elif (info_dict["batch_idx"] + 1) % info_dict["accum_grad"] == 0: + grad_norm = clip_grad_norm_(model.parameters(), info_dict["grad_clip"]) + if torch.isfinite(grad_norm): + optimizer.step() + optimizer.zero_grad() + scheduler.step() + info_dict["lr"] = optimizer.param_groups[0]["lr"] + info_dict["grad_norm"] = grad_norm + return info_dict + + +def log_per_step(writer, info_dict): + tag = info_dict["tag"] + epoch = info_dict.get("epoch", 0) + step = info_dict["step"] + batch_idx = info_dict["batch_idx"] + loss_dict = info_dict["loss_dict"] + rank = int(os.environ.get("RANK", 0)) + + # only rank 0 write to tensorboard to avoid multi-process write + if writer is not None: + if ( + info_dict["train_engine"] == "deepspeed" + and info_dict["is_gradient_accumulation_boundary"] is True + ) or ( + info_dict["train_engine"] == "torch_ddp" + and (info_dict["batch_idx"] + 1) % info_dict["accum_grad"] == 0 + ): + for k in ["epoch", "lr", "grad_norm"]: + writer.add_scalar("{}/{}".format(tag, k), info_dict[k], step + 1) + for k, v in loss_dict.items(): + writer.add_scalar("{}/{}".format(tag, k), v, step + 1) + + # TRAIN & CV, Shell log (stdout) + if (info_dict["batch_idx"] + 1) % info_dict["log_interval"] == 0: + log_str = "{} Batch {}/{} ".format(tag, epoch, batch_idx + 1) + for name, value in loss_dict.items(): + log_str += "{} {:.6f} ".format(name, value) + if tag == "TRAIN": + log_str += "lr {:.8f} grad_norm {:.6f}".format( + info_dict["lr"], info_dict["grad_norm"] + ) + log_str += " rank {}".format(rank) + logging.debug(log_str) + + +def log_per_save(writer, info_dict): + tag = info_dict["tag"] + epoch = info_dict["epoch"] + step = info_dict["step"] + loss_dict = info_dict["loss_dict"] + lr = info_dict["lr"] + rank = int(os.environ.get("RANK", 0)) + logging.info( + "Epoch {} Step {} CV info lr {} {} rank {}".format( + epoch, + step + 1, + lr, + rank, + " ".join(["{}_{}".format(k, v) for k, v in loss_dict.items()]), + ) + ) + + if writer is not None: + for k in ["epoch", "lr"]: + writer.add_scalar("{}/{}".format(tag, k), info_dict[k], step + 1) + for k, v in loss_dict.items(): + writer.add_scalar("{}/{}".format(tag, k), v, step + 1) diff --git a/almeval/models/stepaudio/funasr_detach/frontends/utils/__init__.py b/almeval/models/stepaudio/funasr_detach/frontends/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b7f177368e62a5578b8706300e101f831a3972ac --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/frontends/utils/__init__.py @@ -0,0 +1 @@ +"""Initialize sub package.""" diff --git a/almeval/models/stepaudio/funasr_detach/frontends/utils/complex_utils.py b/almeval/models/stepaudio/funasr_detach/frontends/utils/complex_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3fcd494c7d7611b1e3ef837c78b0001e3d6c5554 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/frontends/utils/complex_utils.py @@ -0,0 +1,194 @@ +"""Beamformer module.""" + +from distutils.version import LooseVersion +from typing import Sequence +from typing import Tuple +from typing import Union + +import torch + +try: + from torch_complex import functional as FC + from torch_complex.tensor import ComplexTensor +except: + print("Please install torch_complex firstly") + + +EPS = torch.finfo(torch.double).eps +is_torch_1_8_plus = LooseVersion(torch.__version__) >= LooseVersion("1.8.0") +is_torch_1_9_plus = LooseVersion(torch.__version__) >= LooseVersion("1.9.0") + + +def new_complex_like( + ref: Union[torch.Tensor, ComplexTensor], + real_imag: Tuple[torch.Tensor, torch.Tensor], +): + if isinstance(ref, ComplexTensor): + return ComplexTensor(*real_imag) + elif is_torch_complex_tensor(ref): + return torch.complex(*real_imag) + else: + raise ValueError( + "Please update your PyTorch version to 1.9+ for complex support." + ) + + +def is_torch_complex_tensor(c): + return ( + not isinstance(c, ComplexTensor) and is_torch_1_9_plus and torch.is_complex(c) + ) + + +def is_complex(c): + return isinstance(c, ComplexTensor) or is_torch_complex_tensor(c) + + +def to_double(c): + if not isinstance(c, ComplexTensor) and is_torch_1_9_plus and torch.is_complex(c): + return c.to(dtype=torch.complex128) + else: + return c.double() + + +def to_float(c): + if not isinstance(c, ComplexTensor) and is_torch_1_9_plus and torch.is_complex(c): + return c.to(dtype=torch.complex64) + else: + return c.float() + + +def cat(seq: Sequence[Union[ComplexTensor, torch.Tensor]], *args, **kwargs): + if not isinstance(seq, (list, tuple)): + raise TypeError( + "cat(): argument 'tensors' (position 1) must be tuple of Tensors, " + "not Tensor" + ) + if isinstance(seq[0], ComplexTensor): + return FC.cat(seq, *args, **kwargs) + else: + return torch.cat(seq, *args, **kwargs) + + +def complex_norm( + c: Union[torch.Tensor, ComplexTensor], dim=-1, keepdim=False +) -> torch.Tensor: + if not is_complex(c): + raise TypeError("Input is not a complex tensor.") + if is_torch_complex_tensor(c): + return torch.norm(c, dim=dim, keepdim=keepdim) + else: + return torch.sqrt((c.real**2 + c.imag**2).sum(dim=dim, keepdim=keepdim) + EPS) + + +def einsum(equation, *operands): + # NOTE: Do not mix ComplexTensor and torch.complex in the input! + # NOTE (wangyou): Until PyTorch 1.9.0, torch.einsum does not support + # mixed input with complex and real tensors. + if len(operands) == 1: + if isinstance(operands[0], (tuple, list)): + operands = operands[0] + complex_module = FC if isinstance(operands[0], ComplexTensor) else torch + return complex_module.einsum(equation, *operands) + elif len(operands) != 2: + op0 = operands[0] + same_type = all(op.dtype == op0.dtype for op in operands[1:]) + if same_type: + _einsum = FC.einsum if isinstance(op0, ComplexTensor) else torch.einsum + return _einsum(equation, *operands) + else: + raise ValueError("0 or More than 2 operands are not supported.") + a, b = operands + if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor): + return FC.einsum(equation, a, b) + elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)): + if not torch.is_complex(a): + o_real = torch.einsum(equation, a, b.real) + o_imag = torch.einsum(equation, a, b.imag) + return torch.complex(o_real, o_imag) + elif not torch.is_complex(b): + o_real = torch.einsum(equation, a.real, b) + o_imag = torch.einsum(equation, a.imag, b) + return torch.complex(o_real, o_imag) + else: + return torch.einsum(equation, a, b) + else: + return torch.einsum(equation, a, b) + + +def inverse( + c: Union[torch.Tensor, ComplexTensor], +) -> Union[torch.Tensor, ComplexTensor]: + if isinstance(c, ComplexTensor): + return c.inverse2() + else: + return c.inverse() + + +def matmul( + a: Union[torch.Tensor, ComplexTensor], b: Union[torch.Tensor, ComplexTensor] +) -> Union[torch.Tensor, ComplexTensor]: + # NOTE: Do not mix ComplexTensor and torch.complex in the input! + # NOTE (wangyou): Until PyTorch 1.9.0, torch.matmul does not support + # multiplication between complex and real tensors. + if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor): + return FC.matmul(a, b) + elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)): + if not torch.is_complex(a): + o_real = torch.matmul(a, b.real) + o_imag = torch.matmul(a, b.imag) + return torch.complex(o_real, o_imag) + elif not torch.is_complex(b): + o_real = torch.matmul(a.real, b) + o_imag = torch.matmul(a.imag, b) + return torch.complex(o_real, o_imag) + else: + return torch.matmul(a, b) + else: + return torch.matmul(a, b) + + +def trace(a: Union[torch.Tensor, ComplexTensor]): + # NOTE (wangyou): until PyTorch 1.9.0, torch.trace does not + # support bacth processing. Use FC.trace() as fallback. + return FC.trace(a) + + +def reverse(a: Union[torch.Tensor, ComplexTensor], dim=0): + if isinstance(a, ComplexTensor): + return FC.reverse(a, dim=dim) + else: + return torch.flip(a, dims=(dim,)) + + +def solve(b: Union[torch.Tensor, ComplexTensor], a: Union[torch.Tensor, ComplexTensor]): + """Solve the linear equation ax = b.""" + # NOTE: Do not mix ComplexTensor and torch.complex in the input! + # NOTE (wangyou): Until PyTorch 1.9.0, torch.solve does not support + # mixed input with complex and real tensors. + if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor): + if isinstance(a, ComplexTensor) and isinstance(b, ComplexTensor): + return FC.solve(b, a, return_LU=False) + else: + return matmul(inverse(a), b) + elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)): + if torch.is_complex(a) and torch.is_complex(b): + return torch.linalg.solve(a, b) + else: + return matmul(inverse(a), b) + else: + if is_torch_1_8_plus: + return torch.linalg.solve(a, b) + else: + return torch.solve(b, a)[0] + + +def stack(seq: Sequence[Union[ComplexTensor, torch.Tensor]], *args, **kwargs): + if not isinstance(seq, (list, tuple)): + raise TypeError( + "stack(): argument 'tensors' (position 1) must be tuple of Tensors, " + "not Tensor" + ) + if isinstance(seq[0], ComplexTensor): + return FC.stack(seq, *args, **kwargs) + else: + return torch.stack(seq, *args, **kwargs) diff --git a/almeval/models/stepaudio/funasr_detach/frontends/utils/dnn_beamformer.py b/almeval/models/stepaudio/funasr_detach/frontends/utils/dnn_beamformer.py new file mode 100644 index 0000000000000000000000000000000000000000..135926971e1bc8e9d7d3c4950cc847d37738583f --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/frontends/utils/dnn_beamformer.py @@ -0,0 +1,173 @@ +"""DNN beamformer module.""" + +from typing import Tuple + +import torch +from torch.nn import functional as F + +from funasr_detach.frontends.utils.beamformer import apply_beamforming_vector +from funasr_detach.frontends.utils.beamformer import get_mvdr_vector +from funasr_detach.frontends.utils.beamformer import ( + get_power_spectral_density_matrix, # noqa: H301 +) +from funasr_detach.frontends.utils.mask_estimator import MaskEstimator +from torch_complex.tensor import ComplexTensor + + +class DNN_Beamformer(torch.nn.Module): + """DNN mask based Beamformer + + Citation: + Multichannel End-to-end Speech Recognition; T. Ochiai et al., 2017; + https://arxiv.org/abs/1703.04783 + + """ + + def __init__( + self, + bidim, + btype="blstmp", + blayers=3, + bunits=300, + bprojs=320, + bnmask=2, + dropout_rate=0.0, + badim=320, + ref_channel: int = -1, + beamformer_type="mvdr", + ): + super().__init__() + self.mask = MaskEstimator( + btype, bidim, blayers, bunits, bprojs, dropout_rate, nmask=bnmask + ) + self.ref = AttentionReference(bidim, badim) + self.ref_channel = ref_channel + + self.nmask = bnmask + + if beamformer_type != "mvdr": + raise ValueError( + "Not supporting beamformer_type={}".format(beamformer_type) + ) + self.beamformer_type = beamformer_type + + def forward( + self, data: ComplexTensor, ilens: torch.LongTensor + ) -> Tuple[ComplexTensor, torch.LongTensor, ComplexTensor]: + """The forward function + + Notation: + B: Batch + C: Channel + T: Time or Sequence length + F: Freq + + Args: + data (ComplexTensor): (B, T, C, F) + ilens (torch.Tensor): (B,) + Returns: + enhanced (ComplexTensor): (B, T, F) + ilens (torch.Tensor): (B,) + + """ + + def apply_beamforming(data, ilens, psd_speech, psd_noise): + # u: (B, C) + if self.ref_channel < 0: + u, _ = self.ref(psd_speech, ilens) + else: + # (optional) Create onehot vector for fixed reference microphone + u = torch.zeros( + *(data.size()[:-3] + (data.size(-2),)), device=data.device + ) + u[..., self.ref_channel].fill_(1) + + ws = get_mvdr_vector(psd_speech, psd_noise, u) + enhanced = apply_beamforming_vector(ws, data) + + return enhanced, ws + + # data (B, T, C, F) -> (B, F, C, T) + data = data.permute(0, 3, 2, 1) + + # mask: (B, F, C, T) + masks, _ = self.mask(data, ilens) + assert self.nmask == len(masks) + + if self.nmask == 2: # (mask_speech, mask_noise) + mask_speech, mask_noise = masks + + psd_speech = get_power_spectral_density_matrix(data, mask_speech) + psd_noise = get_power_spectral_density_matrix(data, mask_noise) + + enhanced, ws = apply_beamforming(data, ilens, psd_speech, psd_noise) + + # (..., F, T) -> (..., T, F) + enhanced = enhanced.transpose(-1, -2) + mask_speech = mask_speech.transpose(-1, -3) + else: # multi-speaker case: (mask_speech1, ..., mask_noise) + mask_speech = list(masks[:-1]) + mask_noise = masks[-1] + + psd_speeches = [ + get_power_spectral_density_matrix(data, mask) for mask in mask_speech + ] + psd_noise = get_power_spectral_density_matrix(data, mask_noise) + + enhanced = [] + ws = [] + for i in range(self.nmask - 1): + psd_speech = psd_speeches.pop(i) + # treat all other speakers' psd_speech as noises + enh, w = apply_beamforming( + data, ilens, psd_speech, sum(psd_speeches) + psd_noise + ) + psd_speeches.insert(i, psd_speech) + + # (..., F, T) -> (..., T, F) + enh = enh.transpose(-1, -2) + mask_speech[i] = mask_speech[i].transpose(-1, -3) + + enhanced.append(enh) + ws.append(w) + + return enhanced, ilens, mask_speech + + +class AttentionReference(torch.nn.Module): + def __init__(self, bidim, att_dim): + super().__init__() + self.mlp_psd = torch.nn.Linear(bidim, att_dim) + self.gvec = torch.nn.Linear(att_dim, 1) + + def forward( + self, psd_in: ComplexTensor, ilens: torch.LongTensor, scaling: float = 2.0 + ) -> Tuple[torch.Tensor, torch.LongTensor]: + """The forward function + + Args: + psd_in (ComplexTensor): (B, F, C, C) + ilens (torch.Tensor): (B,) + scaling (float): + Returns: + u (torch.Tensor): (B, C) + ilens (torch.Tensor): (B,) + """ + B, _, C = psd_in.size()[:3] + assert psd_in.size(2) == psd_in.size(3), psd_in.size() + # psd_in: (B, F, C, C) + psd = psd_in.masked_fill( + torch.eye(C, dtype=torch.bool, device=psd_in.device), 0 + ) + # psd: (B, F, C, C) -> (B, C, F) + psd = (psd.sum(dim=-1) / (C - 1)).transpose(-1, -2) + + # Calculate amplitude + psd_feat = (psd.real**2 + psd.imag**2) ** 0.5 + + # (B, C, F) -> (B, C, F2) + mlp_psd = self.mlp_psd(psd_feat) + # (B, C, F2) -> (B, C, 1) -> (B, C) + e = self.gvec(torch.tanh(mlp_psd)).squeeze(-1) + u = F.softmax(scaling * e, dim=-1) + return u, ilens diff --git a/almeval/models/stepaudio/funasr_detach/frontends/utils/dnn_wpe.py b/almeval/models/stepaudio/funasr_detach/frontends/utils/dnn_wpe.py new file mode 100644 index 0000000000000000000000000000000000000000..33d1ea93624af97b3255b4dffecf3e9ece969c41 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/frontends/utils/dnn_wpe.py @@ -0,0 +1,93 @@ +from typing import Tuple + +from pytorch_wpe import wpe_one_iteration +import torch +from torch_complex.tensor import ComplexTensor + +from funasr_detach.frontends.utils.mask_estimator import MaskEstimator +from funasr_detach.models.transformer.utils.nets_utils import make_pad_mask + + +class DNN_WPE(torch.nn.Module): + def __init__( + self, + wtype: str = "blstmp", + widim: int = 257, + wlayers: int = 3, + wunits: int = 300, + wprojs: int = 320, + dropout_rate: float = 0.0, + taps: int = 5, + delay: int = 3, + use_dnn_mask: bool = True, + iterations: int = 1, + normalization: bool = False, + ): + super().__init__() + self.iterations = iterations + self.taps = taps + self.delay = delay + + self.normalization = normalization + self.use_dnn_mask = use_dnn_mask + + self.inverse_power = True + + if self.use_dnn_mask: + self.mask_est = MaskEstimator( + wtype, widim, wlayers, wunits, wprojs, dropout_rate, nmask=1 + ) + + def forward( + self, data: ComplexTensor, ilens: torch.LongTensor + ) -> Tuple[ComplexTensor, torch.LongTensor, ComplexTensor]: + """The forward function + + Notation: + B: Batch + C: Channel + T: Time or Sequence length + F: Freq or Some dimension of the feature vector + + Args: + data: (B, C, T, F) + ilens: (B,) + Returns: + data: (B, C, T, F) + ilens: (B,) + """ + # (B, T, C, F) -> (B, F, C, T) + enhanced = data = data.permute(0, 3, 2, 1) + mask = None + + for i in range(self.iterations): + # Calculate power: (..., C, T) + power = enhanced.real**2 + enhanced.imag**2 + if i == 0 and self.use_dnn_mask: + # mask: (B, F, C, T) + (mask,), _ = self.mask_est(enhanced, ilens) + if self.normalization: + # Normalize along T + mask = mask / mask.sum(dim=-1)[..., None] + # (..., C, T) * (..., C, T) -> (..., C, T) + power = power * mask + + # Averaging along the channel axis: (..., C, T) -> (..., T) + power = power.mean(dim=-2) + + # enhanced: (..., C, T) -> (..., C, T) + enhanced = wpe_one_iteration( + data.contiguous(), + power, + taps=self.taps, + delay=self.delay, + inverse_power=self.inverse_power, + ) + + enhanced.masked_fill_(make_pad_mask(ilens, enhanced.real), 0) + + # (B, F, C, T) -> (B, T, C, F) + enhanced = enhanced.permute(0, 3, 2, 1) + if mask is not None: + mask = mask.transpose(-1, -3) + return enhanced, ilens, mask diff --git a/almeval/models/stepaudio/funasr_detach/frontends/utils/feature_transform.py b/almeval/models/stepaudio/funasr_detach/frontends/utils/feature_transform.py new file mode 100644 index 0000000000000000000000000000000000000000..f4e04042ab646b2e66a29bfb3827d67c43e9ae5a --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/frontends/utils/feature_transform.py @@ -0,0 +1,263 @@ +from typing import List +from typing import Tuple +from typing import Union + +import librosa +import numpy as np +import torch +from torch_complex.tensor import ComplexTensor + +from funasr_detach.models.transformer.utils.nets_utils import make_pad_mask + + +class FeatureTransform(torch.nn.Module): + def __init__( + self, + # Mel options, + fs: int = 16000, + n_fft: int = 512, + n_mels: int = 80, + fmin: float = 0.0, + fmax: float = None, + # Normalization + stats_file: str = None, + apply_uttmvn: bool = True, + uttmvn_norm_means: bool = True, + uttmvn_norm_vars: bool = False, + ): + super().__init__() + self.apply_uttmvn = apply_uttmvn + + self.logmel = LogMel(fs=fs, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax) + self.stats_file = stats_file + if stats_file is not None: + self.global_mvn = GlobalMVN(stats_file) + else: + self.global_mvn = None + + if self.apply_uttmvn is not None: + self.uttmvn = UtteranceMVN( + norm_means=uttmvn_norm_means, norm_vars=uttmvn_norm_vars + ) + else: + self.uttmvn = None + + def forward( + self, x: ComplexTensor, ilens: Union[torch.LongTensor, np.ndarray, List[int]] + ) -> Tuple[torch.Tensor, torch.LongTensor]: + # (B, T, F) or (B, T, C, F) + if x.dim() not in (3, 4): + raise ValueError(f"Input dim must be 3 or 4: {x.dim()}") + if not torch.is_tensor(ilens): + ilens = torch.from_numpy(np.asarray(ilens)).to(x.device) + + if x.dim() == 4: + # h: (B, T, C, F) -> h: (B, T, F) + if self.training: + # Select 1ch randomly + ch = np.random.randint(x.size(2)) + h = x[:, :, ch, :] + else: + # Use the first channel + h = x[:, :, 0, :] + else: + h = x + + # h: ComplexTensor(B, T, F) -> torch.Tensor(B, T, F) + h = h.real**2 + h.imag**2 + + h, _ = self.logmel(h, ilens) + if self.stats_file is not None: + h, _ = self.global_mvn(h, ilens) + if self.apply_uttmvn: + h, _ = self.uttmvn(h, ilens) + + return h, ilens + + +class LogMel(torch.nn.Module): + """Convert STFT to fbank feats + + The arguments is same as librosa.filters.mel + + Args: + fs: number > 0 [scalar] sampling rate of the incoming signal + n_fft: int > 0 [scalar] number of FFT components + n_mels: int > 0 [scalar] number of Mel bands to generate + fmin: float >= 0 [scalar] lowest frequency (in Hz) + fmax: float >= 0 [scalar] highest frequency (in Hz). + If `None`, use `fmax = fs / 2.0` + htk: use HTK formula instead of Slaney + norm: {None, 1, np.inf} [scalar] + if 1, divide the triangular mel weights by the width of the mel band + (area normalization). Otherwise, leave all the triangles aiming for + a peak value of 1.0 + + """ + + def __init__( + self, + fs: int = 16000, + n_fft: int = 512, + n_mels: int = 80, + fmin: float = 0.0, + fmax: float = None, + htk: bool = False, + norm=1, + ): + super().__init__() + + _mel_options = dict( + sr=fs, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax, htk=htk, norm=norm + ) + self.mel_options = _mel_options + + # Note(kamo): The mel matrix of librosa is different from kaldi. + melmat = librosa.filters.mel(**_mel_options) + # melmat: (D2, D1) -> (D1, D2) + self.register_buffer("melmat", torch.from_numpy(melmat.T).float()) + + def extra_repr(self): + return ", ".join(f"{k}={v}" for k, v in self.mel_options.items()) + + def forward( + self, feat: torch.Tensor, ilens: torch.LongTensor + ) -> Tuple[torch.Tensor, torch.LongTensor]: + # feat: (B, T, D1) x melmat: (D1, D2) -> mel_feat: (B, T, D2) + mel_feat = torch.matmul(feat, self.melmat) + + logmel_feat = (mel_feat + 1e-20).log() + # Zero padding + logmel_feat = logmel_feat.masked_fill(make_pad_mask(ilens, logmel_feat, 1), 0.0) + return logmel_feat, ilens + + +class GlobalMVN(torch.nn.Module): + """Apply global mean and variance normalization + + Args: + stats_file(str): npy file of 1-dim array or text file. + From the _first element to + the {(len(array) - 1) / 2}th element are treated as + the sum of features, + and the rest excluding the last elements are + treated as the sum of the square value of features, + and the last elements eqauls to the number of samples. + std_floor(float): + """ + + def __init__( + self, + stats_file: str, + norm_means: bool = True, + norm_vars: bool = True, + eps: float = 1.0e-20, + ): + super().__init__() + self.norm_means = norm_means + self.norm_vars = norm_vars + + self.stats_file = stats_file + stats = np.load(stats_file) + + stats = stats.astype(float) + assert (len(stats) - 1) % 2 == 0, stats.shape + + count = stats.flatten()[-1] + mean = stats[: (len(stats) - 1) // 2] / count + var = stats[(len(stats) - 1) // 2 : -1] / count - mean * mean + std = np.maximum(np.sqrt(var), eps) + + self.register_buffer("bias", torch.from_numpy(-mean.astype(np.float32))) + self.register_buffer("scale", torch.from_numpy(1 / std.astype(np.float32))) + + def extra_repr(self): + return ( + f"stats_file={self.stats_file}, " + f"norm_means={self.norm_means}, norm_vars={self.norm_vars}" + ) + + def forward( + self, x: torch.Tensor, ilens: torch.LongTensor + ) -> Tuple[torch.Tensor, torch.LongTensor]: + # feat: (B, T, D) + if self.norm_means: + x += self.bias.type_as(x) + x.masked_fill(make_pad_mask(ilens, x, 1), 0.0) + + if self.norm_vars: + x *= self.scale.type_as(x) + return x, ilens + + +class UtteranceMVN(torch.nn.Module): + def __init__( + self, norm_means: bool = True, norm_vars: bool = False, eps: float = 1.0e-20 + ): + super().__init__() + self.norm_means = norm_means + self.norm_vars = norm_vars + self.eps = eps + + def extra_repr(self): + return f"norm_means={self.norm_means}, norm_vars={self.norm_vars}" + + def forward( + self, x: torch.Tensor, ilens: torch.LongTensor + ) -> Tuple[torch.Tensor, torch.LongTensor]: + return utterance_mvn( + x, ilens, norm_means=self.norm_means, norm_vars=self.norm_vars, eps=self.eps + ) + + +def utterance_mvn( + x: torch.Tensor, + ilens: torch.LongTensor, + norm_means: bool = True, + norm_vars: bool = False, + eps: float = 1.0e-20, +) -> Tuple[torch.Tensor, torch.LongTensor]: + """Apply utterance mean and variance normalization + + Args: + x: (B, T, D), assumed zero padded + ilens: (B, T, D) + norm_means: + norm_vars: + eps: + + """ + ilens_ = ilens.type_as(x) + # mean: (B, D) + mean = x.sum(dim=1) / ilens_[:, None] + + if norm_means: + x -= mean[:, None, :] + x_ = x + else: + x_ = x - mean[:, None, :] + + # Zero padding + x_.masked_fill(make_pad_mask(ilens, x_, 1), 0.0) + if norm_vars: + var = x_.pow(2).sum(dim=1) / ilens_[:, None] + var = torch.clamp(var, min=eps) + x /= var.sqrt()[:, None, :] + x_ = x + return x_, ilens + + +def feature_transform_for(args, n_fft): + return FeatureTransform( + # Mel options, + fs=args.fbank_fs, + n_fft=n_fft, + n_mels=args.n_mels, + fmin=args.fbank_fmin, + fmax=args.fbank_fmax, + # Normalization + stats_file=args.stats_file, + apply_uttmvn=args.apply_uttmvn, + uttmvn_norm_means=args.uttmvn_norm_means, + uttmvn_norm_vars=args.uttmvn_norm_vars, + ) diff --git a/almeval/models/stepaudio/funasr_detach/frontends/utils/frontend.py b/almeval/models/stepaudio/funasr_detach/frontends/utils/frontend.py new file mode 100644 index 0000000000000000000000000000000000000000..bd6d730343892a5b98bee32e0b378bada145aa7a --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/frontends/utils/frontend.py @@ -0,0 +1,151 @@ +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union + +import numpy +import torch +import torch.nn as nn +from torch_complex.tensor import ComplexTensor + +from funasr_detach.frontends.utils.dnn_beamformer import DNN_Beamformer +from funasr_detach.frontends.utils.dnn_wpe import DNN_WPE + + +class Frontend(nn.Module): + def __init__( + self, + idim: int, + # WPE options + use_wpe: bool = False, + wtype: str = "blstmp", + wlayers: int = 3, + wunits: int = 300, + wprojs: int = 320, + wdropout_rate: float = 0.0, + taps: int = 5, + delay: int = 3, + use_dnn_mask_for_wpe: bool = True, + # Beamformer options + use_beamformer: bool = False, + btype: str = "blstmp", + blayers: int = 3, + bunits: int = 300, + bprojs: int = 320, + bnmask: int = 2, + badim: int = 320, + ref_channel: int = -1, + bdropout_rate=0.0, + ): + super().__init__() + + self.use_beamformer = use_beamformer + self.use_wpe = use_wpe + self.use_dnn_mask_for_wpe = use_dnn_mask_for_wpe + # use frontend for all the data, + # e.g. in the case of multi-speaker speech separation + self.use_frontend_for_all = bnmask > 2 + + if self.use_wpe: + if self.use_dnn_mask_for_wpe: + # Use DNN for power estimation + # (Not observed significant gains) + iterations = 1 + else: + # Performing as conventional WPE, without DNN Estimator + iterations = 2 + + self.wpe = DNN_WPE( + wtype=wtype, + widim=idim, + wunits=wunits, + wprojs=wprojs, + wlayers=wlayers, + taps=taps, + delay=delay, + dropout_rate=wdropout_rate, + iterations=iterations, + use_dnn_mask=use_dnn_mask_for_wpe, + ) + else: + self.wpe = None + + if self.use_beamformer: + self.beamformer = DNN_Beamformer( + btype=btype, + bidim=idim, + bunits=bunits, + bprojs=bprojs, + blayers=blayers, + bnmask=bnmask, + dropout_rate=bdropout_rate, + badim=badim, + ref_channel=ref_channel, + ) + else: + self.beamformer = None + + def forward( + self, x: ComplexTensor, ilens: Union[torch.LongTensor, numpy.ndarray, List[int]] + ) -> Tuple[ComplexTensor, torch.LongTensor, Optional[ComplexTensor]]: + assert len(x) == len(ilens), (len(x), len(ilens)) + # (B, T, F) or (B, T, C, F) + if x.dim() not in (3, 4): + raise ValueError(f"Input dim must be 3 or 4: {x.dim()}") + if not torch.is_tensor(ilens): + ilens = torch.from_numpy(numpy.asarray(ilens)).to(x.device) + + mask = None + h = x + if h.dim() == 4: + if self.training: + choices = [(False, False)] if not self.use_frontend_for_all else [] + if self.use_wpe: + choices.append((True, False)) + + if self.use_beamformer: + choices.append((False, True)) + + use_wpe, use_beamformer = choices[numpy.random.randint(len(choices))] + + else: + use_wpe = self.use_wpe + use_beamformer = self.use_beamformer + + # 1. WPE + if use_wpe: + # h: (B, T, C, F) -> h: (B, T, C, F) + h, ilens, mask = self.wpe(h, ilens) + + # 2. Beamformer + if use_beamformer: + # h: (B, T, C, F) -> h: (B, T, F) + h, ilens, mask = self.beamformer(h, ilens) + + return h, ilens, mask + + +def frontend_for(args, idim): + return Frontend( + idim=idim, + # WPE options + use_wpe=args.use_wpe, + wtype=args.wtype, + wlayers=args.wlayers, + wunits=args.wunits, + wprojs=args.wprojs, + wdropout_rate=args.wdropout_rate, + taps=args.wpe_taps, + delay=args.wpe_delay, + use_dnn_mask_for_wpe=args.use_dnn_mask_for_wpe, + # Beamformer options + use_beamformer=args.use_beamformer, + btype=args.btype, + blayers=args.blayers, + bunits=args.bunits, + bprojs=args.bprojs, + bnmask=args.bnmask, + badim=args.badim, + ref_channel=args.ref_channel, + bdropout_rate=args.bdropout_rate, + ) diff --git a/almeval/models/stepaudio/funasr_detach/frontends/utils/log_mel.py b/almeval/models/stepaudio/funasr_detach/frontends/utils/log_mel.py new file mode 100644 index 0000000000000000000000000000000000000000..cbec82be1a6bb7627447dd2ae076f22f87c16f80 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/frontends/utils/log_mel.py @@ -0,0 +1,83 @@ +import librosa +import torch +from typing import Tuple + +from funasr_detach.models.transformer.utils.nets_utils import make_pad_mask + + +class LogMel(torch.nn.Module): + """Convert STFT to fbank feats + + The arguments is same as librosa.filters.mel + + Args: + fs: number > 0 [scalar] sampling rate of the incoming signal + n_fft: int > 0 [scalar] number of FFT components + n_mels: int > 0 [scalar] number of Mel bands to generate + fmin: float >= 0 [scalar] lowest frequency (in Hz) + fmax: float >= 0 [scalar] highest frequency (in Hz). + If `None`, use `fmax = fs / 2.0` + htk: use HTK formula instead of Slaney + """ + + def __init__( + self, + fs: int = 16000, + n_fft: int = 512, + n_mels: int = 80, + fmin: float = None, + fmax: float = None, + htk: bool = False, + log_base: float = None, + ): + super().__init__() + + fmin = 0 if fmin is None else fmin + fmax = fs / 2 if fmax is None else fmax + _mel_options = dict( + sr=fs, + n_fft=n_fft, + n_mels=n_mels, + fmin=fmin, + fmax=fmax, + htk=htk, + ) + self.mel_options = _mel_options + self.log_base = log_base + + # Note(kamo): The mel matrix of librosa is different from kaldi. + melmat = librosa.filters.mel(**_mel_options) + # melmat: (D2, D1) -> (D1, D2) + self.register_buffer("melmat", torch.from_numpy(melmat.T).float()) + + def extra_repr(self): + return ", ".join(f"{k}={v}" for k, v in self.mel_options.items()) + + def forward( + self, + feat: torch.Tensor, + ilens: torch.Tensor = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + # feat: (B, T, D1) x melmat: (D1, D2) -> mel_feat: (B, T, D2) + mel_feat = torch.matmul(feat, self.melmat) + mel_feat = torch.clamp(mel_feat, min=1e-10) + + if self.log_base is None: + logmel_feat = mel_feat.log() + elif self.log_base == 2.0: + logmel_feat = mel_feat.log2() + elif self.log_base == 10.0: + logmel_feat = mel_feat.log10() + else: + logmel_feat = mel_feat.log() / torch.log(self.log_base) + + # Zero padding + if ilens is not None: + logmel_feat = logmel_feat.masked_fill( + make_pad_mask(ilens, logmel_feat, 1), 0.0 + ) + else: + ilens = feat.new_full( + [feat.size(0)], fill_value=feat.size(1), dtype=torch.long + ) + return logmel_feat, ilens diff --git a/almeval/models/stepaudio/funasr_detach/frontends/utils/mask_estimator.py b/almeval/models/stepaudio/funasr_detach/frontends/utils/mask_estimator.py new file mode 100644 index 0000000000000000000000000000000000000000..3e4d8d396ded057ec3af6b93837938276b2a8a31 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/frontends/utils/mask_estimator.py @@ -0,0 +1,77 @@ +from typing import Tuple + +import numpy as np +import torch +from torch.nn import functional as F +from torch_complex.tensor import ComplexTensor + +from funasr_detach.models.transformer.utils.nets_utils import make_pad_mask +from funasr_detach.models.language_model.rnn.encoders import RNN +from funasr_detach.models.language_model.rnn.encoders import RNNP + + +class MaskEstimator(torch.nn.Module): + def __init__(self, type, idim, layers, units, projs, dropout, nmask=1): + super().__init__() + subsample = np.ones(layers + 1, dtype=np.int32) + + typ = type.lstrip("vgg").rstrip("p") + if type[-1] == "p": + self.brnn = RNNP(idim, layers, units, projs, subsample, dropout, typ=typ) + else: + self.brnn = RNN(idim, layers, units, projs, dropout, typ=typ) + + self.type = type + self.nmask = nmask + self.linears = torch.nn.ModuleList( + [torch.nn.Linear(projs, idim) for _ in range(nmask)] + ) + + def forward( + self, xs: ComplexTensor, ilens: torch.LongTensor + ) -> Tuple[Tuple[torch.Tensor, ...], torch.LongTensor]: + """The forward function + + Args: + xs: (B, F, C, T) + ilens: (B,) + Returns: + hs (torch.Tensor): The hidden vector (B, F, C, T) + masks: A tuple of the masks. (B, F, C, T) + ilens: (B,) + """ + assert xs.size(0) == ilens.size(0), (xs.size(0), ilens.size(0)) + _, _, C, input_length = xs.size() + # (B, F, C, T) -> (B, C, T, F) + xs = xs.permute(0, 2, 3, 1) + + # Calculate amplitude: (B, C, T, F) -> (B, C, T, F) + xs = (xs.real**2 + xs.imag**2) ** 0.5 + # xs: (B, C, T, F) -> xs: (B * C, T, F) + xs = xs.contiguous().view(-1, xs.size(-2), xs.size(-1)) + # ilens: (B,) -> ilens_: (B * C) + ilens_ = ilens[:, None].expand(-1, C).contiguous().view(-1) + + # xs: (B * C, T, F) -> xs: (B * C, T, D) + xs, _, _ = self.brnn(xs, ilens_) + # xs: (B * C, T, D) -> xs: (B, C, T, D) + xs = xs.view(-1, C, xs.size(-2), xs.size(-1)) + + masks = [] + for linear in self.linears: + # xs: (B, C, T, D) -> mask:(B, C, T, F) + mask = linear(xs) + + mask = torch.sigmoid(mask) + # Zero padding + mask.masked_fill(make_pad_mask(ilens, mask, length_dim=2), 0) + + # (B, C, T, F) -> (B, F, C, T) + mask = mask.permute(0, 3, 1, 2) + + # Take cares of multi gpu cases: If input_length > max(ilens) + if mask.size(-1) < input_length: + mask = F.pad(mask, [0, input_length - mask.size(-1)], value=0) + masks.append(mask) + + return tuple(masks), ilens diff --git a/almeval/models/stepaudio/funasr_detach/frontends/utils/stft.py b/almeval/models/stepaudio/funasr_detach/frontends/utils/stft.py new file mode 100644 index 0000000000000000000000000000000000000000..99d5c5e84ba891c5c8c0c3bb8ffa97538007709c --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/frontends/utils/stft.py @@ -0,0 +1,239 @@ +from distutils.version import LooseVersion +from typing import Optional +from typing import Tuple +from typing import Union + +import torch + +try: + from torch_complex.tensor import ComplexTensor +except: + print("Please install torch_complex firstly") +from funasr_detach.models.transformer.utils.nets_utils import make_pad_mask +from funasr_detach.frontends.utils.complex_utils import is_complex + +import librosa +import numpy as np + +is_torch_1_9_plus = LooseVersion(torch.__version__) >= LooseVersion("1.9.0") + + +is_torch_1_7_plus = LooseVersion(torch.__version__) >= LooseVersion("1.7") + + +class Stft(torch.nn.Module): + def __init__( + self, + n_fft: int = 512, + win_length: int = None, + hop_length: int = 128, + window: Optional[str] = "hann", + center: bool = True, + normalized: bool = False, + onesided: bool = True, + ): + super().__init__() + self.n_fft = n_fft + if win_length is None: + self.win_length = n_fft + else: + self.win_length = win_length + self.hop_length = hop_length + self.center = center + self.normalized = normalized + self.onesided = onesided + if window is not None and not hasattr(torch, f"{window}_window"): + if window.lower() != "povey": + raise ValueError(f"{window} window is not implemented") + self.window = window + + def extra_repr(self): + return ( + f"n_fft={self.n_fft}, " + f"win_length={self.win_length}, " + f"hop_length={self.hop_length}, " + f"center={self.center}, " + f"normalized={self.normalized}, " + f"onesided={self.onesided}" + ) + + def forward( + self, input: torch.Tensor, ilens: torch.Tensor = None + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """STFT forward function. + + Args: + input: (Batch, Nsamples) or (Batch, Nsample, Channels) + ilens: (Batch) + Returns: + output: (Batch, Frames, Freq, 2) or (Batch, Frames, Channels, Freq, 2) + + """ + bs = input.size(0) + if input.dim() == 3: + multi_channel = True + # input: (Batch, Nsample, Channels) -> (Batch * Channels, Nsample) + input = input.transpose(1, 2).reshape(-1, input.size(1)) + else: + multi_channel = False + + # NOTE(kamo): + # The default behaviour of torch.stft is compatible with librosa.stft + # about padding and scaling. + # Note that it's different from scipy.signal.stft + + # output: (Batch, Freq, Frames, 2=real_imag) + # or (Batch, Channel, Freq, Frames, 2=real_imag) + if self.window is not None: + if self.window.lower() == "povey": + window = torch.hann_window( + self.win_length, + periodic=False, + device=input.device, + dtype=input.dtype, + ).pow(0.85) + else: + window_func = getattr(torch, f"{self.window}_window") + window = window_func( + self.win_length, dtype=input.dtype, device=input.device + ) + else: + window = None + + # For the compatibility of ARM devices, which do not support + # torch.stft() due to the lake of MKL. + if input.is_cuda or torch.backends.mkl.is_available(): + stft_kwargs = dict( + n_fft=self.n_fft, + win_length=self.win_length, + hop_length=self.hop_length, + center=self.center, + window=window, + normalized=self.normalized, + onesided=self.onesided, + ) + if is_torch_1_7_plus: + stft_kwargs["return_complex"] = False + output = torch.stft(input, **stft_kwargs) + else: + if self.training: + raise NotImplementedError( + "stft is implemented with librosa on this device, which does not " + "support the training mode." + ) + + # use stft_kwargs to flexibly control different PyTorch versions' kwargs + stft_kwargs = dict( + n_fft=self.n_fft, + win_length=self.win_length, + hop_length=self.hop_length, + center=self.center, + window=window, + ) + + if window is not None: + # pad the given window to n_fft + n_pad_left = (self.n_fft - window.shape[0]) // 2 + n_pad_right = self.n_fft - window.shape[0] - n_pad_left + stft_kwargs["window"] = torch.cat( + [torch.zeros(n_pad_left), window, torch.zeros(n_pad_right)], 0 + ).numpy() + else: + win_length = ( + self.win_length if self.win_length is not None else self.n_fft + ) + stft_kwargs["window"] = torch.ones(win_length) + + output = [] + # iterate over istances in a batch + for i, instance in enumerate(input): + stft = librosa.stft(input[i].numpy(), **stft_kwargs) + output.append(torch.tensor(np.stack([stft.real, stft.imag], -1))) + output = torch.stack(output, 0) + if not self.onesided: + len_conj = self.n_fft - output.shape[1] + conj = output[:, 1 : 1 + len_conj].flip(1) + conj[:, :, :, -1].data *= -1 + output = torch.cat([output, conj], 1) + if self.normalized: + output = output * (stft_kwargs["window"].shape[0] ** (-0.5)) + + # output: (Batch, Freq, Frames, 2=real_imag) + # -> (Batch, Frames, Freq, 2=real_imag) + output = output.transpose(1, 2) + if multi_channel: + # output: (Batch * Channel, Frames, Freq, 2=real_imag) + # -> (Batch, Frame, Channel, Freq, 2=real_imag) + output = output.view(bs, -1, output.size(1), output.size(2), 2).transpose( + 1, 2 + ) + + if ilens is not None: + if self.center: + pad = self.n_fft // 2 + ilens = ilens + 2 * pad + + olens = (ilens - self.n_fft) // self.hop_length + 1 + output.masked_fill_(make_pad_mask(olens, output, 1), 0.0) + else: + olens = None + + return output, olens + + def inverse( + self, input: Union[torch.Tensor, ComplexTensor], ilens: torch.Tensor = None + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Inverse STFT. + + Args: + input: Tensor(batch, T, F, 2) or ComplexTensor(batch, T, F) + ilens: (batch,) + Returns: + wavs: (batch, samples) + ilens: (batch,) + """ + if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): + istft = torch.functional.istft + else: + try: + import torchaudio + except ImportError: + raise ImportError( + "Please install torchaudio>=0.3.0 or use torch>=1.6.0" + ) + + if not hasattr(torchaudio.functional, "istft"): + raise ImportError( + "Please install torchaudio>=0.3.0 or use torch>=1.6.0" + ) + istft = torchaudio.functional.istft + + if self.window is not None: + window_func = getattr(torch, f"{self.window}_window") + if is_complex(input): + datatype = input.real.dtype + else: + datatype = input.dtype + window = window_func(self.win_length, dtype=datatype, device=input.device) + else: + window = None + + if is_complex(input): + input = torch.stack([input.real, input.imag], dim=-1) + elif input.shape[-1] != 2: + raise TypeError("Invalid input type") + input = input.transpose(1, 2) + + wavs = istft( + input, + n_fft=self.n_fft, + hop_length=self.hop_length, + win_length=self.win_length, + window=window, + center=self.center, + normalized=self.normalized, + onesided=self.onesided, + length=ilens.max() if ilens is not None else ilens, + ) + + return wavs, ilens diff --git a/almeval/models/stepaudio/funasr_detach/losses/__init__.py b/almeval/models/stepaudio/funasr_detach/losses/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/almeval/models/stepaudio/funasr_detach/losses/label_smoothing_loss.py b/almeval/models/stepaudio/funasr_detach/losses/label_smoothing_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..d96a29905192e55d9abfdae86d2f75b7f3d0d02f --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/losses/label_smoothing_loss.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +# Copyright 2019 Shigeki Karita +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Label smoothing module.""" + +import torch +from torch import nn +from funasr_detach.models.transformer.utils.nets_utils import make_pad_mask + + +class LabelSmoothingLoss(nn.Module): + """Label-smoothing loss. + + :param int size: the number of class + :param int padding_idx: ignored class id + :param float smoothing: smoothing rate (0.0 means the conventional CE) + :param bool normalize_length: normalize loss by sequence length if True + :param torch.nn.Module criterion: loss function to be smoothed + """ + + def __init__( + self, + size, + padding_idx, + smoothing, + normalize_length=False, + criterion=nn.KLDivLoss(reduction="none"), + ): + """Construct an LabelSmoothingLoss object.""" + super(LabelSmoothingLoss, self).__init__() + self.criterion = criterion + self.padding_idx = padding_idx + self.confidence = 1.0 - smoothing + self.smoothing = smoothing + self.size = size + self.true_dist = None + self.normalize_length = normalize_length + + def forward(self, x, target): + """Compute loss between x and target. + + :param torch.Tensor x: prediction (batch, seqlen, class) + :param torch.Tensor target: + target signal masked with self.padding_id (batch, seqlen) + :return: scalar float value + :rtype torch.Tensor + """ + assert x.size(2) == self.size + batch_size = x.size(0) + x = x.view(-1, self.size) + target = target.view(-1) + with torch.no_grad(): + true_dist = x.clone() + true_dist.fill_(self.smoothing / (self.size - 1)) + ignore = target == self.padding_idx # (B,) + total = len(target) - ignore.sum().item() + target = target.masked_fill(ignore, 0) # avoid -1 index + true_dist.scatter_(1, target.unsqueeze(1), self.confidence) + kl = self.criterion(torch.log_softmax(x, dim=1), true_dist) + denom = total if self.normalize_length else batch_size + return kl.masked_fill(ignore.unsqueeze(1), 0).sum() / denom + + +class SequenceBinaryCrossEntropy(nn.Module): + def __init__( + self, normalize_length=False, criterion=nn.BCEWithLogitsLoss(reduction="none") + ): + super().__init__() + self.normalize_length = normalize_length + self.criterion = criterion + + def forward(self, pred, label, lengths): + pad_mask = make_pad_mask(lengths, maxlen=pred.shape[1]).to(pred.device) + loss = self.criterion(pred, label) + denom = (~pad_mask).sum() if self.normalize_length else pred.shape[0] + return loss.masked_fill(pad_mask.unsqueeze(-1), 0).sum() / denom + + +class NllLoss(nn.Module): + """Nll loss. + + :param int size: the number of class + :param int padding_idx: ignored class id + :param bool normalize_length: normalize loss by sequence length if True + :param torch.nn.Module criterion: loss function + """ + + def __init__( + self, + size, + padding_idx, + normalize_length=False, + criterion=nn.NLLLoss(reduction="none"), + ): + """Construct an NllLoss object.""" + super(NllLoss, self).__init__() + self.criterion = criterion + self.padding_idx = padding_idx + self.size = size + self.true_dist = None + self.normalize_length = normalize_length + + def forward(self, x, target): + """Compute loss between x and target. + + :param torch.Tensor x: prediction (batch, seqlen, class) + :param torch.Tensor target: + target signal masked with self.padding_id (batch, seqlen) + :return: scalar float value + :rtype torch.Tensor + """ + assert x.size(2) == self.size + batch_size = x.size(0) + x = x.view(-1, self.size) + target = target.view(-1) + with torch.no_grad(): + ignore = target == self.padding_idx # (B,) + total = len(target) - ignore.sum().item() + target = target.masked_fill(ignore, 0) # avoid -1 index + kl = self.criterion(x, target) + denom = total if self.normalize_length else batch_size + return kl.masked_fill(ignore, 0).sum() / denom diff --git a/almeval/models/stepaudio/funasr_detach/models/campplus/utils.py b/almeval/models/stepaudio/funasr_detach/models/campplus/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..14f97dceb11410684d92cee669f8264889b3ade5 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/campplus/utils.py @@ -0,0 +1,520 @@ +#!/usr/bin/env python3 +# -*- encoding: utf-8 -*- +# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved. +# MIT License (https://opensource.org/licenses/MIT) +# Modified from 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker) + +import io +import os +import torch +import requests +import tempfile +import contextlib +import numpy as np +import librosa as sf +from typing import Union +from pathlib import Path +from typing import Generator, Union +from abc import ABCMeta, abstractmethod +import torchaudio.compliance.kaldi as Kaldi + +from funasr_detach.models.transformer.utils.nets_utils import pad_list + + +def check_audio_list(audio: list): + audio_dur = 0 + for i in range(len(audio)): + seg = audio[i] + assert seg[1] >= seg[0], "modelscope error: Wrong time stamps." + assert isinstance(seg[2], np.ndarray), "modelscope error: Wrong data type." + assert ( + int(seg[1] * 16000) - int(seg[0] * 16000) == seg[2].shape[0] + ), "modelscope error: audio data in list is inconsistent with time length." + if i > 0: + assert seg[0] >= audio[i - 1][1], "modelscope error: Wrong time stamps." + audio_dur += seg[1] - seg[0] + return audio_dur + # assert audio_dur > 5, 'modelscope error: The effective audio duration is too short.' + + +def sv_preprocess(inputs: Union[np.ndarray, list]): + output = [] + for i in range(len(inputs)): + if isinstance(inputs[i], str): + file_bytes = File.read(inputs[i]) + data, fs = sf.load(io.BytesIO(file_bytes), dtype="float32") + if len(data.shape) == 2: + data = data[:, 0] + data = torch.from_numpy(data).unsqueeze(0) + data = data.squeeze(0) + elif isinstance(inputs[i], np.ndarray): + assert ( + len(inputs[i].shape) == 1 + ), "modelscope error: Input array should be [N, T]" + data = inputs[i] + if data.dtype in ["int16", "int32", "int64"]: + data = (data / (1 << 15)).astype("float32") + else: + data = data.astype("float32") + data = torch.from_numpy(data) + else: + raise ValueError( + "modelscope error: The input type is restricted to audio address and nump array." + ) + output.append(data) + return output + + +def sv_chunk(vad_segments: list, fs=16000) -> list: + config = { + "seg_dur": 1.5, + "seg_shift": 0.75, + } + + def seg_chunk(seg_data): + seg_st = seg_data[0] + data = seg_data[2] + chunk_len = int(config["seg_dur"] * fs) + chunk_shift = int(config["seg_shift"] * fs) + last_chunk_ed = 0 + seg_res = [] + for chunk_st in range(0, data.shape[0], chunk_shift): + chunk_ed = min(chunk_st + chunk_len, data.shape[0]) + if chunk_ed <= last_chunk_ed: + break + last_chunk_ed = chunk_ed + chunk_st = max(0, chunk_ed - chunk_len) + chunk_data = data[chunk_st:chunk_ed] + if chunk_data.shape[0] < chunk_len: + chunk_data = np.pad( + chunk_data, (0, chunk_len - chunk_data.shape[0]), "constant" + ) + seg_res.append([chunk_st / fs + seg_st, chunk_ed / fs + seg_st, chunk_data]) + return seg_res + + segs = [] + for i, s in enumerate(vad_segments): + segs.extend(seg_chunk(s)) + + return segs + + +def extract_feature(audio): + features = [] + feature_times = [] + feature_lengths = [] + for au in audio: + feature = Kaldi.fbank(au.unsqueeze(0), num_mel_bins=80) + feature = feature - feature.mean(dim=0, keepdim=True) + features.append(feature) + feature_times.append(au.shape[0]) + feature_lengths.append(feature.shape[0]) + # padding for batch inference + features_padded = pad_list(features, pad_value=0) + # features = torch.cat(features) + return features_padded, feature_lengths, feature_times + + +def postprocess( + segments: list, vad_segments: list, labels: np.ndarray, embeddings: np.ndarray +) -> list: + assert len(segments) == len(labels) + labels = correct_labels(labels) + distribute_res = [] + for i in range(len(segments)): + distribute_res.append([segments[i][0], segments[i][1], labels[i]]) + # merge the same speakers chronologically + distribute_res = merge_seque(distribute_res) + + # accquire speaker center + spk_embs = [] + for i in range(labels.max() + 1): + spk_emb = embeddings[labels == i].mean(0) + spk_embs.append(spk_emb) + spk_embs = np.stack(spk_embs) + + def is_overlapped(t1, t2): + if t1 > t2 + 1e-4: + return True + return False + + # distribute the overlap region + for i in range(1, len(distribute_res)): + if is_overlapped(distribute_res[i - 1][1], distribute_res[i][0]): + p = (distribute_res[i][0] + distribute_res[i - 1][1]) / 2 + distribute_res[i][0] = p + distribute_res[i - 1][1] = p + + # smooth the result + distribute_res = smooth(distribute_res) + + return distribute_res + + +def correct_labels(labels): + labels_id = 0 + id2id = {} + new_labels = [] + for i in labels: + if i not in id2id: + id2id[i] = labels_id + labels_id += 1 + new_labels.append(id2id[i]) + return np.array(new_labels) + + +def merge_seque(distribute_res): + res = [distribute_res[0]] + for i in range(1, len(distribute_res)): + if distribute_res[i][2] != res[-1][2] or distribute_res[i][0] > res[-1][1]: + res.append(distribute_res[i]) + else: + res[-1][1] = distribute_res[i][1] + return res + + +def smooth(res, mindur=1): + # short segments are assigned to nearest speakers. + for i in range(len(res)): + res[i][0] = round(res[i][0], 2) + res[i][1] = round(res[i][1], 2) + if res[i][1] - res[i][0] < mindur: + if i == 0: + res[i][2] = res[i + 1][2] + elif i == len(res) - 1: + res[i][2] = res[i - 1][2] + elif res[i][0] - res[i - 1][1] <= res[i + 1][0] - res[i][1]: + res[i][2] = res[i - 1][2] + else: + res[i][2] = res[i + 1][2] + # merge the speakers + res = merge_seque(res) + + return res + + +def distribute_spk(sentence_list, sd_time_list): + sd_sentence_list = [] + for d in sentence_list: + sentence_start = d["start"] + sentence_end = d["end"] + sentence_spk = 0 + max_overlap = 0 + for sd_time in sd_time_list: + spk_st, spk_ed, spk = sd_time + spk_st = spk_st * 1000 + spk_ed = spk_ed * 1000 + overlap = max(min(sentence_end, spk_ed) - max(sentence_start, spk_st), 0) + if overlap > max_overlap: + max_overlap = overlap + sentence_spk = spk + d["spk"] = int(sentence_spk) + sd_sentence_list.append(d) + return sd_sentence_list + + +class Storage(metaclass=ABCMeta): + """Abstract class of storage. + + All backends need to implement two apis: ``read()`` and ``read_text()``. + ``read()`` reads the file as a byte stream and ``read_text()`` reads + the file as texts. + """ + + @abstractmethod + def read(self, filepath: str): + pass + + @abstractmethod + def read_text(self, filepath: str): + pass + + @abstractmethod + def write(self, obj: bytes, filepath: Union[str, Path]) -> None: + pass + + @abstractmethod + def write_text( + self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8" + ) -> None: + pass + + +class LocalStorage(Storage): + """Local hard disk storage""" + + def read(self, filepath: Union[str, Path]) -> bytes: + """Read data from a given ``filepath`` with 'rb' mode. + + Args: + filepath (str or Path): Path to read data. + + Returns: + bytes: Expected bytes object. + """ + with open(filepath, "rb") as f: + content = f.read() + return content + + def read_text(self, filepath: Union[str, Path], encoding: str = "utf-8") -> str: + """Read data from a given ``filepath`` with 'r' mode. + + Args: + filepath (str or Path): Path to read data. + encoding (str): The encoding format used to open the ``filepath``. + Default: 'utf-8'. + + Returns: + str: Expected text reading from ``filepath``. + """ + with open(filepath, "r", encoding=encoding) as f: + value_buf = f.read() + return value_buf + + def write(self, obj: bytes, filepath: Union[str, Path]) -> None: + """Write data to a given ``filepath`` with 'wb' mode. + + Note: + ``write`` will create a directory if the directory of ``filepath`` + does not exist. + + Args: + obj (bytes): Data to be written. + filepath (str or Path): Path to write data. + """ + dirname = os.path.dirname(filepath) + if dirname and not os.path.exists(dirname): + os.makedirs(dirname, exist_ok=True) + + with open(filepath, "wb") as f: + f.write(obj) + + def write_text( + self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8" + ) -> None: + """Write data to a given ``filepath`` with 'w' mode. + + Note: + ``write_text`` will create a directory if the directory of + ``filepath`` does not exist. + + Args: + obj (str): Data to be written. + filepath (str or Path): Path to write data. + encoding (str): The encoding format used to open the ``filepath``. + Default: 'utf-8'. + """ + dirname = os.path.dirname(filepath) + if dirname and not os.path.exists(dirname): + os.makedirs(dirname, exist_ok=True) + + with open(filepath, "w", encoding=encoding) as f: + f.write(obj) + + @contextlib.contextmanager + def as_local_path( + self, filepath: Union[str, Path] + ) -> Generator[Union[str, Path], None, None]: + """Only for unified API and do nothing.""" + yield filepath + + +class HTTPStorage(Storage): + """HTTP and HTTPS storage.""" + + def read(self, url): + # TODO @wenmeng.zwm add progress bar if file is too large + r = requests.get(url) + r.raise_for_status() + return r.content + + def read_text(self, url): + r = requests.get(url) + r.raise_for_status() + return r.text + + @contextlib.contextmanager + def as_local_path(self, filepath: str) -> Generator[Union[str, Path], None, None]: + """Download a file from ``filepath``. + + ``as_local_path`` is decorated by :meth:`contextlib.contextmanager`. It + can be called with ``with`` statement, and when exists from the + ``with`` statement, the temporary path will be released. + + Args: + filepath (str): Download a file from ``filepath``. + + Examples: + >>> storage = HTTPStorage() + >>> # After existing from the ``with`` clause, + >>> # the path will be removed + >>> with storage.get_local_path('http://path/to/file') as path: + ... # do something here + """ + try: + f = tempfile.NamedTemporaryFile(delete=False) + f.write(self.read(filepath)) + f.close() + yield f.name + finally: + os.remove(f.name) + + def write(self, obj: bytes, url: Union[str, Path]) -> None: + raise NotImplementedError("write is not supported by HTTP Storage") + + def write_text( + self, obj: str, url: Union[str, Path], encoding: str = "utf-8" + ) -> None: + raise NotImplementedError("write_text is not supported by HTTP Storage") + + +class OSSStorage(Storage): + """OSS storage.""" + + def __init__(self, oss_config_file=None): + # read from config file or env var + raise NotImplementedError("OSSStorage.__init__ to be implemented in the future") + + def read(self, filepath): + raise NotImplementedError("OSSStorage.read to be implemented in the future") + + def read_text(self, filepath, encoding="utf-8"): + raise NotImplementedError( + "OSSStorage.read_text to be implemented in the future" + ) + + @contextlib.contextmanager + def as_local_path(self, filepath: str) -> Generator[Union[str, Path], None, None]: + """Download a file from ``filepath``. + + ``as_local_path`` is decorated by :meth:`contextlib.contextmanager`. It + can be called with ``with`` statement, and when exists from the + ``with`` statement, the temporary path will be released. + + Args: + filepath (str): Download a file from ``filepath``. + + Examples: + >>> storage = OSSStorage() + >>> # After existing from the ``with`` clause, + >>> # the path will be removed + >>> with storage.get_local_path('http://path/to/file') as path: + ... # do something here + """ + try: + f = tempfile.NamedTemporaryFile(delete=False) + f.write(self.read(filepath)) + f.close() + yield f.name + finally: + os.remove(f.name) + + def write(self, obj: bytes, filepath: Union[str, Path]) -> None: + raise NotImplementedError("OSSStorage.write to be implemented in the future") + + def write_text( + self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8" + ) -> None: + raise NotImplementedError( + "OSSStorage.write_text to be implemented in the future" + ) + + +G_STORAGES = {} + + +class File(object): + _prefix_to_storage: dict = { + "oss": OSSStorage, + "http": HTTPStorage, + "https": HTTPStorage, + "local": LocalStorage, + } + + @staticmethod + def _get_storage(uri): + assert isinstance(uri, str), f"uri should be str type, but got {type(uri)}" + + if "://" not in uri: + # local path + storage_type = "local" + else: + prefix, _ = uri.split("://") + storage_type = prefix + + assert storage_type in File._prefix_to_storage, ( + f"Unsupported uri {uri}, valid prefixs: " + f"{list(File._prefix_to_storage.keys())}" + ) + + if storage_type not in G_STORAGES: + G_STORAGES[storage_type] = File._prefix_to_storage[storage_type]() + + return G_STORAGES[storage_type] + + @staticmethod + def read(uri: str) -> bytes: + """Read data from a given ``filepath`` with 'rb' mode. + + Args: + filepath (str or Path): Path to read data. + + Returns: + bytes: Expected bytes object. + """ + storage = File._get_storage(uri) + return storage.read(uri) + + @staticmethod + def read_text(uri: Union[str, Path], encoding: str = "utf-8") -> str: + """Read data from a given ``filepath`` with 'r' mode. + + Args: + filepath (str or Path): Path to read data. + encoding (str): The encoding format used to open the ``filepath``. + Default: 'utf-8'. + + Returns: + str: Expected text reading from ``filepath``. + """ + storage = File._get_storage(uri) + return storage.read_text(uri) + + @staticmethod + def write(obj: bytes, uri: Union[str, Path]) -> None: + """Write data to a given ``filepath`` with 'wb' mode. + + Note: + ``write`` will create a directory if the directory of ``filepath`` + does not exist. + + Args: + obj (bytes): Data to be written. + filepath (str or Path): Path to write data. + """ + storage = File._get_storage(uri) + return storage.write(obj, uri) + + @staticmethod + def write_text(obj: str, uri: str, encoding: str = "utf-8") -> None: + """Write data to a given ``filepath`` with 'w' mode. + + Note: + ``write_text`` will create a directory if the directory of + ``filepath`` does not exist. + + Args: + obj (str): Data to be written. + filepath (str or Path): Path to write data. + encoding (str): The encoding format used to open the ``filepath``. + Default: 'utf-8'. + """ + storage = File._get_storage(uri) + return storage.write_text(obj, uri) + + @contextlib.contextmanager + def as_local_path(uri: str) -> Generator[Union[str, Path], None, None]: + """Only for unified API and do nothing.""" + storage = File._get_storage(uri) + with storage.as_local_path(uri) as local_path: + yield local_path diff --git a/almeval/models/stepaudio/funasr_detach/models/conformer/__init__.py b/almeval/models/stepaudio/funasr_detach/models/conformer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/almeval/models/stepaudio/funasr_detach/models/conformer/encoder.py b/almeval/models/stepaudio/funasr_detach/models/conformer/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..c931996c68279a5768f14e26028b24ae5a6e8b22 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/conformer/encoder.py @@ -0,0 +1,1281 @@ +# Copyright 2020 Tomoki Hayashi +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Conformer encoder definition.""" + +import logging +from typing import Union, Dict, List, Tuple, Optional + +import torch +from torch import nn + +from funasr_detach.models.ctc.ctc import CTC +from funasr_detach.models.transformer.attention import ( + MultiHeadedAttention, # noqa: H301 + RelPositionMultiHeadedAttention, # noqa: H301 + LegacyRelPositionMultiHeadedAttention, # noqa: H301 + RelPositionMultiHeadedAttentionChunk, +) +from funasr_detach.models.transformer.embedding import ( + PositionalEncoding, # noqa: H301 + ScaledPositionalEncoding, # noqa: H301 + RelPositionalEncoding, # noqa: H301 + LegacyRelPositionalEncoding, # noqa: H301 + StreamingRelPositionalEncoding, +) +from funasr_detach.models.transformer.layer_norm import LayerNorm +from funasr_detach.models.transformer.utils.multi_layer_conv import Conv1dLinear +from funasr_detach.models.transformer.utils.multi_layer_conv import MultiLayeredConv1d +from funasr_detach.models.transformer.utils.nets_utils import get_activation +from funasr_detach.models.transformer.utils.nets_utils import make_pad_mask +from funasr_detach.models.transformer.utils.nets_utils import ( + TooShortUttError, + check_short_utt, + make_chunk_mask, + make_source_mask, +) +from funasr_detach.models.transformer.positionwise_feed_forward import ( + PositionwiseFeedForward, # noqa: H301 +) +from funasr_detach.models.transformer.utils.repeat import repeat, MultiBlocks +from funasr_detach.models.transformer.utils.subsampling import Conv2dSubsampling +from funasr_detach.models.transformer.utils.subsampling import Conv2dSubsampling2 +from funasr_detach.models.transformer.utils.subsampling import Conv2dSubsampling6 +from funasr_detach.models.transformer.utils.subsampling import Conv2dSubsampling8 +from funasr_detach.models.transformer.utils.subsampling import TooShortUttError +from funasr_detach.models.transformer.utils.subsampling import check_short_utt +from funasr_detach.models.transformer.utils.subsampling import Conv2dSubsamplingPad +from funasr_detach.models.transformer.utils.subsampling import StreamingConvInput +from funasr_detach.register import tables + + +class ConvolutionModule(nn.Module): + """ConvolutionModule in Conformer model. + + Args: + channels (int): The number of channels of conv layers. + kernel_size (int): Kernerl size of conv layers. + + """ + + def __init__(self, channels, kernel_size, activation=nn.ReLU(), bias=True): + """Construct an ConvolutionModule object.""" + super(ConvolutionModule, self).__init__() + # kernerl_size should be a odd number for 'SAME' padding + assert (kernel_size - 1) % 2 == 0 + + self.pointwise_conv1 = nn.Conv1d( + channels, + 2 * channels, + kernel_size=1, + stride=1, + padding=0, + bias=bias, + ) + self.depthwise_conv = nn.Conv1d( + channels, + channels, + kernel_size, + stride=1, + padding=(kernel_size - 1) // 2, + groups=channels, + bias=bias, + ) + self.norm = nn.BatchNorm1d(channels) + self.pointwise_conv2 = nn.Conv1d( + channels, + channels, + kernel_size=1, + stride=1, + padding=0, + bias=bias, + ) + self.activation = activation + + def forward(self, x): + """Compute convolution module. + + Args: + x (torch.Tensor): Input tensor (#batch, time, channels). + + Returns: + torch.Tensor: Output tensor (#batch, time, channels). + + """ + # exchange the temporal dimension and the feature dimension + x = x.transpose(1, 2) + + # GLU mechanism + x = self.pointwise_conv1(x) # (batch, 2*channel, dim) + x = nn.functional.glu(x, dim=1) # (batch, channel, dim) + + # 1D Depthwise Conv + x = self.depthwise_conv(x) + x = self.activation(self.norm(x)) + + x = self.pointwise_conv2(x) + + return x.transpose(1, 2) + + +class EncoderLayer(nn.Module): + """Encoder layer module. + + Args: + size (int): Input dimension. + self_attn (torch.nn.Module): Self-attention module instance. + `MultiHeadedAttention` or `RelPositionMultiHeadedAttention` instance + can be used as the argument. + feed_forward (torch.nn.Module): Feed-forward module instance. + `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance + can be used as the argument. + feed_forward_macaron (torch.nn.Module): Additional feed-forward module instance. + `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance + can be used as the argument. + conv_module (torch.nn.Module): Convolution module instance. + `ConvlutionModule` instance can be used as the argument. + dropout_rate (float): Dropout rate. + normalize_before (bool): Whether to use layer_norm before the first block. + concat_after (bool): Whether to concat attention layer's input and output. + if True, additional linear will be applied. + i.e. x -> x + linear(concat(x, att(x))) + if False, no additional linear will be applied. i.e. x -> x + att(x) + stochastic_depth_rate (float): Proability to skip this layer. + During training, the layer may skip residual computation and return input + as-is with given probability. + """ + + def __init__( + self, + size, + self_attn, + feed_forward, + feed_forward_macaron, + conv_module, + dropout_rate, + normalize_before=True, + concat_after=False, + stochastic_depth_rate=0.0, + ): + """Construct an EncoderLayer object.""" + super(EncoderLayer, self).__init__() + self.self_attn = self_attn + self.feed_forward = feed_forward + self.feed_forward_macaron = feed_forward_macaron + self.conv_module = conv_module + self.norm_ff = LayerNorm(size) # for the FNN module + self.norm_mha = LayerNorm(size) # for the MHA module + if feed_forward_macaron is not None: + self.norm_ff_macaron = LayerNorm(size) + self.ff_scale = 0.5 + else: + self.ff_scale = 1.0 + if self.conv_module is not None: + self.norm_conv = LayerNorm(size) # for the CNN module + self.norm_final = LayerNorm(size) # for the final output of the block + self.dropout = nn.Dropout(dropout_rate) + self.size = size + self.normalize_before = normalize_before + self.concat_after = concat_after + if self.concat_after: + self.concat_linear = nn.Linear(size + size, size) + self.stochastic_depth_rate = stochastic_depth_rate + + def forward(self, x_input, mask, cache=None): + """Compute encoded features. + + Args: + x_input (Union[Tuple, torch.Tensor]): Input tensor w/ or w/o pos emb. + - w/ pos emb: Tuple of tensors [(#batch, time, size), (1, time, size)]. + - w/o pos emb: Tensor (#batch, time, size). + mask (torch.Tensor): Mask tensor for the input (#batch, time). + cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size). + + Returns: + torch.Tensor: Output tensor (#batch, time, size). + torch.Tensor: Mask tensor (#batch, time). + + """ + if isinstance(x_input, tuple): + x, pos_emb = x_input[0], x_input[1] + else: + x, pos_emb = x_input, None + + skip_layer = False + # with stochastic depth, residual connection `x + f(x)` becomes + # `x <- x + 1 / (1 - p) * f(x)` at training time. + stoch_layer_coeff = 1.0 + if self.training and self.stochastic_depth_rate > 0: + skip_layer = torch.rand(1).item() < self.stochastic_depth_rate + stoch_layer_coeff = 1.0 / (1 - self.stochastic_depth_rate) + + if skip_layer: + if cache is not None: + x = torch.cat([cache, x], dim=1) + if pos_emb is not None: + return (x, pos_emb), mask + return x, mask + + # whether to use macaron style + if self.feed_forward_macaron is not None: + residual = x + if self.normalize_before: + x = self.norm_ff_macaron(x) + x = residual + stoch_layer_coeff * self.ff_scale * self.dropout( + self.feed_forward_macaron(x) + ) + if not self.normalize_before: + x = self.norm_ff_macaron(x) + + # multi-headed self-attention module + residual = x + if self.normalize_before: + x = self.norm_mha(x) + + if cache is None: + x_q = x + else: + assert cache.shape == (x.shape[0], x.shape[1] - 1, self.size) + x_q = x[:, -1:, :] + residual = residual[:, -1:, :] + mask = None if mask is None else mask[:, -1:, :] + + if pos_emb is not None: + x_att = self.self_attn(x_q, x, x, pos_emb, mask) + else: + x_att = self.self_attn(x_q, x, x, mask) + + if self.concat_after: + x_concat = torch.cat((x, x_att), dim=-1) + x = residual + stoch_layer_coeff * self.concat_linear(x_concat) + else: + x = residual + stoch_layer_coeff * self.dropout(x_att) + if not self.normalize_before: + x = self.norm_mha(x) + + # convolution module + if self.conv_module is not None: + residual = x + if self.normalize_before: + x = self.norm_conv(x) + x = residual + stoch_layer_coeff * self.dropout(self.conv_module(x)) + if not self.normalize_before: + x = self.norm_conv(x) + + # feed forward module + residual = x + if self.normalize_before: + x = self.norm_ff(x) + x = residual + stoch_layer_coeff * self.ff_scale * self.dropout( + self.feed_forward(x) + ) + if not self.normalize_before: + x = self.norm_ff(x) + + if self.conv_module is not None: + x = self.norm_final(x) + + if cache is not None: + x = torch.cat([cache, x], dim=1) + + if pos_emb is not None: + return (x, pos_emb), mask + + return x, mask + + +@tables.register("encoder_classes", "ConformerEncoder") +class ConformerEncoder(nn.Module): + """Conformer encoder module. + + Args: + input_size (int): Input dimension. + output_size (int): Dimension of attention. + attention_heads (int): The number of heads of multi head attention. + linear_units (int): The number of units of position-wise feed forward. + num_blocks (int): The number of decoder blocks. + dropout_rate (float): Dropout rate. + attention_dropout_rate (float): Dropout rate in attention. + positional_dropout_rate (float): Dropout rate after adding positional encoding. + input_layer (Union[str, torch.nn.Module]): Input layer type. + normalize_before (bool): Whether to use layer_norm before the first block. + concat_after (bool): Whether to concat attention layer's input and output. + If True, additional linear will be applied. + i.e. x -> x + linear(concat(x, att(x))) + If False, no additional linear will be applied. i.e. x -> x + att(x) + positionwise_layer_type (str): "linear", "conv1d", or "conv1d-linear". + positionwise_conv_kernel_size (int): Kernel size of positionwise conv1d layer. + rel_pos_type (str): Whether to use the latest relative positional encoding or + the legacy one. The legacy relative positional encoding will be deprecated + in the future. More Details can be found in + https://github.com/espnet/espnet/pull/2816. + encoder_pos_enc_layer_type (str): Encoder positional encoding layer type. + encoder_attn_layer_type (str): Encoder attention layer type. + activation_type (str): Encoder activation function type. + macaron_style (bool): Whether to use macaron style for positionwise layer. + use_cnn_module (bool): Whether to use convolution module. + zero_triu (bool): Whether to zero the upper triangular part of attention matrix. + cnn_module_kernel (int): Kernerl size of convolution module. + padding_idx (int): Padding idx for input_layer=embed. + + """ + + def __init__( + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.0, + input_layer: str = "conv2d", + normalize_before: bool = True, + concat_after: bool = False, + positionwise_layer_type: str = "linear", + positionwise_conv_kernel_size: int = 3, + macaron_style: bool = False, + rel_pos_type: str = "legacy", + pos_enc_layer_type: str = "rel_pos", + selfattention_layer_type: str = "rel_selfattn", + activation_type: str = "swish", + use_cnn_module: bool = True, + zero_triu: bool = False, + cnn_module_kernel: int = 31, + padding_idx: int = -1, + interctc_layer_idx: List[int] = [], + interctc_use_conditioning: bool = False, + stochastic_depth_rate: Union[float, List[float]] = 0.0, + ): + super().__init__() + self._output_size = output_size + + if rel_pos_type == "legacy": + if pos_enc_layer_type == "rel_pos": + pos_enc_layer_type = "legacy_rel_pos" + if selfattention_layer_type == "rel_selfattn": + selfattention_layer_type = "legacy_rel_selfattn" + elif rel_pos_type == "latest": + assert selfattention_layer_type != "legacy_rel_selfattn" + assert pos_enc_layer_type != "legacy_rel_pos" + else: + raise ValueError("unknown rel_pos_type: " + rel_pos_type) + + activation = get_activation(activation_type) + if pos_enc_layer_type == "abs_pos": + pos_enc_class = PositionalEncoding + elif pos_enc_layer_type == "scaled_abs_pos": + pos_enc_class = ScaledPositionalEncoding + elif pos_enc_layer_type == "rel_pos": + assert selfattention_layer_type == "rel_selfattn" + pos_enc_class = RelPositionalEncoding + elif pos_enc_layer_type == "legacy_rel_pos": + assert selfattention_layer_type == "legacy_rel_selfattn" + pos_enc_class = LegacyRelPositionalEncoding + logging.warning( + "Using legacy_rel_pos and it will be deprecated in the future." + ) + else: + raise ValueError("unknown pos_enc_layer: " + pos_enc_layer_type) + + if input_layer == "linear": + self.embed = torch.nn.Sequential( + torch.nn.Linear(input_size, output_size), + torch.nn.LayerNorm(output_size), + torch.nn.Dropout(dropout_rate), + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer == "conv2d": + self.embed = Conv2dSubsampling( + input_size, + output_size, + dropout_rate, + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer == "conv2dpad": + self.embed = Conv2dSubsamplingPad( + input_size, + output_size, + dropout_rate, + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer == "conv2d2": + self.embed = Conv2dSubsampling2( + input_size, + output_size, + dropout_rate, + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer == "conv2d6": + self.embed = Conv2dSubsampling6( + input_size, + output_size, + dropout_rate, + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer == "conv2d8": + self.embed = Conv2dSubsampling8( + input_size, + output_size, + dropout_rate, + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer == "embed": + self.embed = torch.nn.Sequential( + torch.nn.Embedding(input_size, output_size, padding_idx=padding_idx), + pos_enc_class(output_size, positional_dropout_rate), + ) + elif isinstance(input_layer, torch.nn.Module): + self.embed = torch.nn.Sequential( + input_layer, + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer is None: + self.embed = torch.nn.Sequential( + pos_enc_class(output_size, positional_dropout_rate) + ) + else: + raise ValueError("unknown input_layer: " + input_layer) + self.normalize_before = normalize_before + if positionwise_layer_type == "linear": + positionwise_layer = PositionwiseFeedForward + positionwise_layer_args = ( + output_size, + linear_units, + dropout_rate, + activation, + ) + elif positionwise_layer_type == "conv1d": + positionwise_layer = MultiLayeredConv1d + positionwise_layer_args = ( + output_size, + linear_units, + positionwise_conv_kernel_size, + dropout_rate, + ) + elif positionwise_layer_type == "conv1d-linear": + positionwise_layer = Conv1dLinear + positionwise_layer_args = ( + output_size, + linear_units, + positionwise_conv_kernel_size, + dropout_rate, + ) + else: + raise NotImplementedError("Support only linear or conv1d.") + + if selfattention_layer_type == "selfattn": + encoder_selfattn_layer = MultiHeadedAttention + encoder_selfattn_layer_args = ( + attention_heads, + output_size, + attention_dropout_rate, + ) + elif selfattention_layer_type == "legacy_rel_selfattn": + assert pos_enc_layer_type == "legacy_rel_pos" + encoder_selfattn_layer = LegacyRelPositionMultiHeadedAttention + encoder_selfattn_layer_args = ( + attention_heads, + output_size, + attention_dropout_rate, + ) + logging.warning( + "Using legacy_rel_selfattn and it will be deprecated in the future." + ) + elif selfattention_layer_type == "rel_selfattn": + assert pos_enc_layer_type == "rel_pos" + encoder_selfattn_layer = RelPositionMultiHeadedAttention + encoder_selfattn_layer_args = ( + attention_heads, + output_size, + attention_dropout_rate, + zero_triu, + ) + else: + raise ValueError("unknown encoder_attn_layer: " + selfattention_layer_type) + + convolution_layer = ConvolutionModule + convolution_layer_args = (output_size, cnn_module_kernel, activation) + + if isinstance(stochastic_depth_rate, float): + stochastic_depth_rate = [stochastic_depth_rate] * num_blocks + + if len(stochastic_depth_rate) != num_blocks: + raise ValueError( + f"Length of stochastic_depth_rate ({len(stochastic_depth_rate)}) " + f"should be equal to num_blocks ({num_blocks})" + ) + + self.encoders = repeat( + num_blocks, + lambda lnum: EncoderLayer( + output_size, + encoder_selfattn_layer(*encoder_selfattn_layer_args), + positionwise_layer(*positionwise_layer_args), + positionwise_layer(*positionwise_layer_args) if macaron_style else None, + convolution_layer(*convolution_layer_args) if use_cnn_module else None, + dropout_rate, + normalize_before, + concat_after, + stochastic_depth_rate[lnum], + ), + ) + if self.normalize_before: + self.after_norm = LayerNorm(output_size) + + self.interctc_layer_idx = interctc_layer_idx + if len(interctc_layer_idx) > 0: + assert 0 < min(interctc_layer_idx) and max(interctc_layer_idx) < num_blocks + self.interctc_use_conditioning = interctc_use_conditioning + self.conditioning_layer = None + + def output_size(self) -> int: + return self._output_size + + def forward( + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, + ctc: CTC = None, + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Calculate forward propagation. + + Args: + xs_pad (torch.Tensor): Input tensor (#batch, L, input_size). + ilens (torch.Tensor): Input length (#batch). + prev_states (torch.Tensor): Not to be used now. + + Returns: + torch.Tensor: Output tensor (#batch, L, output_size). + torch.Tensor: Output length (#batch). + torch.Tensor: Not to be used now. + + """ + masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device) + + if ( + isinstance(self.embed, Conv2dSubsampling) + or isinstance(self.embed, Conv2dSubsampling2) + or isinstance(self.embed, Conv2dSubsampling6) + or isinstance(self.embed, Conv2dSubsampling8) + or isinstance(self.embed, Conv2dSubsamplingPad) + ): + short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1)) + if short_status: + raise TooShortUttError( + f"has {xs_pad.size(1)} frames and is too short for subsampling " + + f"(it needs more than {limit_size} frames), return empty results", + xs_pad.size(1), + limit_size, + ) + xs_pad, masks = self.embed(xs_pad, masks) + else: + xs_pad = self.embed(xs_pad) + + intermediate_outs = [] + if len(self.interctc_layer_idx) == 0: + xs_pad, masks = self.encoders(xs_pad, masks) + else: + for layer_idx, encoder_layer in enumerate(self.encoders): + xs_pad, masks = encoder_layer(xs_pad, masks) + + if layer_idx + 1 in self.interctc_layer_idx: + encoder_out = xs_pad + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + # intermediate outputs are also normalized + if self.normalize_before: + encoder_out = self.after_norm(encoder_out) + + intermediate_outs.append((layer_idx + 1, encoder_out)) + + if self.interctc_use_conditioning: + ctc_out = ctc.softmax(encoder_out) + + if isinstance(xs_pad, tuple): + x, pos_emb = xs_pad + x = x + self.conditioning_layer(ctc_out) + xs_pad = (x, pos_emb) + else: + xs_pad = xs_pad + self.conditioning_layer(ctc_out) + + if isinstance(xs_pad, tuple): + xs_pad = xs_pad[0] + if self.normalize_before: + xs_pad = self.after_norm(xs_pad) + + olens = masks.squeeze(1).sum(1) + if len(intermediate_outs) > 0: + return (xs_pad, intermediate_outs), olens, None + return xs_pad, olens, None + + +class CausalConvolution(torch.nn.Module): + """ConformerConvolution module definition. + Args: + channels: The number of channels. + kernel_size: Size of the convolving kernel. + activation: Type of activation function. + norm_args: Normalization module arguments. + causal: Whether to use causal convolution (set to True if streaming). + """ + + def __init__( + self, + channels: int, + kernel_size: int, + activation: torch.nn.Module = torch.nn.ReLU(), + norm_args: Dict = {}, + causal: bool = False, + ) -> None: + """Construct an ConformerConvolution object.""" + super().__init__() + + assert (kernel_size - 1) % 2 == 0 + + self.kernel_size = kernel_size + + self.pointwise_conv1 = torch.nn.Conv1d( + channels, + 2 * channels, + kernel_size=1, + stride=1, + padding=0, + ) + + if causal: + self.lorder = kernel_size - 1 + padding = 0 + else: + self.lorder = 0 + padding = (kernel_size - 1) // 2 + + self.depthwise_conv = torch.nn.Conv1d( + channels, + channels, + kernel_size, + stride=1, + padding=padding, + groups=channels, + ) + self.norm = torch.nn.BatchNorm1d(channels, **norm_args) + self.pointwise_conv2 = torch.nn.Conv1d( + channels, + channels, + kernel_size=1, + stride=1, + padding=0, + ) + + self.activation = activation + + def forward( + self, + x: torch.Tensor, + cache: Optional[torch.Tensor] = None, + right_context: int = 0, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute convolution module. + Args: + x: ConformerConvolution input sequences. (B, T, D_hidden) + cache: ConformerConvolution input cache. (1, conv_kernel, D_hidden) + right_context: Number of frames in right context. + Returns: + x: ConformerConvolution output sequences. (B, T, D_hidden) + cache: ConformerConvolution output cache. (1, conv_kernel, D_hidden) + """ + x = self.pointwise_conv1(x.transpose(1, 2)) + x = torch.nn.functional.glu(x, dim=1) + + if self.lorder > 0: + if cache is None: + x = torch.nn.functional.pad(x, (self.lorder, 0), "constant", 0.0) + else: + x = torch.cat([cache, x], dim=2) + + if right_context > 0: + cache = x[:, :, -(self.lorder + right_context) : -right_context] + else: + cache = x[:, :, -self.lorder :] + + x = self.depthwise_conv(x) + x = self.activation(self.norm(x)) + + x = self.pointwise_conv2(x).transpose(1, 2) + + return x, cache + + +class ChunkEncoderLayer(torch.nn.Module): + """Chunk Conformer module definition. + Args: + block_size: Input/output size. + self_att: Self-attention module instance. + feed_forward: Feed-forward module instance. + feed_forward_macaron: Feed-forward module instance for macaron network. + conv_mod: Convolution module instance. + norm_class: Normalization module class. + norm_args: Normalization module arguments. + dropout_rate: Dropout rate. + """ + + def __init__( + self, + block_size: int, + self_att: torch.nn.Module, + feed_forward: torch.nn.Module, + feed_forward_macaron: torch.nn.Module, + conv_mod: torch.nn.Module, + norm_class: torch.nn.Module = LayerNorm, + norm_args: Dict = {}, + dropout_rate: float = 0.0, + ) -> None: + """Construct a Conformer object.""" + super().__init__() + + self.self_att = self_att + + self.feed_forward = feed_forward + self.feed_forward_macaron = feed_forward_macaron + self.feed_forward_scale = 0.5 + + self.conv_mod = conv_mod + + self.norm_feed_forward = norm_class(block_size, **norm_args) + self.norm_self_att = norm_class(block_size, **norm_args) + + self.norm_macaron = norm_class(block_size, **norm_args) + self.norm_conv = norm_class(block_size, **norm_args) + self.norm_final = norm_class(block_size, **norm_args) + + self.dropout = torch.nn.Dropout(dropout_rate) + + self.block_size = block_size + self.cache = None + + def reset_streaming_cache(self, left_context: int, device: torch.device) -> None: + """Initialize/Reset self-attention and convolution modules cache for streaming. + Args: + left_context: Number of left frames during chunk-by-chunk inference. + device: Device to use for cache tensor. + """ + self.cache = [ + torch.zeros( + (1, left_context, self.block_size), + device=device, + ), + torch.zeros( + ( + 1, + self.block_size, + self.conv_mod.kernel_size - 1, + ), + device=device, + ), + ] + + def forward( + self, + x: torch.Tensor, + pos_enc: torch.Tensor, + mask: torch.Tensor, + chunk_mask: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Encode input sequences. + Args: + x: Conformer input sequences. (B, T, D_block) + pos_enc: Positional embedding sequences. (B, 2 * (T - 1), D_block) + mask: Source mask. (B, T) + chunk_mask: Chunk mask. (T_2, T_2) + Returns: + x: Conformer output sequences. (B, T, D_block) + mask: Source mask. (B, T) + pos_enc: Positional embedding sequences. (B, 2 * (T - 1), D_block) + """ + residual = x + + x = self.norm_macaron(x) + x = residual + self.feed_forward_scale * self.dropout( + self.feed_forward_macaron(x) + ) + + residual = x + x = self.norm_self_att(x) + x_q = x + x = residual + self.dropout( + self.self_att( + x_q, + x, + x, + pos_enc, + mask, + chunk_mask=chunk_mask, + ) + ) + + residual = x + + x = self.norm_conv(x) + x, _ = self.conv_mod(x) + x = residual + self.dropout(x) + residual = x + + x = self.norm_feed_forward(x) + x = residual + self.feed_forward_scale * self.dropout(self.feed_forward(x)) + + x = self.norm_final(x) + return x, mask, pos_enc + + def chunk_forward( + self, + x: torch.Tensor, + pos_enc: torch.Tensor, + mask: torch.Tensor, + chunk_size: int = 16, + left_context: int = 0, + right_context: int = 0, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Encode chunk of input sequence. + Args: + x: Conformer input sequences. (B, T, D_block) + pos_enc: Positional embedding sequences. (B, 2 * (T - 1), D_block) + mask: Source mask. (B, T_2) + left_context: Number of frames in left context. + right_context: Number of frames in right context. + Returns: + x: Conformer output sequences. (B, T, D_block) + pos_enc: Positional embedding sequences. (B, 2 * (T - 1), D_block) + """ + residual = x + + x = self.norm_macaron(x) + x = residual + self.feed_forward_scale * self.feed_forward_macaron(x) + + residual = x + x = self.norm_self_att(x) + if left_context > 0: + key = torch.cat([self.cache[0], x], dim=1) + else: + key = x + val = key + + if right_context > 0: + att_cache = key[:, -(left_context + right_context) : -right_context, :] + else: + att_cache = key[:, -left_context:, :] + x = residual + self.self_att( + x, + key, + val, + pos_enc, + mask, + left_context=left_context, + ) + + residual = x + x = self.norm_conv(x) + x, conv_cache = self.conv_mod( + x, cache=self.cache[1], right_context=right_context + ) + x = residual + x + residual = x + + x = self.norm_feed_forward(x) + x = residual + self.feed_forward_scale * self.feed_forward(x) + + x = self.norm_final(x) + self.cache = [att_cache, conv_cache] + + return x, pos_enc + + +@tables.register("encoder_classes", "ChunkConformerEncoder") +class ConformerChunkEncoder(torch.nn.Module): + """Encoder module definition. + Args: + input_size: Input size. + body_conf: Encoder body configuration. + input_conf: Encoder input configuration. + main_conf: Encoder main configuration. + """ + + def __init__( + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.0, + embed_vgg_like: bool = False, + normalize_before: bool = True, + concat_after: bool = False, + positionwise_layer_type: str = "linear", + positionwise_conv_kernel_size: int = 3, + macaron_style: bool = False, + rel_pos_type: str = "legacy", + pos_enc_layer_type: str = "rel_pos", + selfattention_layer_type: str = "rel_selfattn", + activation_type: str = "swish", + use_cnn_module: bool = True, + zero_triu: bool = False, + norm_type: str = "layer_norm", + cnn_module_kernel: int = 31, + conv_mod_norm_eps: float = 0.00001, + conv_mod_norm_momentum: float = 0.1, + simplified_att_score: bool = False, + dynamic_chunk_training: bool = False, + short_chunk_threshold: float = 0.75, + short_chunk_size: int = 25, + left_chunk_size: int = 0, + time_reduction_factor: int = 1, + unified_model_training: bool = False, + default_chunk_size: int = 16, + jitter_range: int = 4, + subsampling_factor: int = 1, + ) -> None: + """Construct an Encoder object.""" + super().__init__() + + self.embed = StreamingConvInput( + input_size=input_size, + conv_size=output_size, + subsampling_factor=subsampling_factor, + vgg_like=embed_vgg_like, + output_size=output_size, + ) + + self.pos_enc = StreamingRelPositionalEncoding( + output_size, + positional_dropout_rate, + ) + + activation = get_activation(activation_type) + + pos_wise_args = ( + output_size, + linear_units, + positional_dropout_rate, + activation, + ) + + conv_mod_norm_args = { + "eps": conv_mod_norm_eps, + "momentum": conv_mod_norm_momentum, + } + + conv_mod_args = ( + output_size, + cnn_module_kernel, + activation, + conv_mod_norm_args, + dynamic_chunk_training or unified_model_training, + ) + + mult_att_args = ( + attention_heads, + output_size, + attention_dropout_rate, + simplified_att_score, + ) + + fn_modules = [] + for _ in range(num_blocks): + module = lambda: ChunkEncoderLayer( + output_size, + RelPositionMultiHeadedAttentionChunk(*mult_att_args), + PositionwiseFeedForward(*pos_wise_args), + PositionwiseFeedForward(*pos_wise_args), + CausalConvolution(*conv_mod_args), + dropout_rate=dropout_rate, + ) + fn_modules.append(module) + + self.encoders = MultiBlocks( + [fn() for fn in fn_modules], + output_size, + ) + + self._output_size = output_size + + self.dynamic_chunk_training = dynamic_chunk_training + self.short_chunk_threshold = short_chunk_threshold + self.short_chunk_size = short_chunk_size + self.left_chunk_size = left_chunk_size + + self.unified_model_training = unified_model_training + self.default_chunk_size = default_chunk_size + self.jitter_range = jitter_range + + self.time_reduction_factor = time_reduction_factor + + def output_size(self) -> int: + return self._output_size + + def get_encoder_input_raw_size(self, size: int, hop_length: int) -> int: + """Return the corresponding number of sample for a given chunk size, in frames. + Where size is the number of features frames after applying subsampling. + Args: + size: Number of frames after subsampling. + hop_length: Frontend's hop length + Returns: + : Number of raw samples + """ + return self.embed.get_size_before_subsampling(size) * hop_length + + def get_encoder_input_size(self, size: int) -> int: + """Return the corresponding number of sample for a given chunk size, in frames. + Where size is the number of features frames after applying subsampling. + Args: + size: Number of frames after subsampling. + Returns: + : Number of raw samples + """ + return self.embed.get_size_before_subsampling(size) + + def reset_streaming_cache(self, left_context: int, device: torch.device) -> None: + """Initialize/Reset encoder streaming cache. + Args: + left_context: Number of frames in left context. + device: Device ID. + """ + return self.encoders.reset_streaming_cache(left_context, device) + + def forward( + self, + x: torch.Tensor, + x_len: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Encode input sequences. + Args: + x: Encoder input features. (B, T_in, F) + x_len: Encoder input features lengths. (B,) + Returns: + x: Encoder outputs. (B, T_out, D_enc) + x_len: Encoder outputs lenghts. (B,) + """ + short_status, limit_size = check_short_utt( + self.embed.subsampling_factor, x.size(1) + ) + + if short_status: + raise TooShortUttError( + f"has {x.size(1)} frames and is too short for subsampling " + + f"(it needs more than {limit_size} frames), return empty results", + x.size(1), + limit_size, + ) + + mask = make_source_mask(x_len).to(x.device) + + if self.unified_model_training: + if self.training: + chunk_size = ( + self.default_chunk_size + + torch.randint( + -self.jitter_range, self.jitter_range + 1, (1,) + ).item() + ) + else: + chunk_size = self.default_chunk_size + x, mask = self.embed(x, mask, chunk_size) + pos_enc = self.pos_enc(x) + chunk_mask = make_chunk_mask( + x.size(1), + chunk_size, + left_chunk_size=self.left_chunk_size, + device=x.device, + ) + x_utt = self.encoders( + x, + pos_enc, + mask, + chunk_mask=None, + ) + x_chunk = self.encoders( + x, + pos_enc, + mask, + chunk_mask=chunk_mask, + ) + + olens = mask.eq(0).sum(1) + if self.time_reduction_factor > 1: + x_utt = x_utt[:, :: self.time_reduction_factor, :] + x_chunk = x_chunk[:, :: self.time_reduction_factor, :] + olens = torch.floor_divide(olens - 1, self.time_reduction_factor) + 1 + + return x_utt, x_chunk, olens + + elif self.dynamic_chunk_training: + max_len = x.size(1) + if self.training: + chunk_size = torch.randint(1, max_len, (1,)).item() + + if chunk_size > (max_len * self.short_chunk_threshold): + chunk_size = max_len + else: + chunk_size = (chunk_size % self.short_chunk_size) + 1 + else: + chunk_size = self.default_chunk_size + + x, mask = self.embed(x, mask, chunk_size) + pos_enc = self.pos_enc(x) + + chunk_mask = make_chunk_mask( + x.size(1), + chunk_size, + left_chunk_size=self.left_chunk_size, + device=x.device, + ) + else: + x, mask = self.embed(x, mask, None) + pos_enc = self.pos_enc(x) + chunk_mask = None + x = self.encoders( + x, + pos_enc, + mask, + chunk_mask=chunk_mask, + ) + + olens = mask.eq(0).sum(1) + if self.time_reduction_factor > 1: + x = x[:, :: self.time_reduction_factor, :] + olens = torch.floor_divide(olens - 1, self.time_reduction_factor) + 1 + + return x, olens, None + + def full_utt_forward( + self, + x: torch.Tensor, + x_len: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Encode input sequences. + Args: + x: Encoder input features. (B, T_in, F) + x_len: Encoder input features lengths. (B,) + Returns: + x: Encoder outputs. (B, T_out, D_enc) + x_len: Encoder outputs lenghts. (B,) + """ + short_status, limit_size = check_short_utt( + self.embed.subsampling_factor, x.size(1) + ) + + if short_status: + raise TooShortUttError( + f"has {x.size(1)} frames and is too short for subsampling " + + f"(it needs more than {limit_size} frames), return empty results", + x.size(1), + limit_size, + ) + + mask = make_source_mask(x_len).to(x.device) + x, mask = self.embed(x, mask, None) + pos_enc = self.pos_enc(x) + x_utt = self.encoders( + x, + pos_enc, + mask, + chunk_mask=None, + ) + + if self.time_reduction_factor > 1: + x_utt = x_utt[:, :: self.time_reduction_factor, :] + return x_utt + + def simu_chunk_forward( + self, + x: torch.Tensor, + x_len: torch.Tensor, + chunk_size: int = 16, + left_context: int = 32, + right_context: int = 0, + ) -> torch.Tensor: + short_status, limit_size = check_short_utt( + self.embed.subsampling_factor, x.size(1) + ) + + if short_status: + raise TooShortUttError( + f"has {x.size(1)} frames and is too short for subsampling " + + f"(it needs more than {limit_size} frames), return empty results", + x.size(1), + limit_size, + ) + + mask = make_source_mask(x_len) + + x, mask = self.embed(x, mask, chunk_size) + pos_enc = self.pos_enc(x) + chunk_mask = make_chunk_mask( + x.size(1), + chunk_size, + left_chunk_size=self.left_chunk_size, + device=x.device, + ) + + x = self.encoders( + x, + pos_enc, + mask, + chunk_mask=chunk_mask, + ) + olens = mask.eq(0).sum(1) + if self.time_reduction_factor > 1: + x = x[:, :: self.time_reduction_factor, :] + + return x + + def chunk_forward( + self, + x: torch.Tensor, + x_len: torch.Tensor, + processed_frames: torch.tensor, + chunk_size: int = 16, + left_context: int = 32, + right_context: int = 0, + ) -> torch.Tensor: + """Encode input sequences as chunks. + Args: + x: Encoder input features. (1, T_in, F) + x_len: Encoder input features lengths. (1,) + processed_frames: Number of frames already seen. + left_context: Number of frames in left context. + right_context: Number of frames in right context. + Returns: + x: Encoder outputs. (B, T_out, D_enc) + """ + mask = make_source_mask(x_len) + x, mask = self.embed(x, mask, None) + + if left_context > 0: + processed_mask = ( + torch.arange(left_context, device=x.device) + .view(1, left_context) + .flip(1) + ) + processed_mask = processed_mask >= processed_frames + mask = torch.cat([processed_mask, mask], dim=1) + pos_enc = self.pos_enc(x, left_context=left_context) + x = self.encoders.chunk_forward( + x, + pos_enc, + mask, + chunk_size=chunk_size, + left_context=left_context, + right_context=right_context, + ) + + if right_context > 0: + x = x[:, 0:-right_context, :] + + if self.time_reduction_factor > 1: + x = x[:, :: self.time_reduction_factor, :] + return x diff --git a/almeval/models/stepaudio/funasr_detach/models/conformer/model.py b/almeval/models/stepaudio/funasr_detach/models/conformer/model.py new file mode 100644 index 0000000000000000000000000000000000000000..2240b568e11de8b3ec299f4b5ca0a35586491d51 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/conformer/model.py @@ -0,0 +1,19 @@ +import logging + +import torch + +from funasr_detach.models.transformer.model import Transformer +from funasr_detach.register import tables + + +@tables.register("model_classes", "Conformer") +class Conformer(Transformer): + """CTC-attention hybrid Encoder-Decoder model""" + + def __init__( + self, + *args, + **kwargs, + ): + + super().__init__(*args, **kwargs) diff --git a/almeval/models/stepaudio/funasr_detach/models/conformer/template.yaml b/almeval/models/stepaudio/funasr_detach/models/conformer/template.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f646acc9d99faafc59d9808fd5d05c64a254cd28 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/conformer/template.yaml @@ -0,0 +1,117 @@ +# This is an example that demonstrates how to configure a model file. +# You can modify the configuration according to your own requirements. + +# to print the register_table: +# from funasr.register import tables +# tables.print() + +# network architecture +model: Conformer +model_conf: + ctc_weight: 0.3 + lsm_weight: 0.1 # label smoothing option + length_normalized_loss: false + +# encoder +encoder: ConformerEncoder +encoder_conf: + output_size: 256 + attention_heads: 4 + linear_units: 2048 + num_blocks: 12 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + attention_dropout_rate: 0.0 + input_layer: conv2d + normalize_before: true + pos_enc_layer_type: rel_pos + selfattention_layer_type: rel_selfattn + activation_type: swish + macaron_style: true + use_cnn_module: true + cnn_module_kernel: 15 + +# decoder +decoder: TransformerDecoder +decoder_conf: + attention_heads: 4 + linear_units: 2048 + num_blocks: 6 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + self_attention_dropout_rate: 0.0 + src_attention_dropout_rate: 0.0 + + +# frontend related +frontend: WavFrontend +frontend_conf: + fs: 16000 + window: hamming + n_mels: 80 + frame_length: 25 + frame_shift: 10 + dither: 0.0 + lfr_m: 1 + lfr_n: 1 + +specaug: SpecAug +specaug_conf: + apply_time_warp: true + time_warp_window: 5 + time_warp_mode: bicubic + apply_freq_mask: true + freq_mask_width_range: + - 0 + - 30 + num_freq_mask: 2 + apply_time_mask: true + time_mask_width_range: + - 0 + - 40 + num_time_mask: 2 + +train_conf: + accum_grad: 1 + grad_clip: 5 + max_epoch: 150 + val_scheduler_criterion: + - valid + - acc + best_model_criterion: + - - valid + - acc + - max + keep_nbest_models: 10 + log_interval: 50 + +optim: adam +optim_conf: + lr: 0.0005 +scheduler: warmuplr +scheduler_conf: + warmup_steps: 30000 + +dataset: AudioDataset +dataset_conf: + index_ds: IndexDSJsonl + batch_sampler: DynamicBatchLocalShuffleSampler + batch_type: example # example or length + batch_size: 1 # if batch_type is example, batch_size is the numbers of samples; if length, batch_size is source_token_len+target_token_len; + max_token_length: 2048 # filter samples if source_token_len+target_token_len > max_token_length, + buffer_size: 500 + shuffle: True + num_workers: 0 + +tokenizer: CharTokenizer +tokenizer_conf: + unk_symbol: + split_with_space: true + + +ctc_conf: + dropout_rate: 0.0 + ctc_type: builtin + reduce: true + ignore_nan_grad: true +normalize: null diff --git a/almeval/models/stepaudio/funasr_detach/models/data2vec/__init__.py b/almeval/models/stepaudio/funasr_detach/models/data2vec/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/almeval/models/stepaudio/funasr_detach/models/data2vec/data2vec.py b/almeval/models/stepaudio/funasr_detach/models/data2vec/data2vec.py new file mode 100644 index 0000000000000000000000000000000000000000..3dff699030ca410e2f93b5c2314197af5836ccf4 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/data2vec/data2vec.py @@ -0,0 +1,160 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from contextlib import contextmanager +from distutils.version import LooseVersion +from typing import Dict +from typing import Optional +from typing import Tuple + +import torch +import torch.nn as nn + +# from funasr_detach.layers.abs_normalize import AbsNormalize +# from funasr_detach.models.base_model import FunASRModel +# from funasr_detach.models.encoder.abs_encoder import AbsEncoder +from funasr_detach.frontends.abs_frontend import AbsFrontend + +# from funasr_detach.models.preencoder.abs_preencoder import AbsPreEncoder +# from funasr_detach.models.specaug.abs_specaug import AbsSpecAug +from funasr_detach.train_utils.device_funcs import force_gatherable + +if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): + from torch.cuda.amp import autocast +else: + # Nothing to do if torch<1.6.0 + @contextmanager + def autocast(enabled=True): + yield + + +class Data2VecPretrainModel(nn.Module): + """Data2Vec Pretrain model""" + + def __init__( + self, + frontend=None, + specaug=None, + normalize=None, + encoder=None, + preencoder=None, + ): + + super().__init__() + + self.frontend = frontend + self.specaug = specaug + self.normalize = normalize + self.preencoder = preencoder + self.encoder = encoder + self.num_updates = 0 + + def forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: + """Frontend + Encoder + Calc loss + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + """ + # Check that batch_size is unified + assert speech.shape[0] == speech_lengths.shape[0], ( + speech.shape, + speech_lengths.shape, + ) + + self.encoder.set_num_updates(self.num_updates) + + # 1. Encoder + encoder_out = self.encode(speech, speech_lengths) + + losses = encoder_out["losses"] + loss = sum(losses.values()) + sample_size = encoder_out["sample_size"] + loss = loss.sum() / sample_size + + target_var = float(encoder_out["target_var"]) + pred_var = float(encoder_out["pred_var"]) + ema_decay = float(encoder_out["ema_decay"]) + + stats = dict( + loss=torch.clone(loss.detach()), + target_var=target_var, + pred_var=pred_var, + ema_decay=ema_decay, + ) + + loss, stats, weight = force_gatherable((loss, stats, sample_size), loss.device) + return loss, stats, weight + + def collect_feats( + self, speech: torch.Tensor, speech_lengths: torch.Tensor + ) -> Dict[str, torch.Tensor]: + feats, feats_lengths = self._extract_feats(speech, speech_lengths) + return {"feats": feats, "feats_lengths": feats_lengths} + + def encode( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + ): + """Frontend + Encoder. + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + """ + with autocast(False): + # 1. Extract feats + feats, feats_lengths = self._extract_feats(speech, speech_lengths) + + # 2. Data augmentation + if self.specaug is not None and self.training: + feats, feats_lengths = self.specaug(feats, feats_lengths) + + # 3. Normalization for feature: e.g. Global-CMVN, Utterance-CMVN + if self.normalize is not None: + feats, feats_lengths = self.normalize(feats, feats_lengths) + + # Pre-encoder, e.g. used for raw input data + if self.preencoder is not None: + feats, feats_lengths = self.preencoder(feats, feats_lengths) + + # 4. Forward encoder + if min(speech_lengths) == max( + speech_lengths + ): # for clipping, set speech_lengths as None + speech_lengths = None + encoder_out = self.encoder( + feats, speech_lengths, mask=True, features_only=False + ) + + return encoder_out + + def _extract_feats( + self, speech: torch.Tensor, speech_lengths: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + assert speech_lengths.dim() == 1, speech_lengths.shape + + # for data-parallel + speech = speech[:, : speech_lengths.max()] + + if self.frontend is not None: + # Frontend + # e.g. STFT and Feature extract + # data_loader may send time-domain signal in this case + # speech (Batch, NSamples) -> feats: (Batch, NFrames, Dim) + feats, feats_lengths = self.frontend(speech, speech_lengths) + else: + # No frontend and no feature extract + feats, feats_lengths = speech, speech_lengths + return feats, feats_lengths + + def set_num_updates(self, num_updates): + self.num_updates = num_updates + + def get_num_updates(self): + return self.num_updates diff --git a/almeval/models/stepaudio/funasr_detach/models/data2vec/data2vec_encoder.py b/almeval/models/stepaudio/funasr_detach/models/data2vec/data2vec_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..52b4a20ab19b6835c9ded54f369b6e3b808f0814 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/data2vec/data2vec_encoder.py @@ -0,0 +1,578 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import logging +import math + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F + +from funasr_detach.models.data2vec.data_utils import compute_mask_indices +from funasr_detach.models.data2vec.ema_module import EMAModule +from funasr_detach.models.data2vec.grad_multiply import GradMultiply +from funasr_detach.models.data2vec.wav2vec2 import ( + ConvFeatureExtractionModel, + TransformerEncoder, +) +from funasr_detach.models.transformer.utils.nets_utils import make_pad_mask + + +def get_annealed_rate(start, end, curr_step, total_steps): + r = end - start + pct_remaining = 1 - curr_step / total_steps + return end - r * pct_remaining + + +class Data2VecEncoder(nn.Module): + def __init__( + self, + # for ConvFeatureExtractionModel + input_size: int = None, + extractor_mode: str = None, + conv_feature_layers: str = "[(512,2,2)] + [(512,2,2)]", + # for Transformer Encoder + ## model architecture + layer_type: str = "transformer", + layer_norm_first: bool = False, + encoder_layers: int = 12, + encoder_embed_dim: int = 768, + encoder_ffn_embed_dim: int = 3072, + encoder_attention_heads: int = 12, + activation_fn: str = "gelu", + ## dropouts + dropout: float = 0.1, + attention_dropout: float = 0.1, + activation_dropout: float = 0.0, + encoder_layerdrop: float = 0.0, + dropout_input: float = 0.0, + dropout_features: float = 0.0, + ## grad settings + feature_grad_mult: float = 1.0, + ## masking + mask_prob: float = 0.65, + mask_length: int = 10, + mask_selection: str = "static", + mask_other: int = 0, + no_mask_overlap: bool = False, + mask_min_space: int = 1, + require_same_masks: bool = True, # if set as True, collate_fn should be clipping + mask_dropout: float = 0.0, + ## channel masking + mask_channel_length: int = 10, + mask_channel_prob: float = 0.0, + mask_channel_before: bool = False, + mask_channel_selection: str = "static", + mask_channel_other: int = 0, + no_mask_channel_overlap: bool = False, + mask_channel_min_space: int = 1, + ## positional embeddings + conv_pos: int = 128, + conv_pos_groups: int = 16, + pos_conv_depth: int = 1, + max_positions: int = 100000, + # EMA module + average_top_k_layers: int = 8, + layer_norm_target_layer: bool = False, + instance_norm_target_layer: bool = False, + instance_norm_targets: bool = False, + layer_norm_targets: bool = False, + batch_norm_target_layer: bool = False, + group_norm_target_layer: bool = False, + ema_decay: float = 0.999, + ema_end_decay: float = 0.9999, + ema_anneal_end_step: int = 100000, + ema_transformer_only: bool = True, + ema_layers_only: bool = True, + min_target_var: float = 0.1, + min_pred_var: float = 0.01, + # Loss + loss_beta: float = 0.0, + loss_scale: float = None, + # FP16 optimization + required_seq_len_multiple: int = 2, + ): + super().__init__() + + # ConvFeatureExtractionModel + self.conv_feature_layers = conv_feature_layers + feature_enc_layers = eval(conv_feature_layers) + self.extractor_embed = feature_enc_layers[-1][0] + self.feature_extractor = ConvFeatureExtractionModel( + conv_layers=feature_enc_layers, + dropout=0.0, + mode=extractor_mode, + in_d=input_size, + ) + + # Transformer Encoder + ## model architecture + self.layer_type = layer_type + self.layer_norm_first = layer_norm_first + self.encoder_layers = encoder_layers + self.encoder_embed_dim = encoder_embed_dim + self.encoder_ffn_embed_dim = encoder_ffn_embed_dim + self.encoder_attention_heads = encoder_attention_heads + self.activation_fn = activation_fn + ## dropout + self.dropout = dropout + self.attention_dropout = attention_dropout + self.activation_dropout = activation_dropout + self.encoder_layerdrop = encoder_layerdrop + self.dropout_input = dropout_input + self.dropout_features = dropout_features + ## grad settings + self.feature_grad_mult = feature_grad_mult + ## masking + self.mask_prob = mask_prob + self.mask_length = mask_length + self.mask_selection = mask_selection + self.mask_other = mask_other + self.no_mask_overlap = no_mask_overlap + self.mask_min_space = mask_min_space + self.require_same_masks = ( + require_same_masks # if set as True, collate_fn should be clipping + ) + self.mask_dropout = mask_dropout + ## channel masking + self.mask_channel_length = mask_channel_length + self.mask_channel_prob = mask_channel_prob + self.mask_channel_before = mask_channel_before + self.mask_channel_selection = mask_channel_selection + self.mask_channel_other = mask_channel_other + self.no_mask_channel_overlap = no_mask_channel_overlap + self.mask_channel_min_space = mask_channel_min_space + ## positional embeddings + self.conv_pos = conv_pos + self.conv_pos_groups = conv_pos_groups + self.pos_conv_depth = pos_conv_depth + self.max_positions = max_positions + self.mask_emb = nn.Parameter( + torch.FloatTensor(self.encoder_embed_dim).uniform_() + ) + self.encoder = TransformerEncoder( + dropout=self.dropout, + encoder_embed_dim=self.encoder_embed_dim, + required_seq_len_multiple=required_seq_len_multiple, + pos_conv_depth=self.pos_conv_depth, + conv_pos=self.conv_pos, + conv_pos_groups=self.conv_pos_groups, + # transformer layers + layer_type=self.layer_type, + encoder_layers=self.encoder_layers, + encoder_ffn_embed_dim=self.encoder_ffn_embed_dim, + encoder_attention_heads=self.encoder_attention_heads, + attention_dropout=self.attention_dropout, + activation_dropout=self.activation_dropout, + activation_fn=self.activation_fn, + layer_norm_first=self.layer_norm_first, + encoder_layerdrop=self.encoder_layerdrop, + max_positions=self.max_positions, + ) + ## projections and dropouts + self.post_extract_proj = nn.Linear(self.extractor_embed, self.encoder_embed_dim) + self.dropout_input = nn.Dropout(self.dropout_input) + self.dropout_features = nn.Dropout(self.dropout_features) + self.layer_norm = torch.nn.LayerNorm(self.extractor_embed) + self.final_proj = nn.Linear(self.encoder_embed_dim, self.encoder_embed_dim) + + # EMA module + self.average_top_k_layers = average_top_k_layers + self.layer_norm_target_layer = layer_norm_target_layer + self.instance_norm_target_layer = instance_norm_target_layer + self.instance_norm_targets = instance_norm_targets + self.layer_norm_targets = layer_norm_targets + self.batch_norm_target_layer = batch_norm_target_layer + self.group_norm_target_layer = group_norm_target_layer + self.ema_decay = ema_decay + self.ema_end_decay = ema_end_decay + self.ema_anneal_end_step = ema_anneal_end_step + self.ema_transformer_only = ema_transformer_only + self.ema_layers_only = ema_layers_only + self.min_target_var = min_target_var + self.min_pred_var = min_pred_var + self.ema = None + + # Loss + self.loss_beta = loss_beta + self.loss_scale = loss_scale + + # FP16 optimization + self.required_seq_len_multiple = required_seq_len_multiple + + self.num_updates = 0 + + logging.info("Data2VecEncoder settings: {}".format(self.__dict__)) + + def make_ema_teacher(self): + skip_keys = set() + if self.ema_layers_only: + self.ema_transformer_only = True + for k, _ in self.encoder.pos_conv.named_parameters(): + skip_keys.add(f"pos_conv.{k}") + + self.ema = EMAModule( + self.encoder if self.ema_transformer_only else self, + ema_decay=self.ema_decay, + ema_fp32=True, + skip_keys=skip_keys, + ) + + def set_num_updates(self, num_updates): + if self.ema is None and self.final_proj is not None: + logging.info("Making EMA Teacher") + self.make_ema_teacher() + elif self.training and self.ema is not None: + if self.ema_decay != self.ema_end_decay: + if num_updates >= self.ema_anneal_end_step: + decay = self.ema_end_decay + else: + decay = get_annealed_rate( + self.ema_decay, + self.ema_end_decay, + num_updates, + self.ema_anneal_end_step, + ) + self.ema.set_decay(decay) + if self.ema.get_decay() < 1: + self.ema.step(self.encoder if self.ema_transformer_only else self) + + self.num_updates = num_updates + + def apply_mask( + self, + x, + padding_mask, + mask_indices=None, + mask_channel_indices=None, + ): + B, T, C = x.shape + + if self.mask_channel_prob > 0 and self.mask_channel_before: + mask_channel_indices = compute_mask_indices( + (B, C), + None, + self.mask_channel_prob, + self.mask_channel_length, + self.mask_channel_selection, + self.mask_channel_other, + no_overlap=self.no_mask_channel_overlap, + min_space=self.mask_channel_min_space, + ) + mask_channel_indices = ( + torch.from_numpy(mask_channel_indices) + .to(x.device) + .unsqueeze(1) + .expand(-1, T, -1) + ) + x[mask_channel_indices] = 0 + + if self.mask_prob > 0: + if mask_indices is None: + mask_indices = compute_mask_indices( + (B, T), + padding_mask, + self.mask_prob, + self.mask_length, + self.mask_selection, + self.mask_other, + min_masks=1, + no_overlap=self.no_mask_overlap, + min_space=self.mask_min_space, + require_same_masks=self.require_same_masks, + mask_dropout=self.mask_dropout, + ) + mask_indices = torch.from_numpy(mask_indices).to(x.device) + x[mask_indices] = self.mask_emb + else: + mask_indices = None + + if self.mask_channel_prob > 0 and not self.mask_channel_before: + if mask_channel_indices is None: + mask_channel_indices = compute_mask_indices( + (B, C), + None, + self.mask_channel_prob, + self.mask_channel_length, + self.mask_channel_selection, + self.mask_channel_other, + no_overlap=self.no_mask_channel_overlap, + min_space=self.mask_channel_min_space, + ) + mask_channel_indices = ( + torch.from_numpy(mask_channel_indices) + .to(x.device) + .unsqueeze(1) + .expand(-1, T, -1) + ) + x[mask_channel_indices] = 0 + + return x, mask_indices + + def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor): + """ + Computes the output length of the convolutional layers + """ + + def _conv_out_length(input_length, kernel_size, stride): + return torch.floor( + (input_length - kernel_size).to(torch.float32) / stride + 1 + ) + + conv_cfg_list = eval(self.conv_feature_layers) + + for i in range(len(conv_cfg_list)): + input_lengths = _conv_out_length( + input_lengths, conv_cfg_list[i][1], conv_cfg_list[i][2] + ) + + return input_lengths.to(torch.long) + + def forward( + self, + xs_pad, + ilens=None, + mask=False, + features_only=True, + layer=None, + mask_indices=None, + mask_channel_indices=None, + padding_count=None, + ): + # create padding_mask by ilens + if ilens is not None: + padding_mask = make_pad_mask(lengths=ilens).to(xs_pad.device) + else: + padding_mask = None + + features = xs_pad + + if self.feature_grad_mult > 0: + features = self.feature_extractor(features) + if self.feature_grad_mult != 1.0: + features = GradMultiply.apply(features, self.feature_grad_mult) + else: + with torch.no_grad(): + features = self.feature_extractor(features) + + features = features.transpose(1, 2) + + features = self.layer_norm(features) + + orig_padding_mask = padding_mask + + if padding_mask is not None: + input_lengths = (1 - padding_mask.long()).sum(-1) + # apply conv formula to get real output_lengths + output_lengths = self._get_feat_extract_output_lengths(input_lengths) + + padding_mask = torch.zeros( + features.shape[:2], dtype=features.dtype, device=features.device + ) + # these two operations makes sure that all values + # before the output lengths indices are attended to + padding_mask[ + ( + torch.arange(padding_mask.shape[0], device=padding_mask.device), + output_lengths - 1, + ) + ] = 1 + padding_mask = (1 - padding_mask.flip([-1]).cumsum(-1).flip([-1])).bool() + else: + padding_mask = None + + if self.post_extract_proj is not None: + features = self.post_extract_proj(features) + + pre_encoder_features = None + if self.ema_transformer_only: + pre_encoder_features = features.clone() + + features = self.dropout_input(features) + + if mask: + x, mask_indices = self.apply_mask( + features, + padding_mask, + mask_indices=mask_indices, + mask_channel_indices=mask_channel_indices, + ) + else: + x = features + mask_indices = None + + x, layer_results = self.encoder( + x, + padding_mask=padding_mask, + layer=layer, + ) + + if features_only: + encoder_out_lens = (1 - padding_mask.long()).sum(1) + return x, encoder_out_lens, None + + result = { + "losses": {}, + "padding_mask": padding_mask, + "x": x, + } + + with torch.no_grad(): + self.ema.model.eval() + + if self.ema_transformer_only: + y, layer_results = self.ema.model.extract_features( + pre_encoder_features, + padding_mask=padding_mask, + min_layer=self.encoder_layers - self.average_top_k_layers, + ) + y = { + "x": y, + "padding_mask": padding_mask, + "layer_results": layer_results, + } + else: + y = self.ema.model.extract_features( + source=xs_pad, + padding_mask=orig_padding_mask, + mask=False, + ) + + target_layer_results = [l[2] for l in y["layer_results"]] + + permuted = False + if self.instance_norm_target_layer or self.batch_norm_target_layer: + target_layer_results = [ + tl.permute(1, 2, 0) for tl in target_layer_results # TBC -> BCT + ] + permuted = True + + if self.batch_norm_target_layer: + target_layer_results = [ + F.batch_norm( + tl.float(), running_mean=None, running_var=None, training=True + ) + for tl in target_layer_results + ] + + if self.instance_norm_target_layer: + target_layer_results = [ + F.instance_norm(tl.float()) for tl in target_layer_results + ] + + if permuted: + target_layer_results = [ + tl.transpose(1, 2) for tl in target_layer_results # BCT -> BTC + ] + + if self.group_norm_target_layer: + target_layer_results = [ + F.layer_norm(tl.float(), tl.shape[-2:]) + for tl in target_layer_results + ] + + if self.layer_norm_target_layer: + target_layer_results = [ + F.layer_norm(tl.float(), tl.shape[-1:]) + for tl in target_layer_results + ] + + y = sum(target_layer_results) / len(target_layer_results) + + if self.layer_norm_targets: + y = F.layer_norm(y.float(), y.shape[-1:]) + + if self.instance_norm_targets: + y = F.instance_norm(y.float().transpose(1, 2)).transpose(1, 2) + + if not permuted: + y = y.transpose(0, 1) + + y = y[mask_indices] + + x = x[mask_indices] + x = self.final_proj(x) + + sz = x.size(-1) + + if self.loss_beta == 0: + loss = F.mse_loss(x.float(), y.float(), reduction="none").sum(dim=-1) + else: + loss = F.smooth_l1_loss( + x.float(), y.float(), reduction="none", beta=self.loss_beta + ).sum(dim=-1) + + if self.loss_scale is not None: + scale = self.loss_scale + else: + scale = 1 / math.sqrt(sz) + + result["losses"]["regression"] = loss.sum() * scale + + if "sample_size" not in result: + result["sample_size"] = loss.numel() + + with torch.no_grad(): + result["target_var"] = self.compute_var(y) + result["pred_var"] = self.compute_var(x.float()) + + if self.num_updates > 5000 and result["target_var"] < self.min_target_var: + logging.error( + f"target var is {result['target_var'].item()} < {self.min_target_var}, exiting" + ) + raise Exception( + f"target var is {result['target_var'].item()} < {self.min_target_var}, exiting" + ) + if self.num_updates > 5000 and result["pred_var"] < self.min_pred_var: + logging.error( + f"pred var is {result['pred_var'].item()} < {self.min_pred_var}, exiting" + ) + raise Exception( + f"pred var is {result['pred_var'].item()} < {self.min_pred_var}, exiting" + ) + + if self.ema is not None: + result["ema_decay"] = self.ema.get_decay() * 1000 + + return result + + @staticmethod + def compute_var(y): + y = y.view(-1, y.size(-1)) + if dist.is_initialized(): + zc = torch.tensor(y.size(0)).cuda() + zs = y.sum(dim=0) + zss = (y**2).sum(dim=0) + + dist.all_reduce(zc) + dist.all_reduce(zs) + dist.all_reduce(zss) + + var = zss / (zc - 1) - (zs**2) / (zc * (zc - 1)) + return torch.sqrt(var + 1e-6).mean() + else: + return torch.sqrt(y.var(dim=0) + 1e-6).mean() + + def extract_features(self, xs_pad, ilens, mask=False, layer=None): + res = self.forward( + xs_pad, + ilens, + mask=mask, + features_only=True, + layer=layer, + ) + return res + + def remove_pretraining_modules(self, last_layer=None): + self.final_proj = None + self.ema = None + if last_layer is not None: + self.encoder.layers = nn.ModuleList( + l for i, l in enumerate(self.encoder.layers) if i <= last_layer + ) + + def output_size(self) -> int: + return self.encoder_embed_dim diff --git a/almeval/models/stepaudio/funasr_detach/models/data2vec/data_utils.py b/almeval/models/stepaudio/funasr_detach/models/data2vec/data_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..69c0bbcebedd6c193a31bb7c6f1f32419f2d6fbf --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/data2vec/data_utils.py @@ -0,0 +1,147 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + + +from typing import Optional, Tuple + +import numpy as np +import torch + + +def compute_mask_indices( + shape: Tuple[int, int], + padding_mask: Optional[torch.Tensor], + mask_prob: float, + mask_length: int, + mask_type: str = "static", + mask_other: float = 0.0, + min_masks: int = 0, + no_overlap: bool = False, + min_space: int = 0, + require_same_masks: bool = True, + mask_dropout: float = 0.0, +) -> np.ndarray: + """ + Computes random mask spans for a given shape + + Args: + shape: the the shape for which to compute masks. + should be of size 2 where first element is batch size and 2nd is timesteps + padding_mask: optional padding mask of the same size as shape, which will prevent masking padded elements + mask_prob: probability for each token to be chosen as start of the span to be masked. this will be multiplied by + number of timesteps divided by length of mask span to mask approximately this percentage of all elements. + however due to overlaps, the actual number will be smaller (unless no_overlap is True) + mask_type: how to compute mask lengths + static = fixed size + uniform = sample from uniform distribution [mask_other, mask_length*2] + normal = sample from normal distribution with mean mask_length and stdev mask_other. mask is min 1 element + poisson = sample from possion distribution with lambda = mask length + min_masks: minimum number of masked spans + no_overlap: if false, will switch to an alternative recursive algorithm that prevents spans from overlapping + min_space: only used if no_overlap is True, this is how many elements to keep unmasked between spans + require_same_masks: if true, will randomly drop out masks until same amount of masks remains in each sample + mask_dropout: randomly dropout this percentage of masks in each example + """ + + bsz, all_sz = shape + mask = np.full((bsz, all_sz), False) + + all_num_mask = int( + # add a random number for probabilistic rounding + mask_prob * all_sz / float(mask_length) + + np.random.rand() + ) + + all_num_mask = max(min_masks, all_num_mask) + + mask_idcs = [] + for i in range(bsz): + if padding_mask is not None: + sz = all_sz - padding_mask[i].long().sum().item() + num_mask = int( + # add a random number for probabilistic rounding + mask_prob * sz / float(mask_length) + + np.random.rand() + ) + num_mask = max(min_masks, num_mask) + else: + sz = all_sz + num_mask = all_num_mask + + if mask_type == "static": + lengths = np.full(num_mask, mask_length) + elif mask_type == "uniform": + lengths = np.random.randint(mask_other, mask_length * 2 + 1, size=num_mask) + elif mask_type == "normal": + lengths = np.random.normal(mask_length, mask_other, size=num_mask) + lengths = [max(1, int(round(x))) for x in lengths] + elif mask_type == "poisson": + lengths = np.random.poisson(mask_length, size=num_mask) + lengths = [int(round(x)) for x in lengths] + else: + raise Exception("unknown mask selection " + mask_type) + + if sum(lengths) == 0: + lengths[0] = min(mask_length, sz - 1) + + if no_overlap: + mask_idc = [] + + def arrange(s, e, length, keep_length): + span_start = np.random.randint(s, e - length) + mask_idc.extend(span_start + i for i in range(length)) + + new_parts = [] + if span_start - s - min_space >= keep_length: + new_parts.append((s, span_start - min_space + 1)) + if e - span_start - length - min_space > keep_length: + new_parts.append((span_start + length + min_space, e)) + return new_parts + + parts = [(0, sz)] + min_length = min(lengths) + for length in sorted(lengths, reverse=True): + lens = np.fromiter( + (e - s if e - s >= length + min_space else 0 for s, e in parts), + np.int32, + ) + l_sum = np.sum(lens) + if l_sum == 0: + break + probs = lens / np.sum(lens) + c = np.random.choice(len(parts), p=probs) + s, e = parts.pop(c) + parts.extend(arrange(s, e, length, min_length)) + mask_idc = np.asarray(mask_idc) + else: + min_len = min(lengths) + if sz - min_len <= num_mask: + min_len = sz - num_mask - 1 + + mask_idc = np.random.choice(sz - min_len, num_mask, replace=False) + + mask_idc = np.asarray( + [ + mask_idc[j] + offset + for j in range(len(mask_idc)) + for offset in range(lengths[j]) + ] + ) + + mask_idcs.append(np.unique(mask_idc[mask_idc < sz])) + + min_len = min([len(m) for m in mask_idcs]) + for i, mask_idc in enumerate(mask_idcs): + if len(mask_idc) > min_len and require_same_masks: + mask_idc = np.random.choice(mask_idc, min_len, replace=False) + if mask_dropout > 0: + num_holes = np.rint(len(mask_idc) * mask_dropout).astype(int) + mask_idc = np.random.choice( + mask_idc, len(mask_idc) - num_holes, replace=False + ) + + mask[i, mask_idc] = True + + return mask diff --git a/almeval/models/stepaudio/funasr_detach/models/data2vec/ema_module.py b/almeval/models/stepaudio/funasr_detach/models/data2vec/ema_module.py new file mode 100644 index 0000000000000000000000000000000000000000..a98da5daf758ecdd578f75d27c7b2c39c5c4a257 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/data2vec/ema_module.py @@ -0,0 +1,134 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +""" +Used for EMA tracking a given pytorch module. The user is responsible for calling step() +and setting the appropriate decay +""" + +import copy +import logging + +import torch + + +class EMAModule: + """Exponential Moving Average of Fairseq Models""" + + def __init__( + self, model, ema_decay=0.9999, ema_fp32=False, device=None, skip_keys=None + ): + """ + @param model model to initialize the EMA with + @param config EMAConfig object with configuration like + ema_decay, ema_update_freq, ema_fp32 + @param device If provided, copy EMA to this device (e.g. gpu). + Otherwise EMA is in the same device as the model. + """ + + self.decay = ema_decay + self.ema_fp32 = ema_fp32 + self.model = copy.deepcopy(model) + self.model.requires_grad_(False) + self.skip_keys = skip_keys or set() + self.fp32_params = {} + + if device is not None: + logging.info(f"Copying EMA model to device {device}") + self.model = self.model.to(device=device) + + if self.ema_fp32: + self.build_fp32_params() + + self.update_freq_counter = 0 + + def build_fp32_params(self, state_dict=None): + """ + Store a copy of the EMA params in fp32. + If state dict is passed, the EMA params is copied from + the provided state dict. Otherwise, it is copied from the + current EMA model parameters. + """ + if not self.ema_fp32: + raise RuntimeError( + "build_fp32_params should not be called if ema_fp32=False. " + "Use ema_fp32=True if this is really intended." + ) + + if state_dict is None: + state_dict = self.model.state_dict() + + def _to_float(t): + return t.float() if torch.is_floating_point(t) else t + + for param_key in state_dict: + if param_key in self.fp32_params: + self.fp32_params[param_key].copy_(state_dict[param_key]) + else: + self.fp32_params[param_key] = _to_float(state_dict[param_key]) + + def restore(self, state_dict, build_fp32_params=False): + """Load data from a model spec into EMA model""" + self.model.load_state_dict(state_dict, strict=False) + if build_fp32_params: + self.build_fp32_params(state_dict) + + def set_decay(self, decay): + self.decay = decay + + def get_decay(self): + return self.decay + + def _step_internal(self, new_model): + """One update of the EMA model based on new model weights""" + decay = self.decay + + ema_state_dict = {} + ema_params = self.fp32_params if self.ema_fp32 else self.model.state_dict() + for key, param in new_model.state_dict().items(): + if isinstance(param, dict): + continue + try: + ema_param = ema_params[key] + except KeyError: + ema_param = ( + param.float().clone() if param.ndim == 1 else copy.deepcopy(param) + ) + + if param.shape != ema_param.shape: + raise ValueError( + "incompatible tensor shapes between model param and ema param" + + "{} vs. {}".format(param.shape, ema_param.shape) + ) + + if "version" in key: + # Do not decay a model.version pytorch param + continue + + if key in self.skip_keys or ( + "num_batches_tracked" in key and ema_param.dtype == torch.int64 + ): + ema_param = param.to(dtype=ema_param.dtype).clone() + ema_params[key].copy_(ema_param) + else: + ema_param.mul_(decay) + ema_param.add_(param.to(dtype=ema_param.dtype), alpha=1 - decay) + ema_state_dict[key] = ema_param + self.restore(ema_state_dict, build_fp32_params=False) + + def step(self, new_model): + self._step_internal(new_model) + + def reverse(self, model): + """ + Load the model parameters from EMA model. + Useful for inference or fine-tuning from the EMA model. + """ + d = self.model.state_dict() + if "_ema" in d: + del d["_ema"] + + model.load_state_dict(d, strict=False) + return model diff --git a/almeval/models/stepaudio/funasr_detach/models/data2vec/grad_multiply.py b/almeval/models/stepaudio/funasr_detach/models/data2vec/grad_multiply.py new file mode 100644 index 0000000000000000000000000000000000000000..08d15f55dfda9c61a1cf8641ea31424fe1d97f57 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/data2vec/grad_multiply.py @@ -0,0 +1,18 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch + + +class GradMultiply(torch.autograd.Function): + @staticmethod + def forward(ctx, x, scale): + ctx.scale = scale + res = x.new(x) + return res + + @staticmethod + def backward(ctx, grad): + return grad * ctx.scale, None diff --git a/almeval/models/stepaudio/funasr_detach/models/data2vec/multihead_attention.py b/almeval/models/stepaudio/funasr_detach/models/data2vec/multihead_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..be7d18aefb10e1a43ba7bc6b9c29442c9423d8ef --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/data2vec/multihead_attention.py @@ -0,0 +1,641 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import logging +import math +from typing import Dict, List, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import Tensor, nn +from torch.nn import Parameter + +from funasr_detach.models.data2vec.quant_noise import quant_noise + + +class FairseqDropout(nn.Module): + def __init__(self, p, module_name=None): + super().__init__() + self.p = p + self.module_name = module_name + self.apply_during_inference = False + + def forward(self, x, inplace: bool = False): + if self.p > 0 and (self.training or self.apply_during_inference): + return F.dropout(x, p=self.p, training=True, inplace=inplace) + else: + return x + + def make_generation_fast_( + self, + name: str, + retain_dropout: bool = False, + retain_dropout_modules: Optional[List[str]] = None, + **kwargs, + ): + if retain_dropout: + if retain_dropout_modules is not None and self.module_name is None: + logging.warning( + "Cannot enable dropout during inference for module {} " + "because module_name was not set".format(name) + ) + elif ( + retain_dropout_modules is None # if None, apply to all modules + or self.module_name in retain_dropout_modules + ): + logging.info( + "Enabling dropout during inference for module: {}".format(name) + ) + self.apply_during_inference = True + else: + logging.info("Disabling dropout for module: {}".format(name)) + + +class MultiheadAttention(nn.Module): + """Multi-headed attention. + + See "Attention Is All You Need" for more details. + """ + + def __init__( + self, + embed_dim, + num_heads, + kdim=None, + vdim=None, + dropout=0.0, + bias=True, + add_bias_kv=False, + add_zero_attn=False, + self_attention=False, + encoder_decoder_attention=False, + q_noise=0.0, + qn_block_size=8, + ): + super().__init__() + self.embed_dim = embed_dim + self.kdim = kdim if kdim is not None else embed_dim + self.vdim = vdim if vdim is not None else embed_dim + self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim + + self.num_heads = num_heads + self.dropout_module = FairseqDropout( + dropout, module_name=self.__class__.__name__ + ) + + self.head_dim = embed_dim // num_heads + assert ( + self.head_dim * num_heads == self.embed_dim + ), "embed_dim must be divisible by num_heads" + self.scaling = self.head_dim**-0.5 + + self.self_attention = self_attention + self.encoder_decoder_attention = encoder_decoder_attention + + assert not self.self_attention or self.qkv_same_dim, ( + "Self-attention requires query, key and " "value to be of the same size" + ) + + self.k_proj = quant_noise( + nn.Linear(self.kdim, embed_dim, bias=bias), q_noise, qn_block_size + ) + self.v_proj = quant_noise( + nn.Linear(self.vdim, embed_dim, bias=bias), q_noise, qn_block_size + ) + self.q_proj = quant_noise( + nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size + ) + + self.out_proj = quant_noise( + nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size + ) + + if add_bias_kv: + self.bias_k = Parameter(torch.Tensor(1, 1, embed_dim)) + self.bias_v = Parameter(torch.Tensor(1, 1, embed_dim)) + else: + self.bias_k = self.bias_v = None + + self.add_zero_attn = add_zero_attn + + self.reset_parameters() + + self.onnx_trace = False + self.skip_embed_dim_check = False + + def prepare_for_onnx_export_(self): + self.onnx_trace = True + + def reset_parameters(self): + if self.qkv_same_dim: + # Empirically observed the convergence to be much better with + # the scaled initialization + nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2)) + nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2)) + nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2)) + else: + nn.init.xavier_uniform_(self.k_proj.weight) + nn.init.xavier_uniform_(self.v_proj.weight) + nn.init.xavier_uniform_(self.q_proj.weight) + + nn.init.xavier_uniform_(self.out_proj.weight) + if self.out_proj.bias is not None: + nn.init.constant_(self.out_proj.bias, 0.0) + if self.bias_k is not None: + nn.init.xavier_normal_(self.bias_k) + if self.bias_v is not None: + nn.init.xavier_normal_(self.bias_v) + + def _get_reserve_head_index(self, num_heads_to_keep: int): + k_proj_heads_norm = [] + q_proj_heads_norm = [] + v_proj_heads_norm = [] + + for i in range(self.num_heads): + start_idx = i * self.head_dim + end_idx = (i + 1) * self.head_dim + k_proj_heads_norm.append( + torch.sum(torch.abs(self.k_proj.weight[start_idx:end_idx,])).tolist() + + torch.sum(torch.abs(self.k_proj.bias[start_idx:end_idx])).tolist() + ) + q_proj_heads_norm.append( + torch.sum(torch.abs(self.q_proj.weight[start_idx:end_idx,])).tolist() + + torch.sum(torch.abs(self.q_proj.bias[start_idx:end_idx])).tolist() + ) + v_proj_heads_norm.append( + torch.sum(torch.abs(self.v_proj.weight[start_idx:end_idx,])).tolist() + + torch.sum(torch.abs(self.v_proj.bias[start_idx:end_idx])).tolist() + ) + + heads_norm = [] + for i in range(self.num_heads): + heads_norm.append( + k_proj_heads_norm[i] + q_proj_heads_norm[i] + v_proj_heads_norm[i] + ) + + sorted_head_index = sorted( + range(self.num_heads), key=lambda k: heads_norm[k], reverse=True + ) + reserve_head_index = [] + for i in range(num_heads_to_keep): + start = sorted_head_index[i] * self.head_dim + end = (sorted_head_index[i] + 1) * self.head_dim + reserve_head_index.append((start, end)) + return reserve_head_index + + def _adaptive_prune_heads(self, reserve_head_index: List[Tuple[int, int]]): + new_q_weight = [] + new_q_bias = [] + new_k_weight = [] + new_k_bias = [] + new_v_weight = [] + new_v_bias = [] + new_out_proj_weight = [] + + for ele in reserve_head_index: + start_idx, end_idx = ele + new_q_weight.append(self.q_proj.weight[start_idx:end_idx,]) + new_q_bias.append(self.q_proj.bias[start_idx:end_idx]) + + new_k_weight.append(self.k_proj.weight[start_idx:end_idx,]) + + new_k_bias.append(self.k_proj.bias[start_idx:end_idx]) + + new_v_weight.append(self.v_proj.weight[start_idx:end_idx,]) + new_v_bias.append(self.v_proj.bias[start_idx:end_idx]) + + new_out_proj_weight.append(self.out_proj.weight[:, start_idx:end_idx]) + + new_q_weight = torch.cat(new_q_weight).detach() + new_k_weight = torch.cat(new_k_weight).detach() + new_v_weight = torch.cat(new_v_weight).detach() + new_out_proj_weight = torch.cat(new_out_proj_weight, dim=-1).detach() + new_q_weight.requires_grad = True + new_k_weight.requires_grad = True + new_v_weight.requires_grad = True + new_out_proj_weight.requires_grad = True + + new_q_bias = torch.cat(new_q_bias).detach() + new_q_bias.requires_grad = True + + new_k_bias = torch.cat(new_k_bias).detach() + new_k_bias.requires_grad = True + + new_v_bias = torch.cat(new_v_bias).detach() + new_v_bias.requires_grad = True + + self.q_proj.weight = torch.nn.Parameter(new_q_weight) + self.q_proj.bias = torch.nn.Parameter(new_q_bias) + + self.k_proj.weight = torch.nn.Parameter(new_k_weight) + self.k_proj.bias = torch.nn.Parameter(new_k_bias) + + self.v_proj.weight = torch.nn.Parameter(new_v_weight) + self.v_proj.bias = torch.nn.Parameter(new_v_bias) + + self.out_proj.weight = torch.nn.Parameter(new_out_proj_weight) + + self.num_heads = len(reserve_head_index) + self.embed_dim = self.head_dim * self.num_heads + self.q_proj.out_features = self.embed_dim + self.k_proj.out_features = self.embed_dim + self.v_proj.out_features = self.embed_dim + + def _set_skip_embed_dim_check(self): + self.skip_embed_dim_check = True + + def forward( + self, + query, + key: Optional[Tensor], + value: Optional[Tensor], + key_padding_mask: Optional[Tensor] = None, + incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None, + need_weights: bool = True, + static_kv: bool = False, + attn_mask: Optional[Tensor] = None, + before_softmax: bool = False, + need_head_weights: bool = False, + ) -> Tuple[Tensor, Optional[Tensor]]: + """Input shape: Time x Batch x Channel + + Args: + key_padding_mask (ByteTensor, optional): mask to exclude + keys that are pads, of shape `(batch, src_len)`, where + padding elements are indicated by 1s. + need_weights (bool, optional): return the attention weights, + averaged over heads (default: False). + attn_mask (ByteTensor, optional): typically used to + implement causal attention, where the mask prevents the + attention from looking forward in time (default: None). + before_softmax (bool, optional): return the raw attention + weights and values before the attention softmax. + need_head_weights (bool, optional): return the attention + weights for each head. Implies *need_weights*. Default: + return the average attention weights over all heads. + """ + if need_head_weights: + need_weights = True + + is_tpu = query.device.type == "xla" + + tgt_len, bsz, embed_dim = query.size() + src_len = tgt_len + if not self.skip_embed_dim_check: + assert ( + embed_dim == self.embed_dim + ), f"query dim {embed_dim} != {self.embed_dim}" + assert list(query.size()) == [tgt_len, bsz, embed_dim] + if key is not None: + src_len, key_bsz, _ = key.size() + if not torch.jit.is_scripting(): + assert key_bsz == bsz + assert value is not None + assert src_len, bsz == value.shape[:2] + + if ( + not self.onnx_trace + and not is_tpu # don't use PyTorch version on TPUs + and incremental_state is None + and not static_kv + # A workaround for quantization to work. Otherwise JIT compilation + # treats bias in linear module as method. + and not torch.jit.is_scripting() + # The Multihead attention implemented in pytorch forces strong dimension check + # for input embedding dimention and K,Q,V projection dimension. + # Since pruning will break the dimension check and it is not easy to modify the pytorch API, + # it is preferred to bypass the pytorch MHA when we need to skip embed_dim_check + and not self.skip_embed_dim_check + ): + assert key is not None and value is not None + return F.multi_head_attention_forward( + query, + key, + value, + self.embed_dim, + self.num_heads, + torch.empty([0]), + torch.cat((self.q_proj.bias, self.k_proj.bias, self.v_proj.bias)), + self.bias_k, + self.bias_v, + self.add_zero_attn, + self.dropout_module.p, + self.out_proj.weight, + self.out_proj.bias, + self.training or self.dropout_module.apply_during_inference, + key_padding_mask, + need_weights, + attn_mask, + use_separate_proj_weight=True, + q_proj_weight=self.q_proj.weight, + k_proj_weight=self.k_proj.weight, + v_proj_weight=self.v_proj.weight, + ) + + if incremental_state is not None: + saved_state = self._get_input_buffer(incremental_state) + if saved_state is not None and "prev_key" in saved_state: + # previous time steps are cached - no need to recompute + # key and value if they are static + if static_kv: + assert self.encoder_decoder_attention and not self.self_attention + key = value = None + else: + saved_state = None + + if self.self_attention: + q = self.q_proj(query) + k = self.k_proj(query) + v = self.v_proj(query) + elif self.encoder_decoder_attention: + # encoder-decoder attention + q = self.q_proj(query) + if key is None: + assert value is None + k = v = None + else: + k = self.k_proj(key) + v = self.v_proj(key) + + else: + assert key is not None and value is not None + q = self.q_proj(query) + k = self.k_proj(key) + v = self.v_proj(value) + q *= self.scaling + + if self.bias_k is not None: + assert self.bias_v is not None + k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)]) + v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)]) + if attn_mask is not None: + attn_mask = torch.cat( + [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1 + ) + if key_padding_mask is not None: + key_padding_mask = torch.cat( + [ + key_padding_mask, + key_padding_mask.new_zeros(key_padding_mask.size(0), 1), + ], + dim=1, + ) + + q = ( + q.contiguous() + .view(tgt_len, bsz * self.num_heads, self.head_dim) + .transpose(0, 1) + ) + if k is not None: + k = ( + k.contiguous() + .view(-1, bsz * self.num_heads, self.head_dim) + .transpose(0, 1) + ) + if v is not None: + v = ( + v.contiguous() + .view(-1, bsz * self.num_heads, self.head_dim) + .transpose(0, 1) + ) + + if saved_state is not None: + # saved states are stored with shape (bsz, num_heads, seq_len, head_dim) + if "prev_key" in saved_state: + _prev_key = saved_state["prev_key"] + assert _prev_key is not None + prev_key = _prev_key.view(bsz * self.num_heads, -1, self.head_dim) + if static_kv: + k = prev_key + else: + assert k is not None + k = torch.cat([prev_key, k], dim=1) + src_len = k.size(1) + if "prev_value" in saved_state: + _prev_value = saved_state["prev_value"] + assert _prev_value is not None + prev_value = _prev_value.view(bsz * self.num_heads, -1, self.head_dim) + if static_kv: + v = prev_value + else: + assert v is not None + v = torch.cat([prev_value, v], dim=1) + prev_key_padding_mask: Optional[Tensor] = None + if "prev_key_padding_mask" in saved_state: + prev_key_padding_mask = saved_state["prev_key_padding_mask"] + assert k is not None and v is not None + key_padding_mask = MultiheadAttention._append_prev_key_padding_mask( + key_padding_mask=key_padding_mask, + prev_key_padding_mask=prev_key_padding_mask, + batch_size=bsz, + src_len=k.size(1), + static_kv=static_kv, + ) + + saved_state["prev_key"] = k.view(bsz, self.num_heads, -1, self.head_dim) + saved_state["prev_value"] = v.view(bsz, self.num_heads, -1, self.head_dim) + saved_state["prev_key_padding_mask"] = key_padding_mask + # In this branch incremental_state is never None + assert incremental_state is not None + incremental_state = self._set_input_buffer(incremental_state, saved_state) + assert k is not None + assert k.size(1) == src_len + + # This is part of a workaround to get around fork/join parallelism + # not supporting Optional types. + if key_padding_mask is not None and key_padding_mask.dim() == 0: + key_padding_mask = None + + if key_padding_mask is not None: + assert key_padding_mask.size(0) == bsz + assert key_padding_mask.size(1) == src_len + + if self.add_zero_attn: + assert v is not None + src_len += 1 + k = torch.cat([k, k.new_zeros((k.size(0), 1) + k.size()[2:])], dim=1) + v = torch.cat([v, v.new_zeros((v.size(0), 1) + v.size()[2:])], dim=1) + if attn_mask is not None: + attn_mask = torch.cat( + [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1 + ) + if key_padding_mask is not None: + key_padding_mask = torch.cat( + [ + key_padding_mask, + torch.zeros(key_padding_mask.size(0), 1).type_as( + key_padding_mask + ), + ], + dim=1, + ) + + attn_weights = torch.bmm(q, k.transpose(1, 2)) + attn_weights = self.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz) + + assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len] + + if attn_mask is not None: + attn_mask = attn_mask.unsqueeze(0) + if self.onnx_trace: + attn_mask = attn_mask.repeat(attn_weights.size(0), 1, 1) + attn_weights += attn_mask + + if key_padding_mask is not None: + # don't attend to padding symbols + attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + if not is_tpu: + attn_weights = attn_weights.masked_fill( + key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), + float("-inf"), + ) + else: + attn_weights = attn_weights.transpose(0, 2) + attn_weights = attn_weights.masked_fill(key_padding_mask, float("-inf")) + attn_weights = attn_weights.transpose(0, 2) + attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) + + if before_softmax: + return attn_weights, v + + attn_weights_float = F.softmax(attn_weights, dim=-1, dtype=torch.float32) + attn_weights = attn_weights_float.type_as(attn_weights) + attn_probs = self.dropout_module(attn_weights) + + assert v is not None + attn = torch.bmm(attn_probs, v) + assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim] + if self.onnx_trace and attn.size(1) == 1: + # when ONNX tracing a single decoder step (sequence length == 1) + # the transpose is a no-op copy before view, thus unnecessary + attn = attn.contiguous().view(tgt_len, bsz, self.embed_dim) + else: + attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, self.embed_dim) + attn = self.out_proj(attn) + attn_weights: Optional[Tensor] = None + if need_weights: + attn_weights = attn_weights_float.view( + bsz, self.num_heads, tgt_len, src_len + ).transpose(1, 0) + if not need_head_weights: + # average attention weights over heads + attn_weights = attn_weights.mean(dim=0) + + return attn, attn_weights + + @staticmethod + def _append_prev_key_padding_mask( + key_padding_mask: Optional[Tensor], + prev_key_padding_mask: Optional[Tensor], + batch_size: int, + src_len: int, + static_kv: bool, + ) -> Optional[Tensor]: + # saved key padding masks have shape (bsz, seq_len) + if prev_key_padding_mask is not None and static_kv: + new_key_padding_mask = prev_key_padding_mask + elif prev_key_padding_mask is not None and key_padding_mask is not None: + new_key_padding_mask = torch.cat( + [prev_key_padding_mask.float(), key_padding_mask.float()], dim=1 + ) + # During incremental decoding, as the padding token enters and + # leaves the frame, there will be a time when prev or current + # is None + elif prev_key_padding_mask is not None: + if src_len > prev_key_padding_mask.size(1): + filler = torch.zeros( + (batch_size, src_len - prev_key_padding_mask.size(1)), + device=prev_key_padding_mask.device, + ) + new_key_padding_mask = torch.cat( + [prev_key_padding_mask.float(), filler.float()], dim=1 + ) + else: + new_key_padding_mask = prev_key_padding_mask.float() + elif key_padding_mask is not None: + if src_len > key_padding_mask.size(1): + filler = torch.zeros( + (batch_size, src_len - key_padding_mask.size(1)), + device=key_padding_mask.device, + ) + new_key_padding_mask = torch.cat( + [filler.float(), key_padding_mask.float()], dim=1 + ) + else: + new_key_padding_mask = key_padding_mask.float() + else: + new_key_padding_mask = prev_key_padding_mask + return new_key_padding_mask + + @torch.jit.export + def reorder_incremental_state( + self, + incremental_state: Dict[str, Dict[str, Optional[Tensor]]], + new_order: Tensor, + ): + """Reorder buffered internal state (for incremental generation).""" + input_buffer = self._get_input_buffer(incremental_state) + if input_buffer is not None: + for k in input_buffer.keys(): + input_buffer_k = input_buffer[k] + if input_buffer_k is not None: + if self.encoder_decoder_attention and input_buffer_k.size( + 0 + ) == new_order.size(0): + break + input_buffer[k] = input_buffer_k.index_select(0, new_order) + incremental_state = self._set_input_buffer(incremental_state, input_buffer) + return incremental_state + + def _get_input_buffer( + self, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] + ) -> Dict[str, Optional[Tensor]]: + result = self.get_incremental_state(incremental_state, "attn_state") + if result is not None: + return result + else: + empty_result: Dict[str, Optional[Tensor]] = {} + return empty_result + + def _set_input_buffer( + self, + incremental_state: Dict[str, Dict[str, Optional[Tensor]]], + buffer: Dict[str, Optional[Tensor]], + ): + return self.set_incremental_state(incremental_state, "attn_state", buffer) + + def apply_sparse_mask(self, attn_weights, tgt_len: int, src_len: int, bsz: int): + return attn_weights + + def upgrade_state_dict_named(self, state_dict, name): + prefix = name + "." if name != "" else "" + items_to_add = {} + keys_to_remove = [] + for k in state_dict.keys(): + if k.endswith(prefix + "in_proj_weight"): + # in_proj_weight used to be q + k + v with same dimensions + dim = int(state_dict[k].shape[0] / 3) + items_to_add[prefix + "q_proj.weight"] = state_dict[k][:dim] + items_to_add[prefix + "k_proj.weight"] = state_dict[k][dim : 2 * dim] + items_to_add[prefix + "v_proj.weight"] = state_dict[k][2 * dim :] + + keys_to_remove.append(k) + + k_bias = prefix + "in_proj_bias" + if k_bias in state_dict.keys(): + dim = int(state_dict[k].shape[0] / 3) + items_to_add[prefix + "q_proj.bias"] = state_dict[k_bias][:dim] + items_to_add[prefix + "k_proj.bias"] = state_dict[k_bias][ + dim : 2 * dim + ] + items_to_add[prefix + "v_proj.bias"] = state_dict[k_bias][2 * dim :] + + keys_to_remove.append(prefix + "in_proj_bias") + + for k in keys_to_remove: + del state_dict[k] + + for key, value in items_to_add.items(): + state_dict[key] = value diff --git a/almeval/models/stepaudio/funasr_detach/models/data2vec/quant_noise.py b/almeval/models/stepaudio/funasr_detach/models/data2vec/quant_noise.py new file mode 100644 index 0000000000000000000000000000000000000000..d777dfbb6c1bf6a9b769dfdaec35d5ef084c8a8b --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/data2vec/quant_noise.py @@ -0,0 +1,107 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.nn as nn + + +def quant_noise(module, p, block_size): + """ + Wraps modules and applies quantization noise to the weights for + subsequent quantization with Iterative Product Quantization as + described in "Training with Quantization Noise for Extreme Model Compression" + + Args: + - module: nn.Module + - p: amount of Quantization Noise + - block_size: size of the blocks for subsequent quantization with iPQ + + Remarks: + - Module weights must have the right sizes wrt the block size + - Only Linear, Embedding and Conv2d modules are supported for the moment + - For more detail on how to quantize by blocks with convolutional weights, + see "And the Bit Goes Down: Revisiting the Quantization of Neural Networks" + - We implement the simplest form of noise here as stated in the paper + which consists in randomly dropping blocks + """ + + # if no quantization noise, don't register hook + if p <= 0: + return module + + # supported modules + assert isinstance(module, (nn.Linear, nn.Embedding, nn.Conv2d)) + + # test whether module.weight has the right sizes wrt block_size + is_conv = module.weight.ndim == 4 + + # 2D matrix + if not is_conv: + assert ( + module.weight.size(1) % block_size == 0 + ), "Input features must be a multiple of block sizes" + + # 4D matrix + else: + # 1x1 convolutions + if module.kernel_size == (1, 1): + assert ( + module.in_channels % block_size == 0 + ), "Input channels must be a multiple of block sizes" + # regular convolutions + else: + k = module.kernel_size[0] * module.kernel_size[1] + assert k % block_size == 0, "Kernel size must be a multiple of block size" + + def _forward_pre_hook(mod, input): + # no noise for evaluation + if mod.training: + if not is_conv: + # gather weight and sizes + weight = mod.weight + in_features = weight.size(1) + out_features = weight.size(0) + + # split weight matrix into blocks and randomly drop selected blocks + mask = torch.zeros( + in_features // block_size * out_features, device=weight.device + ) + mask.bernoulli_(p) + mask = mask.repeat_interleave(block_size, -1).view(-1, in_features) + + else: + # gather weight and sizes + weight = mod.weight + in_channels = mod.in_channels + out_channels = mod.out_channels + + # split weight matrix into blocks and randomly drop selected blocks + if mod.kernel_size == (1, 1): + mask = torch.zeros( + int(in_channels // block_size * out_channels), + device=weight.device, + ) + mask.bernoulli_(p) + mask = mask.repeat_interleave(block_size, -1).view(-1, in_channels) + else: + mask = torch.zeros( + weight.size(0), weight.size(1), device=weight.device + ) + mask.bernoulli_(p) + mask = ( + mask.unsqueeze(2) + .unsqueeze(3) + .repeat(1, 1, mod.kernel_size[0], mod.kernel_size[1]) + ) + + # scale weights and apply mask + mask = mask.to( + torch.bool + ) # x.bool() is not currently supported in TorchScript + s = 1 / (1 - p) + mod.weight.data = s * weight.masked_fill(mask, 0) + + module.register_forward_pre_hook(_forward_pre_hook) + return module diff --git a/almeval/models/stepaudio/funasr_detach/models/data2vec/utils.py b/almeval/models/stepaudio/funasr_detach/models/data2vec/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9b92bfeb5ab3f1c8008089754a9a1b36d6fd8a5f --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/data2vec/utils.py @@ -0,0 +1,156 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from funasr_detach.models.data2vec.multihead_attention import MultiheadAttention + + +class Fp32LayerNorm(nn.LayerNorm): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self, input): + output = F.layer_norm( + input.float(), + self.normalized_shape, + self.weight.float() if self.weight is not None else None, + self.bias.float() if self.bias is not None else None, + self.eps, + ) + return output.type_as(input) + + +class Fp32GroupNorm(nn.GroupNorm): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self, input): + output = F.group_norm( + input.float(), + self.num_groups, + self.weight.float() if self.weight is not None else None, + self.bias.float() if self.bias is not None else None, + self.eps, + ) + return output.type_as(input) + + +class TransposeLast(nn.Module): + def __init__(self, deconstruct_idx=None): + super().__init__() + self.deconstruct_idx = deconstruct_idx + + def forward(self, x): + if self.deconstruct_idx is not None: + x = x[self.deconstruct_idx] + return x.transpose(-2, -1) + + +class SamePad(nn.Module): + def __init__(self, kernel_size, causal=False): + super().__init__() + if causal: + self.remove = kernel_size - 1 + else: + self.remove = 1 if kernel_size % 2 == 0 else 0 + + def forward(self, x): + if self.remove > 0: + x = x[:, :, : -self.remove] + return x + + +def pad_to_multiple(x, multiple, dim=-1, value=0): + # Inspired from https://github.com/lucidrains/local-attention/blob/master/local_attention/local_attention.py#L41 + if x is None: + return None, 0 + tsz = x.size(dim) + m = tsz / multiple + remainder = math.ceil(m) * multiple - tsz + if m.is_integer(): + return x, 0 + pad_offset = (0,) * (-1 - dim) * 2 + + return F.pad(x, (*pad_offset, 0, remainder), value=value), remainder + + +def gelu_accurate(x): + if not hasattr(gelu_accurate, "_a"): + gelu_accurate._a = math.sqrt(2 / math.pi) + return ( + 0.5 * x * (1 + torch.tanh(gelu_accurate._a * (x + 0.044715 * torch.pow(x, 3)))) + ) + + +def gelu(x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.gelu(x.float()).type_as(x) + + +def get_available_activation_fns(): + return [ + "relu", + "gelu", + "gelu_fast", # deprecated + "gelu_accurate", + "tanh", + "linear", + ] + + +def get_activation_fn(activation: str): + """Returns the activation function corresponding to `activation`""" + + if activation == "relu": + return F.relu + elif activation == "gelu": + return gelu + elif activation == "gelu_accurate": + return gelu_accurate + elif activation == "tanh": + return torch.tanh + elif activation == "linear": + return lambda x: x + elif activation == "swish": + return torch.nn.SiLU + else: + raise RuntimeError("--activation-fn {} not supported".format(activation)) + + +def init_bert_params(module): + """ + Initialize the weights specific to the BERT Model. + This overrides the default initializations depending on the specified arguments. + 1. If normal_init_linear_weights is set then weights of linear + layer will be initialized using the normal distribution and + bais will be set to the specified value. + 2. If normal_init_embed_weights is set then weights of embedding + layer will be initialized using the normal distribution. + 3. If normal_init_proj_weights is set then weights of + in_project_weight for MultiHeadAttention initialized using + the normal distribution (to be validated). + """ + + def normal_(data): + # with FSDP, module params will be on CUDA, so we cast them back to CPU + # so that the RNG is consistent with and without FSDP + data.copy_(data.cpu().normal_(mean=0.0, std=0.02).to(data.device)) + + if isinstance(module, nn.Linear): + normal_(module.weight.data) + if module.bias is not None: + module.bias.data.zero_() + if isinstance(module, nn.Embedding): + normal_(module.weight.data) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + if isinstance(module, MultiheadAttention): + normal_(module.q_proj.weight.data) + normal_(module.k_proj.weight.data) + normal_(module.v_proj.weight.data) diff --git a/almeval/models/stepaudio/funasr_detach/models/data2vec/wav2vec2.py b/almeval/models/stepaudio/funasr_detach/models/data2vec/wav2vec2.py new file mode 100644 index 0000000000000000000000000000000000000000..cbda714ce0d1184f2cb1f5b1c16a880342b6c901 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/data2vec/wav2vec2.py @@ -0,0 +1,407 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import logging +import math +from typing import List, Tuple + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from funasr_detach.models.data2vec import utils +from funasr_detach.models.data2vec.multihead_attention import MultiheadAttention + + +class ConvFeatureExtractionModel(nn.Module): + def __init__( + self, + conv_layers: List[Tuple[int, int, int]], + dropout: float = 0.0, + mode: str = "default", + conv_bias: bool = False, + in_d: int = 1, + ): + super().__init__() + + assert mode in {"default", "layer_norm"} + + def block( + n_in, + n_out, + k, + stride, + is_layer_norm=False, + is_group_norm=False, + conv_bias=False, + ): + def make_conv(): + conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias) + nn.init.kaiming_normal_(conv.weight) + return conv + + assert ( + is_layer_norm and is_group_norm + ) == False, "layer norm and group norm are exclusive" + + if is_layer_norm: + return nn.Sequential( + make_conv(), + nn.Dropout(p=dropout), + nn.Sequential( + utils.TransposeLast(), + utils.Fp32LayerNorm(dim, elementwise_affine=True), + utils.TransposeLast(), + ), + nn.GELU(), + ) + elif is_group_norm: + return nn.Sequential( + make_conv(), + nn.Dropout(p=dropout), + utils.Fp32GroupNorm(dim, dim, affine=True), + nn.GELU(), + ) + else: + return nn.Sequential(make_conv(), nn.Dropout(p=dropout), nn.GELU()) + + self.conv_layers = nn.ModuleList() + for i, cl in enumerate(conv_layers): + assert len(cl) == 3, "invalid conv definition: " + str(cl) + (dim, k, stride) = cl + + self.conv_layers.append( + block( + in_d, + dim, + k, + stride, + is_layer_norm=mode == "layer_norm", + is_group_norm=mode == "default" and i == 0, + conv_bias=conv_bias, + ) + ) + in_d = dim + + def forward(self, x): + if len(x.shape) == 2: + x = x.unsqueeze(1) + else: + x = x.transpose(1, 2) + + for conv in self.conv_layers: + x = conv(x) + return x + + +def make_conv_pos(e, k, g): + pos_conv = nn.Conv1d( + e, + e, + kernel_size=k, + padding=k // 2, + groups=g, + ) + dropout = 0 + std = math.sqrt((4 * (1.0 - dropout)) / (k * e)) + nn.init.normal_(pos_conv.weight, mean=0, std=std) + nn.init.constant_(pos_conv.bias, 0) + + pos_conv = nn.utils.weight_norm(pos_conv, name="weight", dim=2) + pos_conv = nn.Sequential(pos_conv, utils.SamePad(k), nn.GELU()) + + return pos_conv + + +class TransformerEncoder(nn.Module): + def build_encoder_layer(self): + if self.layer_type == "transformer": + layer = TransformerSentenceEncoderLayer( + embedding_dim=self.embedding_dim, + ffn_embedding_dim=self.encoder_ffn_embed_dim, + num_attention_heads=self.encoder_attention_heads, + dropout=self.dropout, + attention_dropout=self.attention_dropout, + activation_dropout=self.activation_dropout, + activation_fn=self.activation_fn, + layer_norm_first=self.layer_norm_first, + ) + else: + logging.error("Only transformer is supported for data2vec now") + return layer + + def __init__( + self, + # position + dropout, + encoder_embed_dim, + required_seq_len_multiple, + pos_conv_depth, + conv_pos, + conv_pos_groups, + # transformer layers + layer_type, + encoder_layers, + encoder_ffn_embed_dim, + encoder_attention_heads, + attention_dropout, + activation_dropout, + activation_fn, + layer_norm_first, + encoder_layerdrop, + max_positions, + ): + super().__init__() + + # position + self.dropout = dropout + self.embedding_dim = encoder_embed_dim + self.required_seq_len_multiple = required_seq_len_multiple + if pos_conv_depth > 1: + num_layers = pos_conv_depth + k = max(3, conv_pos // num_layers) + + def make_conv_block(e, k, g, l): + return nn.Sequential( + *[ + nn.Sequential( + nn.Conv1d( + e, + e, + kernel_size=k, + padding=k // 2, + groups=g, + ), + utils.SamePad(k), + utils.TransposeLast(), + torch.nn.LayerNorm(e, elementwise_affine=False), + utils.TransposeLast(), + nn.GELU(), + ) + for _ in range(l) + ] + ) + + self.pos_conv = make_conv_block( + self.embedding_dim, k, conv_pos_groups, num_layers + ) + + else: + self.pos_conv = make_conv_pos( + self.embedding_dim, + conv_pos, + conv_pos_groups, + ) + + # transformer layers + self.layer_type = layer_type + self.encoder_ffn_embed_dim = encoder_ffn_embed_dim + self.encoder_attention_heads = encoder_attention_heads + self.attention_dropout = attention_dropout + self.activation_dropout = activation_dropout + self.activation_fn = activation_fn + self.layer_norm_first = layer_norm_first + self.layerdrop = encoder_layerdrop + self.max_positions = max_positions + self.layers = nn.ModuleList( + [self.build_encoder_layer() for _ in range(encoder_layers)] + ) + self.layer_norm = torch.nn.LayerNorm(self.embedding_dim) + + self.apply(utils.init_bert_params) + + def forward(self, x, padding_mask=None, layer=None): + x, layer_results = self.extract_features(x, padding_mask, layer) + + if self.layer_norm_first and layer is None: + x = self.layer_norm(x) + + return x, layer_results + + def extract_features( + self, + x, + padding_mask=None, + tgt_layer=None, + min_layer=0, + ): + + if padding_mask is not None: + x[padding_mask] = 0 + + x_conv = self.pos_conv(x.transpose(1, 2)) + x_conv = x_conv.transpose(1, 2) + x = x + x_conv + + if not self.layer_norm_first: + x = self.layer_norm(x) + + # pad to the sequence length dimension + x, pad_length = utils.pad_to_multiple( + x, self.required_seq_len_multiple, dim=-2, value=0 + ) + if pad_length > 0 and padding_mask is None: + padding_mask = x.new_zeros((x.size(0), x.size(1)), dtype=torch.bool) + padding_mask[:, -pad_length:] = True + else: + padding_mask, _ = utils.pad_to_multiple( + padding_mask, self.required_seq_len_multiple, dim=-1, value=True + ) + x = F.dropout(x, p=self.dropout, training=self.training) + + # B x T x C -> T x B x C + x = x.transpose(0, 1) + + layer_results = [] + r = None + for i, layer in enumerate(self.layers): + dropout_probability = np.random.random() if self.layerdrop > 0 else 1 + if not self.training or (dropout_probability > self.layerdrop): + x, (z, lr) = layer(x, self_attn_padding_mask=padding_mask) + if i >= min_layer: + layer_results.append((x, z, lr)) + if i == tgt_layer: + r = x + break + + if r is not None: + x = r + + # T x B x C -> B x T x C + x = x.transpose(0, 1) + + # undo paddding + if pad_length > 0: + x = x[:, :-pad_length] + + def undo_pad(a, b, c): + return ( + a[:-pad_length], + b[:-pad_length] if b is not None else b, + c[:-pad_length], + ) + + layer_results = [undo_pad(*u) for u in layer_results] + + return x, layer_results + + def max_positions(self): + """Maximum output length supported by the encoder.""" + return self.max_positions + + def upgrade_state_dict_named(self, state_dict, name): + """Upgrade a (possibly old) state dict for new versions of fairseq.""" + return state_dict + + +class TransformerSentenceEncoderLayer(nn.Module): + """ + Implements a Transformer Encoder Layer used in BERT/XLM style pre-trained + models. + """ + + def __init__( + self, + embedding_dim: int = 768, + ffn_embedding_dim: int = 3072, + num_attention_heads: int = 8, + dropout: float = 0.1, + attention_dropout: float = 0.1, + activation_dropout: float = 0.1, + activation_fn: str = "relu", + layer_norm_first: bool = False, + ) -> None: + + super().__init__() + # Initialize parameters + self.embedding_dim = embedding_dim + self.dropout = dropout + self.activation_dropout = activation_dropout + + # Initialize blocks + self.activation_fn = utils.get_activation_fn(activation_fn) + self.self_attn = MultiheadAttention( + self.embedding_dim, + num_attention_heads, + dropout=attention_dropout, + self_attention=True, + ) + + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(self.activation_dropout) + self.dropout3 = nn.Dropout(dropout) + + self.layer_norm_first = layer_norm_first + + # layer norm associated with the self attention layer + self.self_attn_layer_norm = torch.nn.LayerNorm(self.embedding_dim) + self.fc1 = nn.Linear(self.embedding_dim, ffn_embedding_dim) + self.fc2 = nn.Linear(ffn_embedding_dim, self.embedding_dim) + + # layer norm associated with the position wise feed-forward NN + self.final_layer_norm = torch.nn.LayerNorm(self.embedding_dim) + + def forward( + self, + x: torch.Tensor, # (T, B, C) + self_attn_mask: torch.Tensor = None, + self_attn_padding_mask: torch.Tensor = None, + ): + """ + LayerNorm is applied either before or after the self-attention/ffn + modules similar to the original Transformer imlementation. + """ + residual = x + + if self.layer_norm_first: + x = self.self_attn_layer_norm(x) + x, attn = self.self_attn( + query=x, + key=x, + value=x, + key_padding_mask=self_attn_padding_mask, + attn_mask=self_attn_mask, + need_weights=False, + ) + x = self.dropout1(x) + x = residual + x + + residual = x + x = self.final_layer_norm(x) + x = self.activation_fn(self.fc1(x)) + x = self.dropout2(x) + x = self.fc2(x) + + layer_result = x + + x = self.dropout3(x) + x = residual + x + else: + x, attn = self.self_attn( + query=x, + key=x, + value=x, + key_padding_mask=self_attn_padding_mask, + need_weights=False, + ) + + x = self.dropout1(x) + x = residual + x + + x = self.self_attn_layer_norm(x) + + residual = x + x = self.activation_fn(self.fc1(x)) + x = self.dropout2(x) + x = self.fc2(x) + + layer_result = x + + x = self.dropout3(x) + x = residual + x + x = self.final_layer_norm(x) + + return x, (attn, layer_result) diff --git a/almeval/models/stepaudio/funasr_detach/models/e_branchformer/__init__.py b/almeval/models/stepaudio/funasr_detach/models/e_branchformer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/almeval/models/stepaudio/funasr_detach/models/e_branchformer/encoder.py b/almeval/models/stepaudio/funasr_detach/models/e_branchformer/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..33ab040454ffbb48dc8677456187f7a6a5db57e7 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/e_branchformer/encoder.py @@ -0,0 +1,471 @@ +# Copyright 2022 Kwangyoun Kim (ASAPP inc.) +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""E-Branchformer encoder definition. +Reference: + Kwangyoun Kim, Felix Wu, Yifan Peng, Jing Pan, + Prashant Sridhar, Kyu J. Han, Shinji Watanabe, + "E-Branchformer: Branchformer with Enhanced merging + for speech recognition," in SLT 2022. +""" + +import logging +from typing import List, Optional, Tuple + +import torch +import torch.nn as nn +from funasr_detach.models.ctc.ctc import CTC +from funasr_detach.models.branchformer.cgmlp import ConvolutionalGatingMLP +from funasr_detach.models.branchformer.fastformer import FastSelfAttention +from funasr_detach.models.transformer.utils.nets_utils import ( + get_activation, + make_pad_mask, +) +from funasr_detach.models.transformer.attention import ( # noqa: H301 + LegacyRelPositionMultiHeadedAttention, + MultiHeadedAttention, + RelPositionMultiHeadedAttention, +) +from funasr_detach.models.transformer.embedding import ( # noqa: H301 + LegacyRelPositionalEncoding, + PositionalEncoding, + RelPositionalEncoding, + ScaledPositionalEncoding, +) +from funasr_detach.models.transformer.layer_norm import LayerNorm +from funasr_detach.models.transformer.positionwise_feed_forward import ( + PositionwiseFeedForward, +) +from funasr_detach.models.transformer.utils.repeat import repeat +from funasr_detach.models.transformer.utils.subsampling import ( + Conv2dSubsampling, + Conv2dSubsampling2, + Conv2dSubsampling6, + Conv2dSubsampling8, + TooShortUttError, + check_short_utt, +) +from funasr_detach.register import tables + + +class EBranchformerEncoderLayer(torch.nn.Module): + """E-Branchformer encoder layer module. + + Args: + size (int): model dimension + attn: standard self-attention or efficient attention + cgmlp: ConvolutionalGatingMLP + feed_forward: feed-forward module, optional + feed_forward: macaron-style feed-forward module, optional + dropout_rate (float): dropout probability + merge_conv_kernel (int): kernel size of the depth-wise conv in merge module + """ + + def __init__( + self, + size: int, + attn: torch.nn.Module, + cgmlp: torch.nn.Module, + feed_forward: Optional[torch.nn.Module], + feed_forward_macaron: Optional[torch.nn.Module], + dropout_rate: float, + merge_conv_kernel: int = 3, + ): + super().__init__() + + self.size = size + self.attn = attn + self.cgmlp = cgmlp + + self.feed_forward = feed_forward + self.feed_forward_macaron = feed_forward_macaron + self.ff_scale = 1.0 + if self.feed_forward is not None: + self.norm_ff = LayerNorm(size) + if self.feed_forward_macaron is not None: + self.ff_scale = 0.5 + self.norm_ff_macaron = LayerNorm(size) + + self.norm_mha = LayerNorm(size) # for the MHA module + self.norm_mlp = LayerNorm(size) # for the MLP module + self.norm_final = LayerNorm(size) # for the final output of the block + + self.dropout = torch.nn.Dropout(dropout_rate) + + self.depthwise_conv_fusion = torch.nn.Conv1d( + size + size, + size + size, + kernel_size=merge_conv_kernel, + stride=1, + padding=(merge_conv_kernel - 1) // 2, + groups=size + size, + bias=True, + ) + self.merge_proj = torch.nn.Linear(size + size, size) + + def forward(self, x_input, mask, cache=None): + """Compute encoded features. + + Args: + x_input (Union[Tuple, torch.Tensor]): Input tensor w/ or w/o pos emb. + - w/ pos emb: Tuple of tensors [(#batch, time, size), (1, time, size)]. + - w/o pos emb: Tensor (#batch, time, size). + mask (torch.Tensor): Mask tensor for the input (#batch, 1, time). + cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size). + Returns: + torch.Tensor: Output tensor (#batch, time, size). + torch.Tensor: Mask tensor (#batch, time). + """ + + if cache is not None: + raise NotImplementedError("cache is not None, which is not tested") + + if isinstance(x_input, tuple): + x, pos_emb = x_input[0], x_input[1] + else: + x, pos_emb = x_input, None + + if self.feed_forward_macaron is not None: + residual = x + x = self.norm_ff_macaron(x) + x = residual + self.ff_scale * self.dropout(self.feed_forward_macaron(x)) + + # Two branches + x1 = x + x2 = x + + # Branch 1: multi-headed attention module + x1 = self.norm_mha(x1) + + if isinstance(self.attn, FastSelfAttention): + x_att = self.attn(x1, mask) + else: + if pos_emb is not None: + x_att = self.attn(x1, x1, x1, pos_emb, mask) + else: + x_att = self.attn(x1, x1, x1, mask) + + x1 = self.dropout(x_att) + + # Branch 2: convolutional gating mlp + x2 = self.norm_mlp(x2) + + if pos_emb is not None: + x2 = (x2, pos_emb) + x2 = self.cgmlp(x2, mask) + if isinstance(x2, tuple): + x2 = x2[0] + + x2 = self.dropout(x2) + + # Merge two branches + x_concat = torch.cat([x1, x2], dim=-1) + x_tmp = x_concat.transpose(1, 2) + x_tmp = self.depthwise_conv_fusion(x_tmp) + x_tmp = x_tmp.transpose(1, 2) + x = x + self.dropout(self.merge_proj(x_concat + x_tmp)) + + if self.feed_forward is not None: + # feed forward module + residual = x + x = self.norm_ff(x) + x = residual + self.ff_scale * self.dropout(self.feed_forward(x)) + + x = self.norm_final(x) + + if pos_emb is not None: + return (x, pos_emb), mask + + return x, mask + + +@tables.register("encoder_classes", "EBranchformerEncoder") +class EBranchformerEncoder(nn.Module): + """E-Branchformer encoder module.""" + + def __init__( + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + attention_layer_type: str = "rel_selfattn", + pos_enc_layer_type: str = "rel_pos", + rel_pos_type: str = "latest", + cgmlp_linear_units: int = 2048, + cgmlp_conv_kernel: int = 31, + use_linear_after_conv: bool = False, + gate_activation: str = "identity", + num_blocks: int = 12, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.0, + input_layer: Optional[str] = "conv2d", + zero_triu: bool = False, + padding_idx: int = -1, + layer_drop_rate: float = 0.0, + max_pos_emb_len: int = 5000, + use_ffn: bool = False, + macaron_ffn: bool = False, + ffn_activation_type: str = "swish", + linear_units: int = 2048, + positionwise_layer_type: str = "linear", + merge_conv_kernel: int = 3, + interctc_layer_idx=None, + interctc_use_conditioning: bool = False, + ): + super().__init__() + self._output_size = output_size + + if rel_pos_type == "legacy": + if pos_enc_layer_type == "rel_pos": + pos_enc_layer_type = "legacy_rel_pos" + if attention_layer_type == "rel_selfattn": + attention_layer_type = "legacy_rel_selfattn" + elif rel_pos_type == "latest": + assert attention_layer_type != "legacy_rel_selfattn" + assert pos_enc_layer_type != "legacy_rel_pos" + else: + raise ValueError("unknown rel_pos_type: " + rel_pos_type) + + if pos_enc_layer_type == "abs_pos": + pos_enc_class = PositionalEncoding + elif pos_enc_layer_type == "scaled_abs_pos": + pos_enc_class = ScaledPositionalEncoding + elif pos_enc_layer_type == "rel_pos": + assert attention_layer_type == "rel_selfattn" + pos_enc_class = RelPositionalEncoding + elif pos_enc_layer_type == "legacy_rel_pos": + assert attention_layer_type == "legacy_rel_selfattn" + pos_enc_class = LegacyRelPositionalEncoding + logging.warning( + "Using legacy_rel_pos and it will be deprecated in the future." + ) + else: + raise ValueError("unknown pos_enc_layer: " + pos_enc_layer_type) + + if input_layer == "linear": + self.embed = torch.nn.Sequential( + torch.nn.Linear(input_size, output_size), + torch.nn.LayerNorm(output_size), + torch.nn.Dropout(dropout_rate), + pos_enc_class(output_size, positional_dropout_rate, max_pos_emb_len), + ) + elif input_layer == "conv2d": + self.embed = Conv2dSubsampling( + input_size, + output_size, + dropout_rate, + pos_enc_class(output_size, positional_dropout_rate, max_pos_emb_len), + ) + elif input_layer == "conv2d2": + self.embed = Conv2dSubsampling2( + input_size, + output_size, + dropout_rate, + pos_enc_class(output_size, positional_dropout_rate, max_pos_emb_len), + ) + elif input_layer == "conv2d6": + self.embed = Conv2dSubsampling6( + input_size, + output_size, + dropout_rate, + pos_enc_class(output_size, positional_dropout_rate, max_pos_emb_len), + ) + elif input_layer == "conv2d8": + self.embed = Conv2dSubsampling8( + input_size, + output_size, + dropout_rate, + pos_enc_class(output_size, positional_dropout_rate, max_pos_emb_len), + ) + elif input_layer == "embed": + self.embed = torch.nn.Sequential( + torch.nn.Embedding(input_size, output_size, padding_idx=padding_idx), + pos_enc_class(output_size, positional_dropout_rate, max_pos_emb_len), + ) + elif isinstance(input_layer, torch.nn.Module): + self.embed = torch.nn.Sequential( + input_layer, + pos_enc_class(output_size, positional_dropout_rate, max_pos_emb_len), + ) + elif input_layer is None: + if input_size == output_size: + self.embed = None + else: + self.embed = torch.nn.Linear(input_size, output_size) + else: + raise ValueError("unknown input_layer: " + input_layer) + + activation = get_activation(ffn_activation_type) + if positionwise_layer_type == "linear": + positionwise_layer = PositionwiseFeedForward + positionwise_layer_args = ( + output_size, + linear_units, + dropout_rate, + activation, + ) + elif positionwise_layer_type is None: + logging.warning("no macaron ffn") + else: + raise ValueError("Support only linear.") + + if attention_layer_type == "selfattn": + encoder_selfattn_layer = MultiHeadedAttention + encoder_selfattn_layer_args = ( + attention_heads, + output_size, + attention_dropout_rate, + ) + elif attention_layer_type == "legacy_rel_selfattn": + assert pos_enc_layer_type == "legacy_rel_pos" + encoder_selfattn_layer = LegacyRelPositionMultiHeadedAttention + encoder_selfattn_layer_args = ( + attention_heads, + output_size, + attention_dropout_rate, + ) + logging.warning( + "Using legacy_rel_selfattn and it will be deprecated in the future." + ) + elif attention_layer_type == "rel_selfattn": + assert pos_enc_layer_type == "rel_pos" + encoder_selfattn_layer = RelPositionMultiHeadedAttention + encoder_selfattn_layer_args = ( + attention_heads, + output_size, + attention_dropout_rate, + zero_triu, + ) + elif attention_layer_type == "fast_selfattn": + assert pos_enc_layer_type in ["abs_pos", "scaled_abs_pos"] + encoder_selfattn_layer = FastSelfAttention + encoder_selfattn_layer_args = ( + output_size, + attention_heads, + attention_dropout_rate, + ) + else: + raise ValueError("unknown encoder_attn_layer: " + attention_layer_type) + + cgmlp_layer = ConvolutionalGatingMLP + cgmlp_layer_args = ( + output_size, + cgmlp_linear_units, + cgmlp_conv_kernel, + dropout_rate, + use_linear_after_conv, + gate_activation, + ) + + self.encoders = repeat( + num_blocks, + lambda lnum: EBranchformerEncoderLayer( + output_size, + encoder_selfattn_layer(*encoder_selfattn_layer_args), + cgmlp_layer(*cgmlp_layer_args), + positionwise_layer(*positionwise_layer_args) if use_ffn else None, + ( + positionwise_layer(*positionwise_layer_args) + if use_ffn and macaron_ffn + else None + ), + dropout_rate, + merge_conv_kernel, + ), + layer_drop_rate, + ) + self.after_norm = LayerNorm(output_size) + + if interctc_layer_idx is None: + interctc_layer_idx = [] + self.interctc_layer_idx = interctc_layer_idx + if len(interctc_layer_idx) > 0: + assert 0 < min(interctc_layer_idx) and max(interctc_layer_idx) < num_blocks + self.interctc_use_conditioning = interctc_use_conditioning + self.conditioning_layer = None + + def output_size(self) -> int: + return self._output_size + + def forward( + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, + ctc: CTC = None, + max_layer: int = None, + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Calculate forward propagation. + + Args: + xs_pad (torch.Tensor): Input tensor (#batch, L, input_size). + ilens (torch.Tensor): Input length (#batch). + prev_states (torch.Tensor): Not to be used now. + ctc (CTC): Intermediate CTC module. + max_layer (int): Layer depth below which InterCTC is applied. + Returns: + torch.Tensor: Output tensor (#batch, L, output_size). + torch.Tensor: Output length (#batch). + torch.Tensor: Not to be used now. + """ + + masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device) + + if ( + isinstance(self.embed, Conv2dSubsampling) + or isinstance(self.embed, Conv2dSubsampling2) + or isinstance(self.embed, Conv2dSubsampling6) + or isinstance(self.embed, Conv2dSubsampling8) + ): + short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1)) + if short_status: + raise TooShortUttError( + f"has {xs_pad.size(1)} frames and is too short for subsampling " + + f"(it needs more than {limit_size} frames), return empty results", + xs_pad.size(1), + limit_size, + ) + xs_pad, masks = self.embed(xs_pad, masks) + elif self.embed is not None: + xs_pad = self.embed(xs_pad) + + intermediate_outs = [] + if len(self.interctc_layer_idx) == 0: + if max_layer is not None and 0 <= max_layer < len(self.encoders): + for layer_idx, encoder_layer in enumerate(self.encoders): + xs_pad, masks = encoder_layer(xs_pad, masks) + if layer_idx >= max_layer: + break + else: + xs_pad, masks = self.encoders(xs_pad, masks) + else: + for layer_idx, encoder_layer in enumerate(self.encoders): + xs_pad, masks = encoder_layer(xs_pad, masks) + + if layer_idx + 1 in self.interctc_layer_idx: + encoder_out = xs_pad + + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + intermediate_outs.append((layer_idx + 1, encoder_out)) + + if self.interctc_use_conditioning: + ctc_out = ctc.softmax(encoder_out) + + if isinstance(xs_pad, tuple): + xs_pad = list(xs_pad) + xs_pad[0] = xs_pad[0] + self.conditioning_layer(ctc_out) + xs_pad = tuple(xs_pad) + else: + xs_pad = xs_pad + self.conditioning_layer(ctc_out) + + if isinstance(xs_pad, tuple): + xs_pad = xs_pad[0] + + xs_pad = self.after_norm(xs_pad) + olens = masks.squeeze(1).sum(1) + if len(intermediate_outs) > 0: + return (xs_pad, intermediate_outs), olens, None + return xs_pad, olens, None diff --git a/almeval/models/stepaudio/funasr_detach/models/e_branchformer/model.py b/almeval/models/stepaudio/funasr_detach/models/e_branchformer/model.py new file mode 100644 index 0000000000000000000000000000000000000000..c9bcf394ae61a8a7cb619b86fbe3144092ea9185 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/e_branchformer/model.py @@ -0,0 +1,17 @@ +import logging + +from funasr_detach.models.transformer.model import Transformer +from funasr_detach.register import tables + + +@tables.register("model_classes", "EBranchformer") +class EBranchformer(Transformer): + """CTC-attention hybrid Encoder-Decoder model""" + + def __init__( + self, + *args, + **kwargs, + ): + + super().__init__(*args, **kwargs) diff --git a/almeval/models/stepaudio/funasr_detach/models/e_branchformer/template.yaml b/almeval/models/stepaudio/funasr_detach/models/e_branchformer/template.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dea580e5254f23f1c8510de6a6e56a1b3be7495b --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/e_branchformer/template.yaml @@ -0,0 +1,116 @@ +# This is an example that demonstrates how to configure a model file. +# You can modify the configuration according to your own requirements. + +# to print the register_table: +# from funasr.register import tables +# tables.print() + +# network architecture +model: Branchformer +model_conf: + ctc_weight: 0.3 + lsm_weight: 0.1 # label smoothing option + length_normalized_loss: false + +# encoder +encoder: EBranchformerEncoder +encoder_conf: + output_size: 256 + attention_heads: 4 + attention_layer_type: rel_selfattn + pos_enc_layer_type: rel_pos + rel_pos_type: latest + cgmlp_linear_units: 1024 + cgmlp_conv_kernel: 31 + use_linear_after_conv: false + gate_activation: identity + num_blocks: 12 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + attention_dropout_rate: 0.1 + input_layer: conv2d + layer_drop_rate: 0.0 + linear_units: 1024 + positionwise_layer_type: linear + use_ffn: true + macaron_ffn: true + merge_conv_kernel: 31 + +# decoder +decoder: TransformerDecoder +decoder_conf: + attention_heads: 4 + linear_units: 2048 + num_blocks: 6 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + self_attention_dropout_rate: 0. + src_attention_dropout_rate: 0. + + +# frontend related +frontend: WavFrontend +frontend_conf: + fs: 16000 + window: hamming + n_mels: 80 + frame_length: 25 + frame_shift: 10 + dither: 0.0 + lfr_m: 1 + lfr_n: 1 + +specaug: SpecAug +specaug_conf: + apply_time_warp: true + time_warp_window: 5 + time_warp_mode: bicubic + apply_freq_mask: true + freq_mask_width_range: + - 0 + - 30 + num_freq_mask: 2 + apply_time_mask: true + time_mask_width_range: + - 0 + - 40 + num_time_mask: 2 + +train_conf: + accum_grad: 1 + grad_clip: 5 + max_epoch: 180 + keep_nbest_models: 10 + log_interval: 50 + +optim: adam +optim_conf: + lr: 0.001 + weight_decay: 0.000001 +scheduler: warmuplr +scheduler_conf: + warmup_steps: 35000 + +dataset: AudioDataset +dataset_conf: + index_ds: IndexDSJsonl + batch_sampler: DynamicBatchLocalShuffleSampler + batch_type: example # example or length + batch_size: 1 # if batch_type is example, batch_size is the numbers of samples; if length, batch_size is source_token_len+target_token_len; + max_token_length: 2048 # filter samples if source_token_len+target_token_len > max_token_length, + buffer_size: 500 + shuffle: True + num_workers: 4 + +tokenizer: CharTokenizer +tokenizer_conf: + unk_symbol: + split_with_space: true + + +ctc_conf: + dropout_rate: 0.0 + ctc_type: builtin + reduce: true + ignore_nan_grad: true +normalize: null diff --git a/almeval/models/stepaudio/funasr_detach/models/eend/__init__.py b/almeval/models/stepaudio/funasr_detach/models/eend/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/almeval/models/stepaudio/funasr_detach/models/eend/e2e_diar_eend_ola.py b/almeval/models/stepaudio/funasr_detach/models/eend/e2e_diar_eend_ola.py new file mode 100644 index 0000000000000000000000000000000000000000..2fe7982aece28ea42d1c0f69a9069c2d88be3ed0 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/eend/e2e_diar_eend_ola.py @@ -0,0 +1,272 @@ +from contextlib import contextmanager +from distutils.version import LooseVersion +from typing import Dict, List, Tuple, Optional + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from funasr_detach.frontends.wav_frontend import WavFrontendMel23 +from funasr_detach.models.eend.encoder import EENDOLATransformerEncoder +from funasr_detach.models.eend.encoder_decoder_attractor import EncoderDecoderAttractor +from funasr_detach.models.eend.utils.losses import ( + standard_loss, + cal_power_loss, + fast_batch_pit_n_speaker_loss, +) +from funasr_detach.models.eend.utils.power import create_powerlabel +from funasr_detach.models.eend.utils.power import generate_mapping_dict +from funasr_detach.train_utils.device_funcs import force_gatherable + +if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): + pass +else: + # Nothing to do if torch<1.6.0 + @contextmanager + def autocast(enabled=True): + yield + + +def pad_attractor(att, max_n_speakers): + C, D = att.shape + if C < max_n_speakers: + att = torch.cat( + [att, torch.zeros(max_n_speakers - C, D).to(torch.float32).to(att.device)], + dim=0, + ) + return att + + +def pad_labels(ts, out_size): + for i, t in enumerate(ts): + if t.shape[1] < out_size: + ts[i] = F.pad( + t, (0, out_size - t.shape[1], 0, 0), mode="constant", value=0.0 + ) + return ts + + +def pad_results(ys, out_size): + ys_padded = [] + for i, y in enumerate(ys): + if y.shape[1] < out_size: + ys_padded.append( + torch.cat( + [ + y, + torch.zeros(y.shape[0], out_size - y.shape[1]) + .to(torch.float32) + .to(y.device), + ], + dim=1, + ) + ) + else: + ys_padded.append(y) + return ys_padded + + +class DiarEENDOLAModel(nn.Module): + """EEND-OLA diarization model""" + + def __init__( + self, + frontend: Optional[WavFrontendMel23], + encoder: EENDOLATransformerEncoder, + encoder_decoder_attractor: EncoderDecoderAttractor, + n_units: int = 256, + max_n_speaker: int = 8, + attractor_loss_weight: float = 1.0, + mapping_dict=None, + **kwargs, + ): + super().__init__() + self.frontend = frontend + self.enc = encoder + self.encoder_decoder_attractor = encoder_decoder_attractor + self.attractor_loss_weight = attractor_loss_weight + self.max_n_speaker = max_n_speaker + if mapping_dict is None: + mapping_dict = generate_mapping_dict(max_speaker_num=self.max_n_speaker) + self.mapping_dict = mapping_dict + # PostNet + self.postnet = nn.LSTM(self.max_n_speaker, n_units, 1, batch_first=True) + self.output_layer = nn.Linear(n_units, mapping_dict["oov"] + 1) + + def forward_encoder(self, xs, ilens): + xs = nn.utils.rnn.pad_sequence(xs, batch_first=True, padding_value=-1) + pad_shape = xs.shape + xs_mask = [torch.ones(ilen).to(xs.device) for ilen in ilens] + xs_mask = torch.nn.utils.rnn.pad_sequence( + xs_mask, batch_first=True, padding_value=0 + ).unsqueeze(-2) + emb = self.enc(xs, xs_mask) + emb = torch.split(emb.view(pad_shape[0], pad_shape[1], -1), 1, dim=0) + emb = [e[0][:ilen] for e, ilen in zip(emb, ilens)] + return emb + + def forward_post_net(self, logits, ilens): + maxlen = torch.max(ilens).to(torch.int).item() + logits = nn.utils.rnn.pad_sequence(logits, batch_first=True, padding_value=-1) + logits = nn.utils.rnn.pack_padded_sequence( + logits, ilens.cpu().to(torch.int64), batch_first=True, enforce_sorted=False + ) + outputs, (_, _) = self.postnet(logits) + outputs = nn.utils.rnn.pad_packed_sequence( + outputs, batch_first=True, padding_value=-1, total_length=maxlen + )[0] + outputs = [ + output[: ilens[i].to(torch.int).item()] for i, output in enumerate(outputs) + ] + outputs = [self.output_layer(output) for output in outputs] + return outputs + + def forward( + self, + speech: List[torch.Tensor], + speaker_labels: List[torch.Tensor], + orders: torch.Tensor, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: + + # Check that batch_size is unified + assert len(speech) == len(speaker_labels), (len(speech), len(speaker_labels)) + speech_lengths = torch.tensor([len(sph) for sph in speech]).to(torch.int64) + speaker_labels_lengths = torch.tensor( + [spk.shape[-1] for spk in speaker_labels] + ).to(torch.int64) + batch_size = len(speech) + + # Encoder + encoder_out = self.forward_encoder(speech, speech_lengths) + + # Encoder-decoder attractor + attractor_loss, attractors = self.encoder_decoder_attractor( + [e[order] for e, order in zip(encoder_out, orders)], speaker_labels_lengths + ) + speaker_logits = [ + torch.matmul(e, att.permute(1, 0)) + for e, att in zip(encoder_out, attractors) + ] + + # pit loss + pit_speaker_labels = fast_batch_pit_n_speaker_loss( + speaker_logits, speaker_labels + ) + pit_loss = standard_loss(speaker_logits, pit_speaker_labels) + + # pse loss + with torch.no_grad(): + power_ts = [ + create_powerlabel( + label.cpu().numpy(), self.mapping_dict, self.max_n_speaker + ).to(encoder_out[0].device, non_blocking=True) + for label in pit_speaker_labels + ] + pad_attractors = [pad_attractor(att, self.max_n_speaker) for att in attractors] + pse_speaker_logits = [ + torch.matmul(e, pad_att.permute(1, 0)) + for e, pad_att in zip(encoder_out, pad_attractors) + ] + pse_speaker_logits = self.forward_post_net(pse_speaker_logits, speech_lengths) + pse_loss = cal_power_loss(pse_speaker_logits, power_ts) + + loss = pse_loss + pit_loss + self.attractor_loss_weight * attractor_loss + + stats = dict() + stats["pse_loss"] = pse_loss.detach() + stats["pit_loss"] = pit_loss.detach() + stats["attractor_loss"] = attractor_loss.detach() + stats["batch_size"] = batch_size + + # Collect total loss stats + stats["loss"] = torch.clone(loss.detach()) + + # force_gatherable: to-device and to-tensor if scalar for DataParallel + loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) + return loss, stats, weight + + def estimate_sequential( + self, + speech: torch.Tensor, + n_speakers: int = None, + shuffle: bool = True, + threshold: float = 0.5, + **kwargs, + ): + speech_lengths = torch.tensor([len(sph) for sph in speech]).to(torch.int64) + emb = self.forward_encoder(speech, speech_lengths) + if shuffle: + orders = [np.arange(e.shape[0]) for e in emb] + for order in orders: + np.random.shuffle(order) + attractors, probs = self.encoder_decoder_attractor.estimate( + [ + e[torch.from_numpy(order).to(torch.long).to(speech[0].device)] + for e, order in zip(emb, orders) + ] + ) + else: + attractors, probs = self.encoder_decoder_attractor.estimate(emb) + attractors_active = [] + for p, att, e in zip(probs, attractors, emb): + if n_speakers and n_speakers >= 0: + att = att[:n_speakers,] + attractors_active.append(att) + elif threshold is not None: + silence = torch.nonzero(p < threshold)[0] + n_spk = silence[0] if silence.size else None + att = att[:n_spk,] + attractors_active.append(att) + else: + NotImplementedError("n_speakers or threshold has to be given.") + raw_n_speakers = [att.shape[0] for att in attractors_active] + attractors = [ + ( + pad_attractor(att, self.max_n_speaker) + if att.shape[0] <= self.max_n_speaker + else att[: self.max_n_speaker] + ) + for att in attractors_active + ] + ys = [torch.matmul(e, att.permute(1, 0)) for e, att in zip(emb, attractors)] + logits = self.forward_post_net(ys, speech_lengths) + ys = [ + self.recover_y_from_powerlabel(logit, raw_n_speaker) + for logit, raw_n_speaker in zip(logits, raw_n_speakers) + ] + + return ys, emb, attractors, raw_n_speakers + + def recover_y_from_powerlabel(self, logit, n_speaker): + pred = torch.argmax(torch.softmax(logit, dim=-1), dim=-1) + oov_index = torch.where(pred == self.mapping_dict["oov"])[0] + for i in oov_index: + if i > 0: + pred[i] = pred[i - 1] + else: + pred[i] = 0 + pred = [self.inv_mapping_func(i) for i in pred] + decisions = [bin(num)[2:].zfill(self.max_n_speaker)[::-1] for num in pred] + decisions = ( + torch.from_numpy( + np.stack([np.array([int(i) for i in dec]) for dec in decisions], axis=0) + ) + .to(logit.device) + .to(torch.float32) + ) + decisions = decisions[:, :n_speaker] + return decisions + + def inv_mapping_func(self, label): + + if not isinstance(label, int): + label = int(label) + if label in self.mapping_dict["label2dec"].keys(): + num = self.mapping_dict["label2dec"][label] + else: + num = -1 + return num + + def collect_feats(self, **batch: torch.Tensor) -> Dict[str, torch.Tensor]: + pass diff --git a/almeval/models/stepaudio/funasr_detach/models/eend/eend_ola_dataloader.py b/almeval/models/stepaudio/funasr_detach/models/eend/eend_ola_dataloader.py new file mode 100644 index 0000000000000000000000000000000000000000..983046751f4c1fb1133aef714172b0aa7ad1f06a --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/eend/eend_ola_dataloader.py @@ -0,0 +1,61 @@ +import logging + +import kaldiio +import numpy as np +import torch +from torch.utils.data import DataLoader +from torch.utils.data import Dataset + + +def custom_collate(batch): + keys, speech, speaker_labels, orders = zip(*batch) + speech = [torch.from_numpy(np.copy(sph)).to(torch.float32) for sph in speech] + speaker_labels = [ + torch.from_numpy(np.copy(spk)).to(torch.float32) for spk in speaker_labels + ] + orders = [torch.from_numpy(np.copy(o)).to(torch.int64) for o in orders] + batch = dict(speech=speech, speaker_labels=speaker_labels, orders=orders) + + return keys, batch + + +class EENDOLADataset(Dataset): + def __init__( + self, + data_file, + ): + self.data_file = data_file + with open(data_file) as f: + lines = f.readlines() + self.samples = [line.strip().split() for line in lines] + logging.info("total samples: {}".format(len(self.samples))) + + def __len__(self): + return len(self.samples) + + def __getitem__(self, idx): + key, speech_path, speaker_label_path = self.samples[idx] + speech = kaldiio.load_mat(speech_path) + speaker_label = kaldiio.load_mat(speaker_label_path).reshape( + speech.shape[0], -1 + ) + + order = np.arange(speech.shape[0]) + np.random.shuffle(order) + + return key, speech, speaker_label, order + + +class EENDOLADataLoader: + def __init__(self, data_file, batch_size, shuffle=True, num_workers=8): + dataset = EENDOLADataset(data_file) + self.data_loader = DataLoader( + dataset, + batch_size=batch_size, + collate_fn=custom_collate, + shuffle=shuffle, + num_workers=num_workers, + ) + + def build_iter(self, epoch): + return self.data_loader diff --git a/almeval/models/stepaudio/funasr_detach/models/eend/encoder.py b/almeval/models/stepaudio/funasr_detach/models/eend/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..4a8039525c2e1219fad71923d5b87eec416679ca --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/eend/encoder.py @@ -0,0 +1,126 @@ +import math + +import torch +import torch.nn.functional as F +from torch import nn + + +class MultiHeadSelfAttention(nn.Module): + def __init__(self, n_units, h=8, dropout_rate=0.1): + super().__init__() + self.linearQ = nn.Linear(n_units, n_units) + self.linearK = nn.Linear(n_units, n_units) + self.linearV = nn.Linear(n_units, n_units) + self.linearO = nn.Linear(n_units, n_units) + self.d_k = n_units // h + self.h = h + self.dropout = nn.Dropout(dropout_rate) + + def __call__(self, x, batch_size, x_mask): + q = self.linearQ(x).view(batch_size, -1, self.h, self.d_k) + k = self.linearK(x).view(batch_size, -1, self.h, self.d_k) + v = self.linearV(x).view(batch_size, -1, self.h, self.d_k) + scores = torch.matmul(q.permute(0, 2, 1, 3), k.permute(0, 2, 3, 1)) / math.sqrt( + self.d_k + ) + if x_mask is not None: + x_mask = x_mask.unsqueeze(1) + scores = scores.masked_fill(x_mask == 0, -1e9) + self.att = F.softmax(scores, dim=3) + p_att = self.dropout(self.att) + x = torch.matmul(p_att, v.permute(0, 2, 1, 3)) + x = x.permute(0, 2, 1, 3).contiguous().view(-1, self.h * self.d_k) + return self.linearO(x) + + +class PositionwiseFeedForward(nn.Module): + def __init__(self, n_units, d_units, dropout_rate): + super(PositionwiseFeedForward, self).__init__() + self.linear1 = nn.Linear(n_units, d_units) + self.linear2 = nn.Linear(d_units, n_units) + self.dropout = nn.Dropout(dropout_rate) + + def __call__(self, x): + return self.linear2(self.dropout(F.relu(self.linear1(x)))) + + +class PositionalEncoding(torch.nn.Module): + def __init__(self, d_model, dropout_rate, max_len=5000, reverse=False): + super(PositionalEncoding, self).__init__() + self.d_model = d_model + self.reverse = reverse + self.xscale = math.sqrt(self.d_model) + self.dropout = torch.nn.Dropout(p=dropout_rate) + self.pe = None + self.extend_pe(torch.tensor(0.0).expand(1, max_len)) + + def extend_pe(self, x): + if self.pe is not None: + if self.pe.size(1) >= x.size(1): + if self.pe.dtype != x.dtype or self.pe.device != x.device: + self.pe = self.pe.to(dtype=x.dtype, device=x.device) + return + pe = torch.zeros(x.size(1), self.d_model) + if self.reverse: + position = torch.arange( + x.size(1) - 1, -1, -1.0, dtype=torch.float32 + ).unsqueeze(1) + else: + position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1) + div_term = torch.exp( + torch.arange(0, self.d_model, 2, dtype=torch.float32) + * -(math.log(10000.0) / self.d_model) + ) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + pe = pe.unsqueeze(0) + self.pe = pe.to(device=x.device, dtype=x.dtype) + + def forward(self, x: torch.Tensor): + self.extend_pe(x) + x = x * self.xscale + self.pe[:, : x.size(1)] + return self.dropout(x) + + +class EENDOLATransformerEncoder(nn.Module): + def __init__( + self, + idim: int, + n_layers: int, + n_units: int, + e_units: int = 2048, + h: int = 4, + dropout_rate: float = 0.1, + use_pos_emb: bool = False, + ): + super(EENDOLATransformerEncoder, self).__init__() + self.linear_in = nn.Linear(idim, n_units) + self.lnorm_in = nn.LayerNorm(n_units) + self.n_layers = n_layers + self.dropout = nn.Dropout(dropout_rate) + for i in range(n_layers): + setattr(self, "{}{:d}".format("lnorm1_", i), nn.LayerNorm(n_units)) + setattr( + self, + "{}{:d}".format("self_att_", i), + MultiHeadSelfAttention(n_units, h), + ) + setattr(self, "{}{:d}".format("lnorm2_", i), nn.LayerNorm(n_units)) + setattr( + self, + "{}{:d}".format("ff_", i), + PositionwiseFeedForward(n_units, e_units, dropout_rate), + ) + self.lnorm_out = nn.LayerNorm(n_units) + + def __call__(self, x, x_mask=None): + BT_size = x.shape[0] * x.shape[1] + e = self.linear_in(x.reshape(BT_size, -1)) + for i in range(self.n_layers): + e = getattr(self, "{}{:d}".format("lnorm1_", i))(e) + s = getattr(self, "{}{:d}".format("self_att_", i))(e, x.shape[0], x_mask) + e = e + self.dropout(s) + e = getattr(self, "{}{:d}".format("lnorm2_", i))(e) + s = getattr(self, "{}{:d}".format("ff_", i))(e) + e = e + self.dropout(s) + return self.lnorm_out(e) diff --git a/almeval/models/stepaudio/funasr_detach/models/eend/encoder_decoder_attractor.py b/almeval/models/stepaudio/funasr_detach/models/eend/encoder_decoder_attractor.py new file mode 100644 index 0000000000000000000000000000000000000000..d6935ebd3bada28425ea0bd1639bd09e5ab976c4 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/eend/encoder_decoder_attractor.py @@ -0,0 +1,79 @@ +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn + + +class EncoderDecoderAttractor(nn.Module): + + def __init__(self, n_units, encoder_dropout=0.1, decoder_dropout=0.1): + super(EncoderDecoderAttractor, self).__init__() + self.enc0_dropout = nn.Dropout(encoder_dropout) + self.encoder = nn.LSTM( + n_units, n_units, 1, batch_first=True, dropout=encoder_dropout + ) + self.dec0_dropout = nn.Dropout(decoder_dropout) + self.decoder = nn.LSTM( + n_units, n_units, 1, batch_first=True, dropout=decoder_dropout + ) + self.counter = nn.Linear(n_units, 1) + self.n_units = n_units + + def forward_core(self, xs, zeros): + ilens = torch.from_numpy(np.array([x.shape[0] for x in xs])).to(torch.int64) + xs = [self.enc0_dropout(x) for x in xs] + xs = nn.utils.rnn.pad_sequence(xs, batch_first=True, padding_value=-1) + xs = nn.utils.rnn.pack_padded_sequence( + xs, ilens, batch_first=True, enforce_sorted=False + ) + _, (hx, cx) = self.encoder(xs) + zlens = torch.from_numpy(np.array([z.shape[0] for z in zeros])).to(torch.int64) + max_zlen = torch.max(zlens).to(torch.int).item() + zeros = [self.enc0_dropout(z) for z in zeros] + zeros = nn.utils.rnn.pad_sequence(zeros, batch_first=True, padding_value=-1) + zeros = nn.utils.rnn.pack_padded_sequence( + zeros, zlens, batch_first=True, enforce_sorted=False + ) + attractors, (_, _) = self.decoder(zeros, (hx, cx)) + attractors = nn.utils.rnn.pad_packed_sequence( + attractors, batch_first=True, padding_value=-1, total_length=max_zlen + )[0] + attractors = [ + att[: zlens[i].to(torch.int).item()] for i, att in enumerate(attractors) + ] + return attractors + + def forward(self, xs, n_speakers): + zeros = [ + torch.zeros(n_spk + 1, self.n_units).to(torch.float32).to(xs[0].device) + for n_spk in n_speakers + ] + attractors = self.forward_core(xs, zeros) + labels = torch.cat( + [ + torch.from_numpy(np.array([[1] * n_spk + [0]], np.float32)) + for n_spk in n_speakers + ], + dim=1, + ) + labels = labels.to(xs[0].device) + logit = torch.cat( + [ + self.counter(att).view(-1, n_spk + 1) + for att, n_spk in zip(attractors, n_speakers) + ], + dim=1, + ) + loss = F.binary_cross_entropy(torch.sigmoid(logit), labels) + + attractors = [att[slice(0, att.shape[0] - 1)] for att in attractors] + return loss, attractors + + def estimate(self, xs, max_n_speakers=15): + zeros = [ + torch.zeros(max_n_speakers, self.n_units).to(torch.float32).to(xs[0].device) + for _ in xs + ] + attractors = self.forward_core(xs, zeros) + probs = [torch.sigmoid(torch.flatten(self.counter(att))) for att in attractors] + return attractors, probs diff --git a/almeval/models/stepaudio/funasr_detach/models/eend/utils/__init__.py b/almeval/models/stepaudio/funasr_detach/models/eend/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/almeval/models/stepaudio/funasr_detach/models/eend/utils/feature.py b/almeval/models/stepaudio/funasr_detach/models/eend/utils/feature.py new file mode 100644 index 0000000000000000000000000000000000000000..d3e6cc59e2cb4ce1f8e04127975baedb009284a4 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/eend/utils/feature.py @@ -0,0 +1,275 @@ +# Copyright 2019 Hitachi, Ltd. (author: Yusuke Fujita) +# Licensed under the MIT license. +# +# This module is for computing audio features + +import numpy as np +import librosa + + +def get_input_dim( + frame_size, + context_size, + transform_type, +): + if transform_type.startswith("logmel23"): + frame_size = 23 + elif transform_type.startswith("logmel"): + frame_size = 40 + else: + fft_size = 1 << (frame_size - 1).bit_length() + frame_size = int(fft_size / 2) + 1 + input_dim = (2 * context_size + 1) * frame_size + return input_dim + + +def transform(Y, transform_type=None, dtype=np.float32): + """Transform STFT feature + + Args: + Y: STFT + (n_frames, n_bins)-shaped np.complex array + transform_type: + None, "log" + dtype: output data type + np.float32 is expected + Returns: + Y (numpy.array): transformed feature + """ + Y = np.abs(Y) + if not transform_type: + pass + elif transform_type == "log": + Y = np.log(np.maximum(Y, 1e-10)) + elif transform_type == "logmel": + n_fft = 2 * (Y.shape[1] - 1) + sr = 16000 + n_mels = 40 + mel_basis = librosa.filters.mel(sr, n_fft, n_mels) + Y = np.dot(Y**2, mel_basis.T) + Y = np.log10(np.maximum(Y, 1e-10)) + elif transform_type == "logmel23": + n_fft = 2 * (Y.shape[1] - 1) + sr = 8000 + n_mels = 23 + mel_basis = librosa.filters.mel(sr, n_fft, n_mels) + Y = np.dot(Y**2, mel_basis.T) + Y = np.log10(np.maximum(Y, 1e-10)) + elif transform_type == "logmel23_mn": + n_fft = 2 * (Y.shape[1] - 1) + sr = 8000 + n_mels = 23 + mel_basis = librosa.filters.mel(sr, n_fft, n_mels) + Y = np.dot(Y**2, mel_basis.T) + Y = np.log10(np.maximum(Y, 1e-10)) + mean = np.mean(Y, axis=0) + Y = Y - mean + elif transform_type == "logmel23_swn": + n_fft = 2 * (Y.shape[1] - 1) + sr = 8000 + n_mels = 23 + mel_basis = librosa.filters.mel(sr, n_fft, n_mels) + Y = np.dot(Y**2, mel_basis.T) + Y = np.log10(np.maximum(Y, 1e-10)) + # b = np.ones(300)/300 + # mean = scipy.signal.convolve2d(Y, b[:, None], mode='same') + + # simple 2-means based threshoding for mean calculation + powers = np.sum(Y, axis=1) + th = (np.max(powers) + np.min(powers)) / 2.0 + for i in range(10): + th = (np.mean(powers[powers >= th]) + np.mean(powers[powers < th])) / 2 + mean = np.mean(Y[powers > th, :], axis=0) + Y = Y - mean + elif transform_type == "logmel23_mvn": + n_fft = 2 * (Y.shape[1] - 1) + sr = 8000 + n_mels = 23 + mel_basis = librosa.filters.mel(sr, n_fft, n_mels) + Y = np.dot(Y**2, mel_basis.T) + Y = np.log10(np.maximum(Y, 1e-10)) + mean = np.mean(Y, axis=0) + Y = Y - mean + std = np.maximum(np.std(Y, axis=0), 1e-10) + Y = Y / std + else: + raise ValueError("Unknown transform_type: %s" % transform_type) + return Y.astype(dtype) + + +def subsample(Y, T, subsampling=1): + """Frame subsampling""" + Y_ss = Y[::subsampling] + T_ss = T[::subsampling] + return Y_ss, T_ss + + +def splice(Y, context_size=0): + """Frame splicing + + Args: + Y: feature + (n_frames, n_featdim)-shaped numpy array + context_size: + number of frames concatenated on left-side + if context_size = 5, 11 frames are concatenated. + + Returns: + Y_spliced: spliced feature + (n_frames, n_featdim * (2 * context_size + 1))-shaped + """ + Y_pad = np.pad(Y, [(context_size, context_size), (0, 0)], "constant") + Y_spliced = np.lib.stride_tricks.as_strided( + np.ascontiguousarray(Y_pad), + (Y.shape[0], Y.shape[1] * (2 * context_size + 1)), + (Y.itemsize * Y.shape[1], Y.itemsize), + writeable=False, + ) + return Y_spliced + + +def stft(data, frame_size=1024, frame_shift=256): + """Compute STFT features + + Args: + data: audio signal + (n_samples,)-shaped np.float32 array + frame_size: number of samples in a frame (must be a power of two) + frame_shift: number of samples between frames + + Returns: + stft: STFT frames + (n_frames, n_bins)-shaped np.complex64 array + """ + # round up to nearest power of 2 + fft_size = 1 << (frame_size - 1).bit_length() + # HACK: The last frame is ommited + # as librosa.stft produces such an excessive frame + if len(data) % frame_shift == 0: + return librosa.stft( + data, n_fft=fft_size, win_length=frame_size, hop_length=frame_shift + ).T[:-1] + else: + return librosa.stft( + data, n_fft=fft_size, win_length=frame_size, hop_length=frame_shift + ).T + + +def _count_frames(data_len, size, shift): + # HACK: Assuming librosa.stft(..., center=True) + n_frames = 1 + int(data_len / shift) + if data_len % shift == 0: + n_frames = n_frames - 1 + return n_frames + + +def get_frame_labels( + kaldi_obj, rec, start=0, end=None, frame_size=1024, frame_shift=256, n_speakers=None +): + """Get frame-aligned labels of given recording + Args: + kaldi_obj (KaldiData) + rec (str): recording id + start (int): start frame index + end (int): end frame index + None means the last frame of recording + frame_size (int): number of frames in a frame + frame_shift (int): number of shift samples + n_speakers (int): number of speakers + if None, the value is given from data + Returns: + T: label + (n_frames, n_speakers)-shaped np.int32 array + """ + filtered_segments = kaldi_obj.segments[kaldi_obj.segments["rec"] == rec] + speakers = np.unique( + [kaldi_obj.utt2spk[seg["utt"]] for seg in filtered_segments] + ).tolist() + if n_speakers is None: + n_speakers = len(speakers) + es = end * frame_shift if end is not None else None + data, rate = kaldi_obj.load_wav(rec, start * frame_shift, es) + n_frames = _count_frames(len(data), frame_size, frame_shift) + T = np.zeros((n_frames, n_speakers), dtype=np.int32) + if end is None: + end = n_frames + + for seg in filtered_segments: + speaker_index = speakers.index(kaldi_obj.utt2spk[seg["utt"]]) + start_frame = np.rint(seg["st"] * rate / frame_shift).astype(int) + end_frame = np.rint(seg["et"] * rate / frame_shift).astype(int) + rel_start = rel_end = None + if start <= start_frame and start_frame < end: + rel_start = start_frame - start + if start < end_frame and end_frame <= end: + rel_end = end_frame - start + if rel_start is not None or rel_end is not None: + T[rel_start:rel_end, speaker_index] = 1 + return T + + +def get_labeledSTFT( + kaldi_obj, + rec, + start, + end, + frame_size, + frame_shift, + n_speakers=None, + use_speaker_id=False, +): + """Extracts STFT and corresponding labels + + Extracts STFT and corresponding diarization labels for + given recording id and start/end times + + Args: + kaldi_obj (KaldiData) + rec (str): recording id + start (int): start frame index + end (int): end frame index + frame_size (int): number of samples in a frame + frame_shift (int): number of shift samples + n_speakers (int): number of speakers + if None, the value is given from data + Returns: + Y: STFT + (n_frames, n_bins)-shaped np.complex64 array, + T: label + (n_frmaes, n_speakers)-shaped np.int32 array. + """ + data, rate = kaldi_obj.load_wav(rec, start * frame_shift, end * frame_shift) + Y = stft(data, frame_size, frame_shift) + filtered_segments = kaldi_obj.segments[rec] + # filtered_segments = kaldi_obj.segments[kaldi_obj.segments['rec'] == rec] + speakers = np.unique( + [kaldi_obj.utt2spk[seg["utt"]] for seg in filtered_segments] + ).tolist() + if n_speakers is None: + n_speakers = len(speakers) + T = np.zeros((Y.shape[0], n_speakers), dtype=np.int32) + + if use_speaker_id: + all_speakers = sorted(kaldi_obj.spk2utt.keys()) + S = np.zeros((Y.shape[0], len(all_speakers)), dtype=np.int32) + + for seg in filtered_segments: + speaker_index = speakers.index(kaldi_obj.utt2spk[seg["utt"]]) + if use_speaker_id: + all_speaker_index = all_speakers.index(kaldi_obj.utt2spk[seg["utt"]]) + start_frame = np.rint(seg["st"] * rate / frame_shift).astype(int) + end_frame = np.rint(seg["et"] * rate / frame_shift).astype(int) + rel_start = rel_end = None + if start <= start_frame and start_frame < end: + rel_start = start_frame - start + if start < end_frame and end_frame <= end: + rel_end = end_frame - start + if rel_start is not None or rel_end is not None: + T[rel_start:rel_end, speaker_index] = 1 + if use_speaker_id: + S[rel_start:rel_end, all_speaker_index] = 1 + + if use_speaker_id: + return Y, T, S + else: + return Y, T diff --git a/almeval/models/stepaudio/funasr_detach/models/eend/utils/kaldi_data.py b/almeval/models/stepaudio/funasr_detach/models/eend/utils/kaldi_data.py new file mode 100644 index 0000000000000000000000000000000000000000..59e7a16f8b55d7c27ad7bd0075b6a66338d8bbb2 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/eend/utils/kaldi_data.py @@ -0,0 +1,152 @@ +# Copyright 2019 Hitachi, Ltd. (author: Yusuke Fujita) +# Licensed under the MIT license. +# +# This library provides utilities for kaldi-style data directory. + + +from __future__ import print_function +import os +import sys +import numpy as np +import subprocess +import librosa as sf +import io +from functools import lru_cache + + +def load_segments(segments_file): + """load segments file as array""" + if not os.path.exists(segments_file): + return None + return np.loadtxt( + segments_file, + dtype=[("utt", "object"), ("rec", "object"), ("st", "f"), ("et", "f")], + ndmin=1, + ) + + +def load_segments_hash(segments_file): + ret = {} + if not os.path.exists(segments_file): + return None + for line in open(segments_file): + utt, rec, st, et = line.strip().split() + ret[utt] = (rec, float(st), float(et)) + return ret + + +def load_segments_rechash(segments_file): + ret = {} + if not os.path.exists(segments_file): + return None + for line in open(segments_file): + utt, rec, st, et = line.strip().split() + if rec not in ret: + ret[rec] = [] + ret[rec].append({"utt": utt, "st": float(st), "et": float(et)}) + return ret + + +def load_wav_scp(wav_scp_file): + """return dictionary { rec: wav_rxfilename }""" + lines = [line.strip().split(None, 1) for line in open(wav_scp_file)] + return {x[0]: x[1] for x in lines} + + +@lru_cache(maxsize=1) +def load_wav(wav_rxfilename, start=0, end=None): + """This function reads audio file and return data in numpy.float32 array. + "lru_cache" holds recently loaded audio so that can be called + many times on the same audio file. + OPTIMIZE: controls lru_cache size for random access, + considering memory size + """ + if wav_rxfilename.endswith("|"): + # input piped command + p = subprocess.Popen(wav_rxfilename[:-1], shell=True, stdout=subprocess.PIPE) + data, samplerate = sf.load(io.BytesIO(p.stdout.read()), dtype="float32") + # cannot seek + data = data[start:end] + elif wav_rxfilename == "-": + # stdin + data, samplerate = sf.load(sys.stdin, dtype="float32") + # cannot seek + data = data[start:end] + else: + # normal wav file + data, samplerate = sf.load(wav_rxfilename, start=start, stop=end) + return data, samplerate + + +def load_utt2spk(utt2spk_file): + """returns dictionary { uttid: spkid }""" + lines = [line.strip().split(None, 1) for line in open(utt2spk_file)] + return {x[0]: x[1] for x in lines} + + +def load_spk2utt(spk2utt_file): + """returns dictionary { spkid: list of uttids }""" + if not os.path.exists(spk2utt_file): + return None + lines = [line.strip().split() for line in open(spk2utt_file)] + return {x[0]: x[1:] for x in lines} + + +def load_reco2dur(reco2dur_file): + """returns dictionary { recid: duration }""" + if not os.path.exists(reco2dur_file): + return None + lines = [line.strip().split(None, 1) for line in open(reco2dur_file)] + return {x[0]: float(x[1]) for x in lines} + + +def process_wav(wav_rxfilename, process): + """This function returns preprocessed wav_rxfilename + Args: + wav_rxfilename: input + process: command which can be connected via pipe, + use stdin and stdout + Returns: + wav_rxfilename: output piped command + """ + if wav_rxfilename.endswith("|"): + # input piped command + return wav_rxfilename + process + "|" + else: + # stdin "-" or normal file + return "cat {} | {} |".format(wav_rxfilename, process) + + +def extract_segments(wavs, segments=None): + """This function returns generator of segmented audio as + (utterance id, numpy.float32 array) + TODO?: sampling rate is not converted. + """ + if segments is not None: + # segments should be sorted by rec-id + for seg in segments: + wav = wavs[seg["rec"]] + data, samplerate = load_wav(wav) + st_sample = np.rint(seg["st"] * samplerate).astype(int) + et_sample = np.rint(seg["et"] * samplerate).astype(int) + yield seg["utt"], data[st_sample:et_sample] + else: + # segments file not found, + # wav.scp is used as segmented audio list + for rec in wavs: + data, samplerate = load_wav(wavs[rec]) + yield rec, data + + +class KaldiData: + def __init__(self, data_dir): + self.data_dir = data_dir + self.segments = load_segments_rechash(os.path.join(self.data_dir, "segments")) + self.utt2spk = load_utt2spk(os.path.join(self.data_dir, "utt2spk")) + self.wavs = load_wav_scp(os.path.join(self.data_dir, "wav.scp")) + self.reco2dur = load_reco2dur(os.path.join(self.data_dir, "reco2dur")) + self.spk2utt = load_spk2utt(os.path.join(self.data_dir, "spk2utt")) + + def load_wav(self, recid, start=0, end=None): + data, rate = load_wav(self.wavs[recid], start, end) + return data, rate diff --git a/almeval/models/stepaudio/funasr_detach/models/eend/utils/losses.py b/almeval/models/stepaudio/funasr_detach/models/eend/utils/losses.py new file mode 100644 index 0000000000000000000000000000000000000000..957e8b291784edbc3b9e2a5076b8156206c0560d --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/eend/utils/losses.py @@ -0,0 +1,53 @@ +import numpy as np +import torch +import torch.nn.functional as F +from scipy.optimize import linear_sum_assignment + + +def standard_loss(ys, ts): + losses = [ + F.binary_cross_entropy(torch.sigmoid(y), t) * len(y) for y, t in zip(ys, ts) + ] + loss = torch.sum(torch.stack(losses)) + n_frames = ( + torch.from_numpy(np.array(np.sum([t.shape[0] for t in ts]))) + .to(torch.float32) + .to(ys[0].device) + ) + loss = loss / n_frames + return loss + + +def fast_batch_pit_n_speaker_loss(ys, ts): + with torch.no_grad(): + bs = len(ys) + indices = [] + for b in range(bs): + y = ys[b].transpose(0, 1) + t = ts[b].transpose(0, 1) + C, _ = t.shape + y = y[:, None, :].repeat(1, C, 1) + t = t[None, :, :].repeat(C, 1, 1) + bce_loss = F.binary_cross_entropy( + torch.sigmoid(y), t, reduction="none" + ).mean(-1) + C = bce_loss.cpu() + indices.append(linear_sum_assignment(C)) + labels_perm = [t[:, idx[1]] for t, idx in zip(ts, indices)] + + return labels_perm + + +def cal_power_loss(logits, power_ts): + losses = [ + F.cross_entropy(input=logit, target=power_t.to(torch.long)) * len(logit) + for logit, power_t in zip(logits, power_ts) + ] + loss = torch.sum(torch.stack(losses)) + n_frames = ( + torch.from_numpy(np.array(np.sum([power_t.shape[0] for power_t in power_ts]))) + .to(torch.float32) + .to(power_ts[0].device) + ) + loss = loss / n_frames + return loss diff --git a/almeval/models/stepaudio/funasr_detach/models/eend/utils/power.py b/almeval/models/stepaudio/funasr_detach/models/eend/utils/power.py new file mode 100644 index 0000000000000000000000000000000000000000..b0b3203bb90302d85904acc590ce3bae3330208b --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/eend/utils/power.py @@ -0,0 +1,114 @@ +import numpy as np +import torch +import torch.multiprocessing +import torch.nn.functional as F +from itertools import combinations +from itertools import permutations + + +def generate_mapping_dict(max_speaker_num=6, max_olp_speaker_num=3): + all_kinds = [] + all_kinds.append(0) + for i in range(max_olp_speaker_num): + selected_num = i + 1 + coms = np.array(list(combinations(np.arange(max_speaker_num), selected_num))) + for com in coms: + tmp = np.zeros(max_speaker_num) + tmp[com] = 1 + item = int(raw_dec_trans(tmp.reshape(1, -1), max_speaker_num)[0]) + all_kinds.append(item) + all_kinds_order = sorted(all_kinds) + + mapping_dict = {} + mapping_dict["dec2label"] = {} + mapping_dict["label2dec"] = {} + for i in range(len(all_kinds_order)): + dec = all_kinds_order[i] + mapping_dict["dec2label"][dec] = i + mapping_dict["label2dec"][i] = dec + oov_id = len(all_kinds_order) + mapping_dict["oov"] = oov_id + return mapping_dict + + +def raw_dec_trans(x, max_speaker_num): + num_list = [] + for i in range(max_speaker_num): + num_list.append(x[:, i]) + base = 1 + T = x.shape[0] + res = np.zeros((T)) + for num in num_list: + res += num * base + base = base * 2 + return res + + +def mapping_func(num, mapping_dict): + if num in mapping_dict["dec2label"].keys(): + label = mapping_dict["dec2label"][num] + else: + label = mapping_dict["oov"] + return label + + +def dec_trans(x, max_speaker_num, mapping_dict): + num_list = [] + for i in range(max_speaker_num): + num_list.append(x[:, i]) + base = 1 + T = x.shape[0] + res = np.zeros((T)) + for num in num_list: + res += num * base + base = base * 2 + res = np.array([mapping_func(i, mapping_dict) for i in res]) + return res + + +def create_powerlabel(label, mapping_dict, max_speaker_num=6, max_olp_speaker_num=3): + T, C = label.shape + padding_label = np.zeros((T, max_speaker_num)) + padding_label[:, :C] = label + out_label = dec_trans(padding_label, max_speaker_num, mapping_dict) + out_label = torch.from_numpy(out_label) + return out_label + + +def generate_perm_pse( + label, n_speaker, mapping_dict, max_speaker_num, max_olp_speaker_num=3 +): + perms = np.array(list(permutations(range(n_speaker)))).astype(np.float32) + perms = torch.from_numpy(perms).to(label.device).to(torch.int64) + perm_labels = [label[:, perm] for perm in perms] + perm_pse_labels = [ + create_powerlabel(perm_label.cpu().numpy(), mapping_dict, max_speaker_num).to( + perm_label.device, non_blocking=True + ) + for perm_label in perm_labels + ] + return perm_labels, perm_pse_labels + + +def generate_min_pse( + label, n_speaker, mapping_dict, max_speaker_num, pse_logit, max_olp_speaker_num=3 +): + perm_labels, perm_pse_labels = generate_perm_pse( + label, + n_speaker, + mapping_dict, + max_speaker_num, + max_olp_speaker_num=max_olp_speaker_num, + ) + losses = [ + F.cross_entropy(input=pse_logit, target=perm_pse_label.to(torch.long)) + * len(pse_logit) + for perm_pse_label in perm_pse_labels + ] + loss = torch.stack(losses) + min_index = torch.argmin(loss) + selected_perm_label, selected_pse_label = ( + perm_labels[min_index], + perm_pse_labels[min_index], + ) + return selected_perm_label, selected_pse_label diff --git a/almeval/models/stepaudio/funasr_detach/models/eend/utils/report.py b/almeval/models/stepaudio/funasr_detach/models/eend/utils/report.py new file mode 100644 index 0000000000000000000000000000000000000000..23382072427580b7b1cf168653ae59bc7fd983a8 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/eend/utils/report.py @@ -0,0 +1,201 @@ +import copy +import numpy as np +import time +import torch +from funasr_detach.models.eend.utils.power import create_powerlabel +from itertools import combinations + +metrics = [ + ("diarization_error", "speaker_scored", "DER"), + ("speech_miss", "speech_scored", "SAD_MR"), + ("speech_falarm", "speech_scored", "SAD_FR"), + ("speaker_miss", "speaker_scored", "MI"), + ("speaker_falarm", "speaker_scored", "FA"), + ("speaker_error", "speaker_scored", "CF"), + ("correct", "frames", "accuracy"), +] + + +def recover_prediction(y, n_speaker): + if n_speaker <= 1: + return y + elif n_speaker == 2: + com_index = torch.from_numpy( + np.array(list(combinations(np.arange(n_speaker), 2))) + ).to(y.dtype) + num_coms = com_index.shape[0] + y_single = y[:, :-num_coms] + y_olp = y[:, -num_coms:] + olp_map_index = torch.where(y_olp > 0.5) + olp_map_index = torch.stack(olp_map_index, dim=1) + com_map_index = com_index[olp_map_index[:, -1]] + speaker_map_index = ( + torch.from_numpy(np.array(com_map_index)).view(-1).to(torch.int64) + ) + frame_map_index = ( + olp_map_index[:, 0][:, None].repeat([1, 2]).view(-1).to(torch.int64) + ) + y_single[frame_map_index] = 0 + y_single[frame_map_index, speaker_map_index] = 1 + return y_single + else: + olp2_com_index = torch.from_numpy( + np.array(list(combinations(np.arange(n_speaker), 2))) + ).to(y.dtype) + olp2_num_coms = olp2_com_index.shape[0] + olp3_com_index = torch.from_numpy( + np.array(list(combinations(np.arange(n_speaker), 3))) + ).to(y.dtype) + olp3_num_coms = olp3_com_index.shape[0] + y_single = y[:, :n_speaker] + y_olp2 = y[:, n_speaker : n_speaker + olp2_num_coms] + y_olp3 = y[:, -olp3_num_coms:] + + olp3_map_index = torch.where(y_olp3 > 0.5) + olp3_map_index = torch.stack(olp3_map_index, dim=1) + olp3_com_map_index = olp3_com_index[olp3_map_index[:, -1]] + olp3_speaker_map_index = ( + torch.from_numpy(np.array(olp3_com_map_index)).view(-1).to(torch.int64) + ) + olp3_frame_map_index = ( + olp3_map_index[:, 0][:, None].repeat([1, 3]).view(-1).to(torch.int64) + ) + y_single[olp3_frame_map_index] = 0 + y_single[olp3_frame_map_index, olp3_speaker_map_index] = 1 + y_olp2[olp3_frame_map_index] = 0 + + olp2_map_index = torch.where(y_olp2 > 0.5) + olp2_map_index = torch.stack(olp2_map_index, dim=1) + olp2_com_map_index = olp2_com_index[olp2_map_index[:, -1]] + olp2_speaker_map_index = ( + torch.from_numpy(np.array(olp2_com_map_index)).view(-1).to(torch.int64) + ) + olp2_frame_map_index = ( + olp2_map_index[:, 0][:, None].repeat([1, 2]).view(-1).to(torch.int64) + ) + y_single[olp2_frame_map_index] = 0 + y_single[olp2_frame_map_index, olp2_speaker_map_index] = 1 + return y_single + + +class PowerReporter: + def __init__(self, valid_data_loader, mapping_dict, max_n_speaker): + valid_data_loader_cp = copy.deepcopy(valid_data_loader) + self.valid_data_loader = valid_data_loader_cp + del valid_data_loader + self.mapping_dict = mapping_dict + self.max_n_speaker = max_n_speaker + + def report(self, model, eidx, device): + self.report_val(model, eidx, device) + + def report_val(self, model, eidx, device): + model.eval() + ud_valid_start = time.time() + valid_res, valid_loss, stats_keys, vad_valid_accuracy = self.report_core( + model, self.valid_data_loader, device + ) + + # Epoch Display + valid_der = valid_res["diarization_error"] / valid_res["speaker_scored"] + valid_accuracy = ( + valid_res["correct"].to(torch.float32) / valid_res["frames"] * 100 + ) + vad_valid_accuracy = vad_valid_accuracy * 100 + print( + "Epoch ", + eidx + 1, + "Valid Loss ", + valid_loss, + "Valid_DER %.5f" % valid_der, + "Valid_Accuracy %.5f%% " % valid_accuracy, + "VAD_Valid_Accuracy %.5f%% " % vad_valid_accuracy, + ) + ud_valid = (time.time() - ud_valid_start) / 60.0 + print("Valid cost time ... ", ud_valid) + + def inv_mapping_func(self, label, mapping_dict): + if not isinstance(label, int): + label = int(label) + if label in mapping_dict["label2dec"].keys(): + num = mapping_dict["label2dec"][label] + else: + num = -1 + return num + + def report_core(self, model, data_loader, device): + res = {} + for item in metrics: + res[item[0]] = 0.0 + res[item[1]] = 0.0 + with torch.no_grad(): + loss_s = 0.0 + uidx = 0 + for xs, ts, orders in data_loader: + xs = [x.to(device) for x in xs] + ts = [t.to(device) for t in ts] + orders = [o.to(device) for o in orders] + loss, pit_loss, mpit_loss, att_loss, ys, logits, labels, attractors = ( + model(xs, ts, orders) + ) + loss_s += loss.item() + uidx += 1 + + for logit, t, att in zip(logits, labels, attractors): + pred = torch.argmax(torch.softmax(logit, dim=-1), dim=-1) # (T, ) + oov_index = torch.where(pred == self.mapping_dict["oov"])[0] + for i in oov_index: + if i > 0: + pred[i] = pred[i - 1] + else: + pred[i] = 0 + pred = [self.inv_mapping_func(i, self.mapping_dict) for i in pred] + decisions = [ + bin(num)[2:].zfill(self.max_n_speaker)[::-1] for num in pred + ] + decisions = ( + torch.from_numpy( + np.stack( + [np.array([int(i) for i in dec]) for dec in decisions], + axis=0, + ) + ) + .to(att.device) + .to(torch.float32) + ) + decisions = decisions[:, : att.shape[0]] + + stats = self.calc_diarization_error(decisions, t) + res["speaker_scored"] += stats["speaker_scored"] + res["speech_scored"] += stats["speech_scored"] + res["frames"] += stats["frames"] + for item in metrics: + res[item[0]] += stats[item[0]] + loss_s /= uidx + vad_acc = 0 + + return res, loss_s, stats.keys(), vad_acc + + def calc_diarization_error(self, decisions, label, label_delay=0): + label = label[: len(label) - label_delay, ...] + n_ref = torch.sum(label, dim=-1) + n_sys = torch.sum(decisions, dim=-1) + res = {} + res["speech_scored"] = torch.sum(n_ref > 0) + res["speech_miss"] = torch.sum((n_ref > 0) & (n_sys == 0)) + res["speech_falarm"] = torch.sum((n_ref == 0) & (n_sys > 0)) + res["speaker_scored"] = torch.sum(n_ref) + res["speaker_miss"] = torch.sum( + torch.max(n_ref - n_sys, torch.zeros_like(n_ref)) + ) + res["speaker_falarm"] = torch.sum( + torch.max(n_sys - n_ref, torch.zeros_like(n_ref)) + ) + n_map = torch.sum(((label == 1) & (decisions == 1)), dim=-1).to(torch.float32) + res["speaker_error"] = torch.sum(torch.min(n_ref, n_sys) - n_map) + res["correct"] = torch.sum(label == decisions) / label.shape[1] + res["diarization_error"] = ( + res["speaker_miss"] + res["speaker_falarm"] + res["speaker_error"] + ) + res["frames"] = len(label) + return res diff --git a/almeval/models/stepaudio/funasr_detach/models/emotion2vec/__init__.py b/almeval/models/stepaudio/funasr_detach/models/emotion2vec/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/almeval/models/stepaudio/funasr_detach/models/emotion2vec/audio.py b/almeval/models/stepaudio/funasr_detach/models/emotion2vec/audio.py new file mode 100644 index 0000000000000000000000000000000000000000..0f68ab497360713f4ff01cbca4b3cee4e231ac2d --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/emotion2vec/audio.py @@ -0,0 +1,165 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import numpy as np +import torch.nn as nn +from functools import partial +import torch.nn.functional as F +from typing import Callable, Dict + +from funasr_detach.models.emotion2vec.fairseq_modules import ( + LayerNorm, + SamePad, + TransposeLast, + ConvFeatureExtractionModel, +) +from funasr_detach.models.emotion2vec.modules import Modality, BlockEncoder, Decoder1d +from funasr_detach.models.emotion2vec.base import ( + ModalitySpecificEncoder, + get_alibi_bias, +) + + +class AudioEncoder(ModalitySpecificEncoder): + + def __init__( + self, + modality_cfg, + embed_dim: int, + make_block: Callable[[float], nn.ModuleList], + norm_layer: Callable[[int], nn.LayerNorm], + layer_norm_first: bool, + alibi_biases: Dict, + ): + + self.feature_enc_layers = eval(modality_cfg.feature_encoder_spec) + feature_embed_dim = self.feature_enc_layers[-1][0] + + local_encoder = ConvFeatureExtractionModel( + conv_layers=self.feature_enc_layers, + dropout=0.0, + mode=modality_cfg.extractor_mode, + conv_bias=False, + ) + + project_features = nn.Sequential( + TransposeLast(), + nn.LayerNorm(feature_embed_dim), + nn.Linear(feature_embed_dim, embed_dim), + ) + + num_pos_layers = modality_cfg.conv_pos_depth + k = max(3, modality_cfg.conv_pos_width // num_pos_layers) + + positional_encoder = nn.Sequential( + TransposeLast(), + *[ + nn.Sequential( + nn.Conv1d( + embed_dim, + embed_dim, + kernel_size=k, + padding=k // 2, + groups=modality_cfg.conv_pos_groups, + ), + SamePad(k), + TransposeLast(), + LayerNorm(embed_dim, elementwise_affine=False), + TransposeLast(), + nn.GELU(), + ) + for _ in range(num_pos_layers) + ], + TransposeLast(), + ) + + if modality_cfg.conv_pos_pre_ln: + positional_encoder = nn.Sequential(LayerNorm(embed_dim), positional_encoder) + + dpr = np.linspace( + modality_cfg.start_drop_path_rate, + modality_cfg.end_drop_path_rate, + modality_cfg.prenet_depth, + ) + context_encoder = BlockEncoder( + nn.ModuleList(make_block(dpr[i]) for i in range(modality_cfg.prenet_depth)), + norm_layer(embed_dim) if not layer_norm_first else None, + layer_norm_first, + modality_cfg.prenet_layerdrop, + modality_cfg.prenet_dropout, + ) + + decoder = ( + Decoder1d(modality_cfg.decoder, embed_dim) + if modality_cfg.decoder is not None + else None + ) + + alibi_bias_fn = partial(get_alibi_bias, alibi_biases=alibi_biases) + + super().__init__( + modality_cfg=modality_cfg, + embed_dim=embed_dim, + local_encoder=local_encoder, + project_features=project_features, + fixed_positional_encoder=None, + relative_positional_encoder=positional_encoder, + context_encoder=context_encoder, + decoder=decoder, + get_alibi_bias=alibi_bias_fn, + ) + + def convert_padding_mask(self, x, padding_mask): + def get_feat_extract_output_lengths(input_lengths: torch.LongTensor): + """ + Computes the output length of the convolutional layers + """ + + def _conv_out_length(input_length, kernel_size, stride): + return torch.floor((input_length - kernel_size) / stride + 1) + + for i in range(len(self.feature_enc_layers)): + input_lengths = _conv_out_length( + input_lengths, + self.feature_enc_layers[i][1], + self.feature_enc_layers[i][2], + ) + + return input_lengths.to(torch.long) + + if padding_mask is not None: + input_lengths = (1 - padding_mask.long()).sum(-1) + # apply conv formula to get real output_lengths + output_lengths = get_feat_extract_output_lengths(input_lengths) + + if padding_mask.any(): + padding_mask = torch.zeros(x.shape[:2], dtype=x.dtype, device=x.device) + + # these two operations makes sure that all values + # before the output lengths indices are attended to + padding_mask[ + ( + torch.arange(padding_mask.shape[0], device=padding_mask.device), + output_lengths - 1, + ) + ] = 1 + padding_mask = ( + 1 - padding_mask.flip([-1]).cumsum(-1).flip([-1]) + ).bool() + else: + padding_mask = torch.zeros( + x.shape[:2], dtype=torch.bool, device=x.device + ) + + return padding_mask + + def reset_parameters(self): + super().reset_parameters() + for mod in self.project_features.children(): + if isinstance(mod, nn.Linear): + mod.reset_parameters() + if self.decoder is not None: + self.decoder.reset_parameters() diff --git a/almeval/models/stepaudio/funasr_detach/models/emotion2vec/base.py b/almeval/models/stepaudio/funasr_detach/models/emotion2vec/base.py new file mode 100644 index 0000000000000000000000000000000000000000..0c006a1f73301ceec7c063d57b609c2e476197b2 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/emotion2vec/base.py @@ -0,0 +1,646 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import logging +import math +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from collections import namedtuple +from dataclasses import dataclass +from functools import partial +from omegaconf import MISSING, II +from typing import Optional, Callable +from funasr_detach.models.emotion2vec.fairseq_modules import compute_mask_indices +from funasr_detach.models.emotion2vec.fairseq_modules import GradMultiply +from funasr_detach.models.emotion2vec.fairseq_modules import index_put + + +logger = logging.getLogger(__name__) + + +MaskSeed = namedtuple("MaskSeed", ["seed", "update", "ids"]) +MaskInfo = namedtuple("MaskInfo", ["x_unmasked", "mask", "ids_restore", "ids_keep"]) + + +class ModalitySpecificEncoder(nn.Module): + def __init__( + self, + modality_cfg, + embed_dim: int, + local_encoder: nn.Module, + project_features: nn.Module, + fixed_positional_encoder: Optional[nn.Module], + relative_positional_encoder: Optional[nn.Module], + context_encoder: nn.Module, + decoder: nn.Module, + get_alibi_bias: Optional[Callable[[int, int, str, str], torch.Tensor]], + ): + super().__init__() + + self.modality_cfg = modality_cfg + self.local_encoder = local_encoder + self.project_features = project_features + self.fixed_positional_encoder = fixed_positional_encoder + self.relative_positional_encoder = relative_positional_encoder + self.context_encoder = context_encoder + + self.decoder = decoder + self.get_alibi_bias = get_alibi_bias if modality_cfg.use_alibi_encoder else None + + self.local_grad_mult = self.modality_cfg.local_grad_mult + + self.extra_tokens = None + if modality_cfg.num_extra_tokens > 0: + self.extra_tokens = nn.Parameter( + torch.zeros(1, modality_cfg.num_extra_tokens, embed_dim) + ) + if not modality_cfg.init_extra_token_zero: + nn.init.normal_(self.extra_tokens) + elif self.extra_tokens.size(1) > 1: + nn.init.normal_(self.extra_tokens[:, 1:]) + + self.alibi_scale = None + if self.get_alibi_bias is not None: + self.alibi_scale = nn.Parameter( + torch.full( + ( + ( + (modality_cfg.prenet_depth + modality_cfg.model_depth) + if modality_cfg.learned_alibi_scale_per_layer + else 1 + ), + 1, + ( + self.modality_cfg.num_alibi_heads + if modality_cfg.learned_alibi_scale_per_head + else 1 + ), + 1, + 1, + ), + modality_cfg.alibi_scale, + dtype=torch.float, + ), + requires_grad=modality_cfg.learned_alibi_scale, + ) + + if modality_cfg.learned_alibi and self.get_alibi_bias is not None: + assert modality_cfg.alibi_max_pos is not None + alibi_bias = self.get_alibi_bias( + batch_size=1, + time_steps=modality_cfg.alibi_max_pos, + heads=modality_cfg.num_alibi_heads, + scale=1.0, + dtype=torch.float, + device="cpu", + ) + self.alibi_bias = nn.Parameter(alibi_bias) + self.get_alibi_bias = partial( + _learned_alibi_bias, alibi_bias=self.alibi_bias + ) + + def upgrade_state_dict_named(self, state_dict, name): + k = f"{name}.alibi_scale" + if k in state_dict and state_dict[k].dim() == 4: + state_dict[k] = state_dict[k].unsqueeze(0) + + return state_dict + + def convert_padding_mask(self, x, padding_mask): + return padding_mask + + def decoder_input(self, x, mask_info: MaskInfo): + inp_drop = self.modality_cfg.decoder.input_dropout + if inp_drop > 0: + x = F.dropout(x, inp_drop, training=self.training, inplace=True) + + num_extra = self.modality_cfg.num_extra_tokens + + if mask_info is not None: + num_masked = mask_info.ids_restore.shape[1] - x.shape[1] + num_extra + + mask_tokens = x.new_empty( + x.size(0), + num_masked, + x.size(-1), + ).normal_(0, self.modality_cfg.mask_noise_std) + + x_ = torch.cat([x[:, num_extra:], mask_tokens], dim=1) + x = torch.gather(x_, dim=1, index=mask_info.ids_restore) + + if self.modality_cfg.decoder.add_positions_masked: + assert self.fixed_positional_encoder is not None + pos = self.fixed_positional_encoder(x, None) + x = x + (pos * mask_info.mask.unsqueeze(-1)) + else: + x = x[:, num_extra:] + + if self.modality_cfg.decoder.add_positions_all: + assert self.fixed_positional_encoder is not None + x = x + self.fixed_positional_encoder(x, None) + + return x, mask_info + + def local_features(self, features): + if self.local_grad_mult > 0: + if self.local_grad_mult == 1.0: + x = self.local_encoder(features) + else: + x = GradMultiply.apply( + self.local_encoder(features), self.local_grad_mult + ) + else: + with torch.no_grad(): + x = self.local_encoder(features) + + x = self.project_features(x) + return x + + def contextualized_features( + self, + x, + padding_mask, + mask, + remove_masked, + clone_batch: int = 1, + mask_seeds: Optional[torch.Tensor] = None, + precomputed_mask=None, + ): + + if padding_mask is not None: + padding_mask = self.convert_padding_mask(x, padding_mask) + + local_features = x + if mask and clone_batch == 1: + local_features = local_features.clone() + + orig_B, orig_T, _ = x.shape + pre_mask_B = orig_B + mask_info = None + + x_pos = None + if self.fixed_positional_encoder is not None: + x = x + self.fixed_positional_encoder(x, padding_mask) + + if mask: + if clone_batch > 1: + x = x.repeat_interleave(clone_batch, 0) + if mask_seeds is not None: + clone_hash = [ + int(hash((mask_seeds.seed, ind)) % 1e10) + for ind in range(clone_batch - 1) + ] + clone_hash = torch.tensor([0] + clone_hash).long().view(1, -1) + + id = mask_seeds.ids + id = id.repeat_interleave(clone_batch, 0) + id = id.view(-1, clone_batch) + clone_hash.to(id) + id = id.view(-1) + mask_seeds = MaskSeed( + seed=mask_seeds.seed, update=mask_seeds.update, ids=id + ) + if padding_mask is not None: + padding_mask = padding_mask.repeat_interleave(clone_batch, 0) + + x, mask_info = self.compute_mask( + x, + padding_mask, + mask_seed=mask_seeds, + apply=self.relative_positional_encoder is not None or not remove_masked, + precomputed_mask=precomputed_mask, + ) + + if self.relative_positional_encoder is not None: + x_pos = self.relative_positional_encoder(x) + + masked_padding_mask = padding_mask + if mask and remove_masked: + x = mask_info.x_unmasked + if x_pos is not None: + x = x + gather_unmasked(x_pos, mask_info) + + if padding_mask is not None and padding_mask.any(): + masked_padding_mask = gather_unmasked_mask(padding_mask, mask_info) + if not masked_padding_mask.any(): + masked_padding_mask = None + else: + masked_padding_mask = None + + elif x_pos is not None: + x = x + x_pos + + alibi_bias = None + alibi_scale = self.alibi_scale + + if self.get_alibi_bias is not None: + alibi_bias = self.get_alibi_bias( + batch_size=pre_mask_B, + time_steps=orig_T, + heads=self.modality_cfg.num_alibi_heads, + dtype=torch.float32, + device=x.device, + ) + + if alibi_scale is not None: + alibi_scale = alibi_scale.clamp_min(0) + if alibi_scale.size(0) == 1: + alibi_bias = alibi_bias * alibi_scale.squeeze(0).type_as(alibi_bias) + alibi_scale = None + + if clone_batch > 1: + alibi_bias = alibi_bias.repeat_interleave(clone_batch, 0) + + if mask_info is not None and remove_masked: + alibi_bias = masked_alibi(alibi_bias, mask_info) + + if self.extra_tokens is not None: + num = self.extra_tokens.size(1) + x = torch.cat([self.extra_tokens.expand(x.size(0), -1, -1), x], dim=1) + if masked_padding_mask is not None: + # B x T + masked_padding_mask = F.pad(masked_padding_mask, (num, 0)) + if alibi_bias is not None: + # B x H x T x T + alibi_bias = F.pad(alibi_bias, (num, 0, num, 0)) + + x = self.context_encoder( + x, + masked_padding_mask, + alibi_bias, + ( + alibi_scale[: self.modality_cfg.prenet_depth] + if alibi_scale is not None + else None + ), + ) + + return { + "x": x, + "local_features": local_features, + "padding_mask": masked_padding_mask, + "alibi_bias": alibi_bias, + "alibi_scale": ( + alibi_scale[self.modality_cfg.prenet_depth :] + if alibi_scale is not None and alibi_scale.size(0) > 1 + else alibi_scale + ), + "encoder_mask": mask_info, + } + + def forward( + self, + features, + padding_mask, + mask: bool, + remove_masked: bool, + clone_batch: int = 1, + mask_seeds: Optional[torch.Tensor] = None, + precomputed_mask=None, + ): + x = self.local_features(features) + return self.contextualized_features( + x, + padding_mask, + mask, + remove_masked, + clone_batch, + mask_seeds, + precomputed_mask, + ) + + def reset_parameters(self): + pass + + def compute_mask( + self, + x, + padding_mask, + mask_seed: Optional[MaskSeed], + apply, + precomputed_mask, + ): + if precomputed_mask is not None: + mask = precomputed_mask + mask_info = self.make_maskinfo(x, mask) + else: + B, T, C = x.shape + cfg = self.modality_cfg + + mask_prob = cfg.mask_prob + + if ( + cfg.mask_prob_min is not None + and cfg.mask_prob_min >= 0 + and cfg.mask_prob_min < mask_prob + ): + mask_prob = np.random.uniform(cfg.mask_prob_min, mask_prob) + + if mask_prob > 0: + if cfg.mask_length == 1: + mask_info = random_masking(x, mask_prob, mask_seed) + else: + if self.modality_cfg.inverse_mask: + mask_prob = 1 - mask_prob + + mask = compute_mask_indices( + (B, T), + padding_mask, + mask_prob, + cfg.mask_length, + min_masks=1, + require_same_masks=True, + mask_dropout=cfg.mask_dropout, + add_masks=cfg.add_masks, + seed=mask_seed.seed if mask_seed is not None else None, + epoch=mask_seed.update if mask_seed is not None else None, + indices=mask_seed.ids if mask_seed is not None else None, + ) + + mask = torch.from_numpy(mask).to(device=x.device) + if self.modality_cfg.inverse_mask: + mask = 1 - mask + mask_info = self.make_maskinfo(x, mask) + else: + mask_info = None + + if apply: + x = self.apply_mask(x, mask_info) + + return x, mask_info + + def make_maskinfo(self, x, mask, shape=None): + if shape is None: + B, T, D = x.shape + else: + B, T, D = shape + + mask = mask.to(torch.uint8) + ids_shuffle = mask.argsort(dim=1) + ids_restore = ids_shuffle.argsort(dim=1).unsqueeze(-1).expand(-1, -1, D) + + len_keep = T - mask[0].sum() + if self.modality_cfg.keep_masked_pct > 0: + len_keep += round((T - int(len_keep)) * self.modality_cfg.keep_masked_pct) + + ids_keep = ids_shuffle[:, :len_keep] + + if shape is not None: + x_unmasked = None + else: + ids_keep = ids_keep.unsqueeze(-1).expand(-1, -1, D) + x_unmasked = torch.gather(x, dim=1, index=ids_keep) + + mask_info = MaskInfo( + x_unmasked=x_unmasked, + mask=mask, + ids_restore=ids_restore, + ids_keep=ids_keep, + ) + return mask_info + + def apply_mask(self, x, mask_info): + cfg = self.modality_cfg + B, T, C = x.shape + + if mask_info is not None: + mask = mask_info.mask + if cfg.encoder_zero_mask: + x = x * (1 - mask.type_as(x).unsqueeze(-1)) + else: + num_masks = mask.sum().item() + masks = x.new_empty(num_masks, x.size(-1)).normal_( + 0, cfg.mask_noise_std + ) + x = index_put(x, mask, masks) + if cfg.mask_channel_prob > 0: + mask_channel = compute_mask_indices( + (B, C), + None, + cfg.mask_channel_prob, + cfg.mask_channel_length, + ) + mask_channel = ( + torch.from_numpy(mask_channel) + .to(x.device) + .unsqueeze(1) + .expand(-1, T, -1) + ) + x = index_put(x, mask_channel, 0) + return x + + def remove_pretraining_modules(self, keep_decoder=False): + if not keep_decoder: + self.decoder = None + + +def get_annealed_rate(start, end, curr_step, total_steps): + if curr_step >= total_steps: + return end + r = end - start + pct_remaining = 1 - curr_step / total_steps + return end - r * pct_remaining + + +# adapted from MAE +def random_masking(x, mask_ratio, mask_seed: Optional[MaskSeed]): + N, L, D = x.shape # batch, length, dim + len_keep = int(L * (1 - mask_ratio)) + + generator = None + if mask_seed is not None: + seed = int( + hash((mask_seed.seed, mask_seed.update, mask_seed.ids.sum().item())) % 1e6 + ) + generator = torch.Generator(device=x.device) + generator.manual_seed(seed) + + noise = torch.rand(N, L, generator=generator, device=x.device) # noise in [0, 1] + + # sort noise for each sample + ids_shuffle = noise.argsort(dim=1) # ascend: small is keep, large is remove + ids_restore = ids_shuffle.argsort(dim=1) + + # keep the first subset + ids_keep = ids_shuffle[:, :len_keep] + ids_keep = ids_keep.unsqueeze(-1).expand(-1, -1, D) + x_unmasked = torch.gather(x, dim=1, index=ids_keep) + + # generate the binary mask: 0 is keep, 1 is remove + mask = torch.ones([N, L], dtype=x.dtype, device=x.device) + mask[:, :len_keep] = 0 + # unshuffle to get the binary mask + mask = torch.gather(mask, dim=1, index=ids_restore) + + ids_restore = ids_restore.unsqueeze(-1).expand(-1, -1, D) + + return MaskInfo( + x_unmasked=x_unmasked, mask=mask, ids_restore=ids_restore, ids_keep=ids_keep + ) + + +def gather_unmasked(x: torch.Tensor, mask_info: MaskInfo) -> torch.Tensor: + return torch.gather( + x, + dim=1, + index=mask_info.ids_keep, + ) + + +def gather_unmasked_mask(x: torch.Tensor, mask_info: MaskInfo) -> torch.Tensor: + return torch.gather( + x, + dim=1, + index=mask_info.ids_keep[..., 0], # ignore the feature dimension + ) + + +def get_alibi( + max_positions: int, + attention_heads: int, + dims: int = 1, + distance: str = "manhattan", +): + def get_slopes(n): + def get_slopes_power_of_2(n): + start = 2 ** (-(2 ** -(math.log2(n) - 3))) + ratio = start + return [start * ratio**i for i in range(n)] + + # In the paper, we only train models that have 2^a heads for some + # a. This function has some good properties that only occur when + # the input is a power of 2. To maintain that even when the number + # of heads is not a power of 2, we use this workaround. + if math.log2(n).is_integer(): + return get_slopes_power_of_2(n) + else: + closest_power_of_2 = 2 ** math.floor(math.log2(n)) + return ( + get_slopes_power_of_2(closest_power_of_2) + + get_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2] + ) + + maxpos = max_positions + attn_heads = attention_heads + slopes = torch.Tensor(get_slopes(attn_heads)) + + if dims == 1: + # prepare alibi position linear bias. Note that wav2vec2 is non + # autoregressive model so we want a symmetric mask with 0 on the + # diagonal and other wise linear decreasing valuees + pos_bias = ( + torch.abs( + torch.arange(maxpos).unsqueeze(0) - torch.arange(maxpos).unsqueeze(1) + ) + * -1 + ) + elif dims == 2: + if distance == "manhattan": + df = lambda x1, y1, x2, y2: abs(x1 - x2) + abs(y1 - y2) + elif distance == "euclidean": + df = lambda x1, y1, x2, y2: math.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2) + + n = math.sqrt(max_positions) + assert n.is_integer(), n + n = int(n) + + pos_bias = torch.zeros((max_positions, max_positions)) + + for i in range(n): + for j in range(n): + for k in range(n): + for l in range(n): + new_x = i * n + j + new_y = k * n + l + pos_bias[new_x, new_y] = -df(i, j, k, l) + + else: + raise Exception(f"unsupported number of alibi dims: {dims}") + + alibi_bias = slopes.unsqueeze(1).unsqueeze(1) * pos_bias.unsqueeze(0).expand( + attn_heads, -1, -1 + ) + + return alibi_bias + + +def get_alibi_bias( + alibi_biases, + batch_size, + time_steps, + heads, + dtype, + device, + dims=1, + distance="manhattan", +): + cache_key = f"{dims}_{heads}_{distance}" + + buffered = alibi_biases.get(cache_key, None) + + target_size = heads * batch_size + if ( + buffered is None + or buffered.size(0) < target_size + or buffered.size(1) < time_steps + or buffered.dtype != dtype + or buffered.device != device + ): + bt = max(time_steps, buffered.size(1) if buffered is not None else 0) + bn = max(target_size, buffered.size(0) if buffered is not None else 0) // heads + + buffered = ( + get_alibi(bt, heads, dims=dims, distance=distance) + .to(dtype=dtype, device=device) + .repeat(bn, 1, 1) + ) + + alibi_biases[cache_key] = buffered + + b = buffered[:target_size, :time_steps, :time_steps] + b = b.view(batch_size, heads, time_steps, time_steps) + return b + + +def _learned_alibi_bias( + alibi_bias, + batch_size, + time_steps, + heads, + scale, + dtype, + device, +): + assert alibi_bias.size(1) == heads, alibi_bias.shape + assert alibi_bias.dtype == dtype, alibi_bias.dtype + assert alibi_bias.device == device, alibi_bias.device + + if alibi_bias.size(-1) < time_steps: + psz = math.ceil((time_steps - alibi_bias.size(-1)) / 2) + alibi_bias = F.pad(alibi_bias, (psz, psz, psz, psz), mode="replicate") + + alibi_bias = alibi_bias.expand(batch_size, -1, -1, -1) * scale + return alibi_bias[..., :time_steps, :time_steps] + + +def masked_alibi(alibi_bias, mask_info): + H = alibi_bias.size(1) + + orig_bias = alibi_bias + + index = mask_info.ids_keep.unsqueeze(1)[..., 0].unsqueeze(-1) + alibi_bias = torch.gather( + orig_bias, + dim=-2, + index=index.expand(-1, H, -1, mask_info.ids_restore.size(1)), + ) + alibi_bias = torch.gather( + alibi_bias, + dim=-1, + index=index.transpose(-1, -2).expand(-1, H, alibi_bias.size(-2), -1), + ) + + return alibi_bias diff --git a/almeval/models/stepaudio/funasr_detach/models/emotion2vec/fairseq_modules.py b/almeval/models/stepaudio/funasr_detach/models/emotion2vec/fairseq_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..46dd225dc82898ee163f5f50228978591970b670 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/emotion2vec/fairseq_modules.py @@ -0,0 +1,310 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import Optional, Tuple, List +import numpy as np + + +def LayerNorm(normalized_shape, eps=1e-5, elementwise_affine=True, export=False): + return torch.nn.LayerNorm(normalized_shape, eps, elementwise_affine) + + +class SamePad(nn.Module): + def __init__(self, kernel_size, causal=False): + super().__init__() + if causal: + self.remove = kernel_size - 1 + else: + self.remove = 1 if kernel_size % 2 == 0 else 0 + + def forward(self, x): + if self.remove > 0: + x = x[:, :, : -self.remove] + return x + + +class TransposeLast(nn.Module): + def __init__(self, deconstruct_idx=None): + super().__init__() + self.deconstruct_idx = deconstruct_idx + + def forward(self, x): + if self.deconstruct_idx is not None: + x = x[self.deconstruct_idx] + return x.transpose(-2, -1) + + +class Fp32LayerNorm(nn.LayerNorm): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self, input): + output = F.layer_norm( + input.float(), + self.normalized_shape, + self.weight.float() if self.weight is not None else None, + self.bias.float() if self.bias is not None else None, + self.eps, + ) + return output.type_as(input) + + +class Fp32GroupNorm(nn.GroupNorm): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self, input): + output = F.group_norm( + input.float(), + self.num_groups, + self.weight.float() if self.weight is not None else None, + self.bias.float() if self.bias is not None else None, + self.eps, + ) + return output.type_as(input) + + +class ConvFeatureExtractionModel(nn.Module): + def __init__( + self, + conv_layers: List[Tuple[int, int, int]], + dropout: float = 0.0, + mode: str = "default", + conv_bias: bool = False, + ): + super().__init__() + + assert mode in {"default", "layer_norm"} + + def block( + n_in, + n_out, + k, + stride, + is_layer_norm=False, + is_group_norm=False, + conv_bias=False, + ): + def make_conv(): + conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias) + nn.init.kaiming_normal_(conv.weight) + return conv + + assert ( + is_layer_norm and is_group_norm + ) == False, "layer norm and group norm are exclusive" + + if is_layer_norm: + return nn.Sequential( + make_conv(), + nn.Dropout(p=dropout), + nn.Sequential( + TransposeLast(), + Fp32LayerNorm(dim, elementwise_affine=True), + TransposeLast(), + ), + nn.GELU(), + ) + elif is_group_norm: + return nn.Sequential( + make_conv(), + nn.Dropout(p=dropout), + Fp32GroupNorm(dim, dim, affine=True), + nn.GELU(), + ) + else: + return nn.Sequential(make_conv(), nn.Dropout(p=dropout), nn.GELU()) + + in_d = 1 + self.conv_layers = nn.ModuleList() + for i, cl in enumerate(conv_layers): + assert len(cl) == 3, "invalid conv definition: " + str(cl) + (dim, k, stride) = cl + + self.conv_layers.append( + block( + in_d, + dim, + k, + stride, + is_layer_norm=mode == "layer_norm", + is_group_norm=mode == "default" and i == 0, + conv_bias=conv_bias, + ) + ) + in_d = dim + + def forward(self, x): + + # BxT -> BxCxT + x = x.unsqueeze(1) + + for conv in self.conv_layers: + x = conv(x) + + return x + + +def compute_mask_indices( + shape: Tuple[int, int], + padding_mask: Optional[torch.Tensor], + mask_prob: float, + mask_length: int, + mask_type: str = "static", + mask_other: float = 0.0, + min_masks: int = 0, + no_overlap: bool = False, + min_space: int = 0, + require_same_masks: bool = True, + mask_dropout: float = 0.0, +) -> np.ndarray: + """ + Computes random mask spans for a given shape + + Args: + shape: the the shape for which to compute masks. + should be of size 2 where first element is batch size and 2nd is timesteps + padding_mask: optional padding mask of the same size as shape, which will prevent masking padded elements + mask_prob: probability for each token to be chosen as start of the span to be masked. this will be multiplied by + number of timesteps divided by length of mask span to mask approximately this percentage of all elements. + however due to overlaps, the actual number will be smaller (unless no_overlap is True) + mask_type: how to compute mask lengths + static = fixed size + uniform = sample from uniform distribution [mask_other, mask_length*2] + normal = sample from normal distribution with mean mask_length and stdev mask_other. mask is min 1 element + poisson = sample from possion distribution with lambda = mask length + min_masks: minimum number of masked spans + no_overlap: if false, will switch to an alternative recursive algorithm that prevents spans from overlapping + min_space: only used if no_overlap is True, this is how many elements to keep unmasked between spans + require_same_masks: if true, will randomly drop out masks until same amount of masks remains in each sample + mask_dropout: randomly dropout this percentage of masks in each example + """ + + bsz, all_sz = shape + mask = np.full((bsz, all_sz), False) + + all_num_mask = int( + # add a random number for probabilistic rounding + mask_prob * all_sz / float(mask_length) + + np.random.rand() + ) + + all_num_mask = max(min_masks, all_num_mask) + + mask_idcs = [] + for i in range(bsz): + if padding_mask is not None: + sz = all_sz - padding_mask[i].long().sum().item() + num_mask = int( + # add a random number for probabilistic rounding + mask_prob * sz / float(mask_length) + + np.random.rand() + ) + num_mask = max(min_masks, num_mask) + else: + sz = all_sz + num_mask = all_num_mask + + if mask_type == "static": + lengths = np.full(num_mask, mask_length) + elif mask_type == "uniform": + lengths = np.random.randint(mask_other, mask_length * 2 + 1, size=num_mask) + elif mask_type == "normal": + lengths = np.random.normal(mask_length, mask_other, size=num_mask) + lengths = [max(1, int(round(x))) for x in lengths] + elif mask_type == "poisson": + lengths = np.random.poisson(mask_length, size=num_mask) + lengths = [int(round(x)) for x in lengths] + else: + raise Exception("unknown mask selection " + mask_type) + + if sum(lengths) == 0: + lengths[0] = min(mask_length, sz - 1) + + if no_overlap: + mask_idc = [] + + def arrange(s, e, length, keep_length): + span_start = np.random.randint(s, e - length) + mask_idc.extend(span_start + i for i in range(length)) + + new_parts = [] + if span_start - s - min_space >= keep_length: + new_parts.append((s, span_start - min_space + 1)) + if e - span_start - length - min_space > keep_length: + new_parts.append((span_start + length + min_space, e)) + return new_parts + + parts = [(0, sz)] + min_length = min(lengths) + for length in sorted(lengths, reverse=True): + lens = np.fromiter( + (e - s if e - s >= length + min_space else 0 for s, e in parts), + np.int, + ) + l_sum = np.sum(lens) + if l_sum == 0: + break + probs = lens / np.sum(lens) + c = np.random.choice(len(parts), p=probs) + s, e = parts.pop(c) + parts.extend(arrange(s, e, length, min_length)) + mask_idc = np.asarray(mask_idc) + else: + min_len = min(lengths) + if sz - min_len <= num_mask: + min_len = sz - num_mask - 1 + + mask_idc = np.random.choice(sz - min_len, num_mask, replace=False) + + mask_idc = np.asarray( + [ + mask_idc[j] + offset + for j in range(len(mask_idc)) + for offset in range(lengths[j]) + ] + ) + + mask_idcs.append(np.unique(mask_idc[mask_idc < sz])) + + min_len = min([len(m) for m in mask_idcs]) + for i, mask_idc in enumerate(mask_idcs): + if len(mask_idc) > min_len and require_same_masks: + mask_idc = np.random.choice(mask_idc, min_len, replace=False) + if mask_dropout > 0: + num_holes = np.rint(len(mask_idc) * mask_dropout).astype(int) + mask_idc = np.random.choice( + mask_idc, len(mask_idc) - num_holes, replace=False + ) + + mask[i, mask_idc] = True + + return mask + + +class GradMultiply(torch.autograd.Function): + @staticmethod + def forward(ctx, x, scale): + ctx.scale = scale + res = x.new(x) + return res + + @staticmethod + def backward(ctx, grad): + return grad * ctx.scale, None + + +def is_xla_tensor(tensor): + return torch.is_tensor(tensor) and tensor.device.type == "xla" + + +def index_put(tensor, indices, value): + if is_xla_tensor(tensor): + for _ in range(indices.dim(), tensor.dim()): + indices = indices.unsqueeze(-1) + if indices.size(-1) < tensor.size(-1): + indices = indices.expand_as(tensor) + tensor = torch.mul(tensor, ~indices) + torch.mul(value, indices) + else: + tensor[indices] = value + return tensor diff --git a/almeval/models/stepaudio/funasr_detach/models/emotion2vec/model.py b/almeval/models/stepaudio/funasr_detach/models/emotion2vec/model.py new file mode 100644 index 0000000000000000000000000000000000000000..f0d971bcd01a0663375c472258762882f14ff4d8 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/emotion2vec/model.py @@ -0,0 +1,272 @@ +#!/usr/bin/env python3 +# -*- encoding: utf-8 -*- +# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved. +# MIT License (https://opensource.org/licenses/MIT) +# Modified from https://github.com/ddlBoJack/emotion2vec/tree/main + +import os +import time +import torch +import logging +import numpy as np +from functools import partial +from omegaconf import OmegaConf +import torch.nn.functional as F +from contextlib import contextmanager +from distutils.version import LooseVersion + +from funasr_detach.register import tables +from funasr_detach.models.emotion2vec.modules import AltBlock +from funasr_detach.models.emotion2vec.audio import AudioEncoder +from funasr_detach.utils.load_utils import load_audio_text_image_video + + +logger = logging.getLogger(__name__) +if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): + from torch.cuda.amp import autocast +else: + # Nothing to do if torch<1.6.0 + @contextmanager + def autocast(enabled=True): + yield + + +@tables.register("model_classes", "Emotion2vec") +class Emotion2vec(torch.nn.Module): + """ + Author: Ziyang Ma, Zhisheng Zheng, Jiaxin Ye, Jinchao Li, Zhifu Gao, Shiliang Zhang, Xie Chen + emotion2vec: Self-Supervised Pre-Training for Speech Emotion Representation + https://arxiv.org/abs/2312.15185 + """ + + def __init__(self, **kwargs): + super().__init__() + # import pdb; pdb.set_trace() + cfg = OmegaConf.create(kwargs["model_conf"]) + self.cfg = cfg + + make_layer_norm = partial( + torch.nn.LayerNorm, + eps=cfg.get("norm_eps"), + elementwise_affine=cfg.get("norm_affine"), + ) + + def make_block(drop_path, dim=None, heads=None): + return AltBlock( + cfg.get("embed_dim") if dim is None else dim, + cfg.get("num_heads") if heads is None else heads, + cfg.get("mlp_ratio"), + qkv_bias=True, + drop=cfg.get("encoder_dropout"), + attn_drop=cfg.get("attention_dropout"), + mlp_drop=cfg.get("activation_dropout"), + post_mlp_drop=cfg.get("post_mlp_drop"), + drop_path=drop_path, + norm_layer=make_layer_norm, + layer_norm_first=cfg.get("layer_norm_first"), + ffn_targets=not cfg.get("end_of_block_targets"), + ) + + self.alibi_biases = {} + self.modality_encoders = torch.nn.ModuleDict() + + enc = AudioEncoder( + cfg.modalities.audio, + cfg.get("embed_dim"), + make_block, + make_layer_norm, + cfg.get("layer_norm_first"), + self.alibi_biases, + ) + self.modality_encoders["AUDIO"] = enc + + self.ema = None + + self.average_top_k_layers = cfg.get("average_top_k_layers") + self.loss_beta = cfg.get("loss_beta") + self.loss_scale = cfg.get("loss_scale") + + self.dropout_input = torch.nn.Dropout(cfg.get("dropout_input")) + + dpr = np.linspace( + cfg.get("start_drop_path_rate"), + cfg.get("end_drop_path_rate"), + cfg.get("depth"), + ) + + self.blocks = torch.nn.ModuleList( + [make_block(dpr[i]) for i in range(cfg.get("depth"))] + ) + + self.norm = None + if cfg.get("layer_norm_first"): + self.norm = make_layer_norm(cfg.get("embed_dim")) + + vocab_size = kwargs.get("vocab_size", -1) + self.proj = None + if vocab_size > 0: + self.proj = torch.nn.Linear(cfg.get("embed_dim"), vocab_size) + + def forward( + self, + source, + target=None, + id=None, + mode=None, + padding_mask=None, + mask=True, + features_only=False, + force_remove_masked=False, + remove_extra_tokens=True, + precomputed_mask=None, + **kwargs, + ): + + feature_extractor = self.modality_encoders["AUDIO"] + + mask_seeds = None + + extractor_out = feature_extractor( + source, + padding_mask, + mask, + remove_masked=not features_only or force_remove_masked, + clone_batch=self.cfg.get("clone_batch") if not features_only else 1, + mask_seeds=mask_seeds, + precomputed_mask=precomputed_mask, + ) + + x = extractor_out["x"] + encoder_mask = extractor_out["encoder_mask"] + masked_padding_mask = extractor_out["padding_mask"] + masked_alibi_bias = extractor_out.get("alibi_bias", None) + alibi_scale = extractor_out.get("alibi_scale", None) + + if self.dropout_input is not None: + x = self.dropout_input(x) + + layer_results = [] + for i, blk in enumerate(self.blocks): + if ( + not self.training + or self.cfg.get("layerdrop", 0) == 0 + or (np.random.random() > self.cfg.get("layerdrop", 0)) + ): + ab = masked_alibi_bias + if ab is not None and alibi_scale is not None: + scale = ( + alibi_scale[i] + if alibi_scale.size(0) > 1 + else alibi_scale.squeeze(0) + ) + ab = ab * scale.type_as(ab) + + x, lr = blk( + x, + padding_mask=masked_padding_mask, + alibi_bias=ab, + ) + if features_only: + layer_results.append(lr) + + if self.norm is not None: + x = self.norm(x) + + if features_only: + if remove_extra_tokens: + x = x[:, feature_extractor.modality_cfg.num_extra_tokens :] + if masked_padding_mask is not None: + masked_padding_mask = masked_padding_mask[ + :, feature_extractor.modality_cfg.num_extra_tokens : + ] + + return { + "x": x, + "padding_mask": masked_padding_mask, + "layer_results": layer_results, + "mask": encoder_mask, + } + + def extract_features( + self, source, mode=None, padding_mask=None, mask=False, remove_extra_tokens=True + ): + res = self.forward( + source, + mode=mode, + padding_mask=padding_mask, + mask=mask, + features_only=True, + remove_extra_tokens=remove_extra_tokens, + ) + return res + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + + # if source_file.endswith('.wav'): + # wav, sr = sf.read(source_file) + # channel = sf.info(source_file).channels + # assert sr == 16e3, "Sample rate should be 16kHz, but got {}in file {}".format(sr, source_file) + # assert channel == 1, "Channel should be 1, but got {} in file {}".format(channel, source_file) + granularity = kwargs.get("granularity", "utterance") + extract_embedding = kwargs.get("extract_embedding", True) + if self.proj is None: + extract_embedding = True + meta_data = {} + # extract fbank feats + time1 = time.perf_counter() + audio_sample_list = load_audio_text_image_video( + data_in, + fs=16000, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + ) + time2 = time.perf_counter() + meta_data["load_data"] = f"{time2 - time1:0.3f}" + meta_data["batch_data_time"] = len(audio_sample_list[0]) / kwargs.get( + "fs", 16000 + ) + + results = [] + output_dir = kwargs.get("output_dir") + if output_dir: + os.makedirs(output_dir, exist_ok=True) + for i, wav in enumerate(audio_sample_list): + source = wav.to(device=kwargs["device"]) + if self.cfg.normalize: + source = F.layer_norm(source, source.shape) + source = source.view(1, -1) + + feats = self.extract_features(source, padding_mask=None) + x = feats["x"] + feats = feats["x"].squeeze(0).cpu().numpy() + if granularity == "frame": + feats = feats + elif granularity == "utterance": + feats = np.mean(feats, axis=0) + + if output_dir and extract_embedding: + np.save(os.path.join(output_dir, "{}.npy".format(key[i])), feats) + + labels = tokenizer.token_list if tokenizer is not None else [] + scores = [] + if self.proj: + x = x.mean(dim=1) + x = self.proj(x) + x = torch.softmax(x, dim=-1) + scores = x[0].tolist() + + result_i = {"key": key[i], "labels": labels, "scores": scores} + if extract_embedding: + result_i["feats"] = feats + results.append(result_i) + + return results, meta_data diff --git a/almeval/models/stepaudio/funasr_detach/models/emotion2vec/modules.py b/almeval/models/stepaudio/funasr_detach/models/emotion2vec/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..9d2bb6454e5bd1f2e3178a3d66c18a2eac23e592 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/emotion2vec/modules.py @@ -0,0 +1,323 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import numpy as np +import torch.nn as nn +from enum import Enum, auto +import torch.nn.functional as F +from dataclasses import dataclass +from funasr_detach.models.emotion2vec.fairseq_modules import ( + LayerNorm, + SamePad, + TransposeLast, +) + + +class Modality(Enum): + AUDIO = auto() + + +@dataclass +class D2vDecoderConfig: + decoder_dim: int = 384 + decoder_groups: int = 16 + decoder_kernel: int = 5 + decoder_layers: int = 5 + input_dropout: float = 0.1 + + add_positions_masked: bool = False + add_positions_all: bool = False + + decoder_residual: bool = True + projection_layers: int = 1 + projection_ratio: float = 2.0 + + +class FixedPositionalEncoder(nn.Module): + def __init__(self, pos_embed): + super().__init__() + self.positions = pos_embed + + def forward(self, x, padding_mask): + return self.positions + + +class TextFeatPositionalEncoder(nn.Module): + """ + Original encoder expects (B, T) long input. This module wraps it to take + local_encoder output which are (B, T, D) float tensors + """ + + def __init__(self, pos_encoder): + super().__init__() + self.pos_encoder = pos_encoder + + def forward(self, x, padding_mask): + # assume padded token embeddings are 0s + # TODO: consider using padding_mask as input + return self.pos_encoder(x[..., 0]) + + +class BlockEncoder(nn.Module): + def __init__(self, blocks, norm_layer, layer_norm_first, layerdrop, dropout): + super().__init__() + self.blocks = blocks + self.norm = norm_layer + self.layer_norm_first = layer_norm_first + self.layerdrop = layerdrop + self.dropout = nn.Dropout(dropout, inplace=True) + + def forward(self, x, padding_mask, alibi_bias, alibi_scale): + if self.norm is not None and not self.layer_norm_first: + x = self.norm(x) + + x = self.dropout(x) + + for i, blk in enumerate(self.blocks): + if ( + not self.training + or self.layerdrop == 0 + or (np.random.random() > self.layerdrop) + ): + ab = alibi_bias + if ab is not None and alibi_scale is not None: + scale = ( + alibi_scale[i] + if alibi_scale.size(0) > 1 + else alibi_scale.squeeze(0) + ) + ab = ab * scale.type_as(ab) + x, _ = blk(x, padding_mask, ab) + + if self.norm is not None and self.layer_norm_first: + x = self.norm(x) + + return x + + +class DecoderBase(nn.Module): + decoder_cfg: D2vDecoderConfig + + def __init__(self, cfg: D2vDecoderConfig): + super().__init__() + + self.decoder_cfg = cfg + + def reset_parameters(self): + for mod in self.proj.modules(): + if isinstance(mod, nn.Linear): + mod.reset_parameters() + + def add_residual(self, x, residual, i, mask_info): + if ( + residual is None + or not self.decoder_cfg.decoder_residual + or residual.size(1) != x.size(1) + ): + return x + + ret = x + residual + + return ret + + +class Decoder1d(DecoderBase): + def __init__(self, cfg: D2vDecoderConfig, input_dim): + super().__init__(cfg) + + def make_block(in_dim): + block = [ + nn.Conv1d( + in_dim, + cfg.decoder_dim, + kernel_size=cfg.decoder_kernel, + padding=cfg.decoder_kernel // 2, + groups=cfg.decoder_groups, + ), + SamePad(cfg.decoder_kernel), + TransposeLast(), + LayerNorm(cfg.decoder_dim, elementwise_affine=False), + TransposeLast(), + nn.GELU(), + ] + + return nn.Sequential(*block) + + self.blocks = nn.Sequential( + *[ + make_block(input_dim if i == 0 else cfg.decoder_dim) + for i in range(cfg.decoder_layers) + ] + ) + + projs = [] + curr_dim = cfg.decoder_dim + for i in range(cfg.projection_layers - 1): + next_dim = int(curr_dim * cfg.projection_ratio) if i == 0 else curr_dim + projs.append(nn.Linear(curr_dim, next_dim)) + projs.append(nn.GELU()) + curr_dim = next_dim + projs.append(nn.Linear(curr_dim, input_dim)) + if len(projs) == 1: + self.proj = projs[0] + else: + self.proj = nn.Sequential(*projs) + + def forward(self, x, mask_info): + + x = x.transpose(1, 2) + + residual = x + + for i, layer in enumerate(self.blocks): + x = layer(x) + x = self.add_residual(x, residual, i, mask_info) + residual = x + + x = x.transpose(1, 2) + x = self.proj(x) + return x + + +class AltBlock(nn.Module): + def __init__( + self, + dim, + num_heads, + mlp_ratio=4.0, + qkv_bias=False, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + mlp_drop=0.0, + post_mlp_drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + layer_norm_first=True, + ffn_targets=False, + cosine_attention=False, + ): + super().__init__() + + self.layer_norm_first = layer_norm_first + self.ffn_targets = ffn_targets + + from funasr_detach.models.emotion2vec.timm_modules import DropPath, Mlp + + self.norm1 = norm_layer(dim) + self.attn = AltAttention( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=drop, + cosine_attention=cosine_attention, + ) + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=mlp_drop, + ) + self.post_mlp_dropout = nn.Dropout(post_mlp_drop, inplace=False) + + def forward(self, x, padding_mask=None, alibi_bias=None): + if self.layer_norm_first: + x = x + self.drop_path(self.attn(self.norm1(x), padding_mask, alibi_bias)) + r = x = self.mlp(self.norm2(x)) + t = x + x = r + self.drop_path(self.post_mlp_dropout(x)) + if not self.ffn_targets: + t = x + else: + x = x + self.drop_path(self.attn(x, padding_mask, alibi_bias)) + r = x = self.norm1(x) + x = self.mlp(x) + t = x + x = self.norm2(r + self.drop_path(self.post_mlp_dropout(x))) + if not self.ffn_targets: + t = x + + return x, t + + +class AltAttention(nn.Module): + def __init__( + self, + dim, + num_heads=8, + qkv_bias=False, + qk_scale=None, + attn_drop=0.0, + proj_drop=0.0, + cosine_attention=False, + ): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim**-0.5 + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + self.cosine_attention = cosine_attention + + if cosine_attention: + self.logit_scale = nn.Parameter( + torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True + ) + + def forward(self, x, padding_mask=None, alibi_bias=None): + B, N, C = x.shape + qkv = ( + self.qkv(x) + .reshape(B, N, 3, self.num_heads, C // self.num_heads) + .permute(2, 0, 3, 1, 4) # qkv x B x H x L x D + ) + q, k, v = ( + qkv[0], + qkv[1], + qkv[2], + ) # make torchscript happy (cannot use tensor as tuple) + + dtype = q.dtype + + if self.cosine_attention: + # cosine attention + attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1) + logit_scale = torch.clamp( + self.logit_scale, max=torch.log(torch.tensor(1.0 / 0.01)) + ).exp() + attn = attn * logit_scale + else: + q = q * self.scale + attn = q @ k.transpose(-2, -1) + + if alibi_bias is not None: + attn = attn.type_as(alibi_bias) + attn[:, : alibi_bias.size(1)] += alibi_bias + + if padding_mask is not None and padding_mask.any(): + attn = attn.masked_fill( + padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), + float("-inf"), + ) + + attn = attn.softmax(dim=-1, dtype=torch.float32).to(dtype=dtype) + attn = self.attn_drop(attn) + x = (attn @ v).transpose(1, 2) # + x = x.reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x diff --git a/almeval/models/stepaudio/funasr_detach/models/emotion2vec/template.yaml b/almeval/models/stepaudio/funasr_detach/models/emotion2vec/template.yaml new file mode 100644 index 0000000000000000000000000000000000000000..53bca63742d2954229afd58ff406e5d3d7ae97a6 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/emotion2vec/template.yaml @@ -0,0 +1,113 @@ +# This is an example that demonstrates how to configure a model file. +# You can modify the configuration according to your own requirements. + +# to print the register_table: +# from funasr.register import tables +# tables.print() + +# network architecture +model: Emotion2vec +model_conf: + loss_beta: 0.0 + loss_scale: null + depth: 8 + start_drop_path_rate: 0.0 + end_drop_path_rate: 0.0 + num_heads: 12 + norm_eps: 1e-05 + norm_affine: true + encoder_dropout: 0.1 + post_mlp_drop: 0.1 + attention_dropout: 0.1 + activation_dropout: 0.0 + dropout_input: 0.0 + layerdrop: 0.05 + embed_dim: 768 + mlp_ratio: 4.0 + layer_norm_first: false + average_top_k_layers: 8 + end_of_block_targets: false + clone_batch: 8 + layer_norm_target_layer: false + batch_norm_target_layer: false + instance_norm_target_layer: true + instance_norm_targets: false + layer_norm_targets: false + ema_decay: 0.999 + ema_same_dtype: true + log_norms: true + ema_end_decay: 0.99999 + ema_anneal_end_step: 20000 + ema_encoder_only: false + max_update: 100000 + extractor_mode: layer_norm + shared_decoder: null + min_target_var: 0.1 + min_pred_var: 0.01 + supported_modality: AUDIO + mae_init: false + seed: 1 + skip_ema: false + cls_loss: 1.0 + recon_loss: 0.0 + d2v_loss: 1.0 + decoder_group: false + adversarial_training: false + adversarial_hidden_dim: 128 + adversarial_weight: 0.1 + cls_type: chunk + normalize: true + + modalities: + audio: + type: AUDIO + prenet_depth: 4 + prenet_layerdrop: 0.05 + prenet_dropout: 0.1 + start_drop_path_rate: 0.0 + end_drop_path_rate: 0.0 + num_extra_tokens: 10 + init_extra_token_zero: true + mask_noise_std: 0.01 + mask_prob_min: null + mask_prob: 0.5 + inverse_mask: false + mask_prob_adjust: 0.05 + keep_masked_pct: 0.0 + mask_length: 5 + add_masks: false + remove_masks: false + mask_dropout: 0.0 + encoder_zero_mask: true + mask_channel_prob: 0.0 + mask_channel_length: 64 + ema_local_encoder: false + local_grad_mult: 1.0 + use_alibi_encoder: true + alibi_scale: 1.0 + learned_alibi: false + alibi_max_pos: null + learned_alibi_scale: true + learned_alibi_scale_per_head: true + learned_alibi_scale_per_layer: false + num_alibi_heads: 12 + model_depth: 8 + decoder: + decoder_dim: 384 + decoder_groups: 16 + decoder_kernel: 7 + decoder_layers: 4 + input_dropout: 0.1 + add_positions_masked: false + add_positions_all: false + decoder_residual: true + projection_layers: 1 + projection_ratio: 2.0 + extractor_mode: layer_norm + feature_encoder_spec: '[(512, 10, 5)] + [(512, 3, 2)] * 4 + [(512,2,2)] + [(512,2,2)]' + conv_pos_width: 95 + conv_pos_groups: 16 + conv_pos_depth: 5 + conv_pos_pre_ln: false + + diff --git a/almeval/models/stepaudio/funasr_detach/models/emotion2vec/timm_modules.py b/almeval/models/stepaudio/funasr_detach/models/emotion2vec/timm_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..60d7076a7a5beeab86f8f0f98e7d92d5eea4225d --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/emotion2vec/timm_modules.py @@ -0,0 +1,100 @@ +import torch.nn as nn +import collections.abc +from itertools import repeat +from functools import partial + + +def drop_path( + x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True +): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + + This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, + the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... + See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for + changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use + 'survival rate' as the argument. + + """ + if drop_prob == 0.0 or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * ( + x.ndim - 1 + ) # work with diff dim tensors, not just 2D ConvNets + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + if keep_prob > 0.0 and scale_by_keep: + random_tensor.div_(keep_prob) + return x * random_tensor + + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + self.scale_by_keep = scale_by_keep + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training, self.scale_by_keep) + + def extra_repr(self): + return f"drop_prob={round(self.drop_prob,3):0.3f}" + + +# From PyTorch internals +def _ntuple(n): + def parse(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + return tuple(x) + return tuple(repeat(x, n)) + + return parse + + +to_1tuple = _ntuple(1) +to_2tuple = _ntuple(2) +to_3tuple = _ntuple(3) +to_4tuple = _ntuple(4) +to_ntuple = _ntuple + + +class Mlp(nn.Module): + """MLP as used in Vision Transformer, MLP-Mixer and related networks""" + + def __init__( + self, + in_features, + hidden_features=None, + out_features=None, + act_layer=nn.GELU, + norm_layer=None, + bias=True, + drop=0.0, + use_conv=False, + ): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + bias = to_2tuple(bias) + drop_probs = to_2tuple(drop) + linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear + + self.fc1 = linear_layer(in_features, hidden_features, bias=bias[0]) + self.act = act_layer() + self.drop1 = nn.Dropout(drop_probs[0]) + self.norm = ( + norm_layer(hidden_features) if norm_layer is not None else nn.Identity() + ) + self.fc2 = linear_layer(hidden_features, out_features, bias=bias[1]) + self.drop2 = nn.Dropout(drop_probs[1]) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop1(x) + x = self.norm(x) + x = self.fc2(x) + x = self.drop2(x) + return x diff --git a/almeval/models/stepaudio/funasr_detach/models/eres2net/__init__.py b/almeval/models/stepaudio/funasr_detach/models/eres2net/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..570bbb1f155ab288ceec70c73a377d2e46e31fac --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/eres2net/__init__.py @@ -0,0 +1,2 @@ +from .eres2net import ERes2Net +from .eres2net_aug import ERes2NetAug diff --git a/almeval/models/stepaudio/funasr_detach/models/eres2net/eres2net.py b/almeval/models/stepaudio/funasr_detach/models/eres2net/eres2net.py new file mode 100644 index 0000000000000000000000000000000000000000..afe981510c3e644caf5cf7e0dded941c1ec679fd --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/eres2net/eres2net.py @@ -0,0 +1,431 @@ +# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Res2Net implementation is adapted from https://github.com/wenet-e2e/wespeaker. +ERes2Net incorporates both local and global feature fusion techniques to improve the performance. +The local feature fusion (LFF) fuses the features within one single residual block to extract the local signal. +The global feature fusion (GFF) takes acoustic features of different scales as input to aggregate global signal. +ERes2Net-Large is an upgraded version of ERes2Net that uses a larger number of parameters to achieve better +recognition performance. Parameters expansion, baseWidth, and scale can be modified to obtain optimal performance. +""" + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import funasr_detach.models.sond.pooling.pooling_layers as pooling_layers + +from funasr_detach.models.eres2net.fusion import AFF + + +class ReLU(nn.Hardtanh): + + def __init__(self, inplace=False): + super(ReLU, self).__init__(0, 20, inplace) + + def __repr__(self): + inplace_str = "inplace" if self.inplace else "" + return self.__class__.__name__ + " (" + inplace_str + ")" + + +def conv1x1(in_planes, out_planes, stride=1): + "1x1 convolution without padding" + return nn.Conv2d( + in_planes, out_planes, kernel_size=1, stride=stride, padding=0, bias=False + ) + + +def conv3x3(in_planes, out_planes, stride=1): + "3x3 convolution with padding" + return nn.Conv2d( + in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False + ) + + +class BasicBlockERes2Net(nn.Module): + expansion = 2 + + def __init__(self, in_planes, planes, stride=1, baseWidth=32, scale=2): + super(BasicBlockERes2Net, self).__init__() + width = int(math.floor(planes * (baseWidth / 64.0))) + self.conv1 = conv1x1(in_planes, width * scale, stride) + self.bn1 = nn.BatchNorm2d(width * scale) + self.nums = scale + + convs = [] + bns = [] + for i in range(self.nums): + convs.append(conv3x3(width, width)) + bns.append(nn.BatchNorm2d(width)) + self.convs = nn.ModuleList(convs) + self.bns = nn.ModuleList(bns) + self.relu = ReLU(inplace=True) + + self.conv3 = conv1x1(width * scale, planes * self.expansion) + self.bn3 = nn.BatchNorm2d(planes * self.expansion) + self.shortcut = nn.Sequential() + if stride != 1 or in_planes != self.expansion * planes: + self.shortcut = nn.Sequential( + nn.Conv2d( + in_planes, + self.expansion * planes, + kernel_size=1, + stride=stride, + bias=False, + ), + nn.BatchNorm2d(self.expansion * planes), + ) + self.stride = stride + self.width = width + self.scale = scale + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + spx = torch.split(out, self.width, 1) + for i in range(self.nums): + if i == 0: + sp = spx[i] + else: + sp = sp + spx[i] + sp = self.convs[i](sp) + sp = self.relu(self.bns[i](sp)) + if i == 0: + out = sp + else: + out = torch.cat((out, sp), 1) + + out = self.conv3(out) + out = self.bn3(out) + + residual = self.shortcut(x) + out += residual + out = self.relu(out) + + return out + + +class BasicBlockERes2Net_diff_AFF(nn.Module): + expansion = 2 + + def __init__(self, in_planes, planes, stride=1, baseWidth=32, scale=2): + super(BasicBlockERes2Net_diff_AFF, self).__init__() + width = int(math.floor(planes * (baseWidth / 64.0))) + self.conv1 = conv1x1(in_planes, width * scale, stride) + self.bn1 = nn.BatchNorm2d(width * scale) + self.nums = scale + + convs = [] + fuse_models = [] + bns = [] + for i in range(self.nums): + convs.append(conv3x3(width, width)) + bns.append(nn.BatchNorm2d(width)) + for j in range(self.nums - 1): + fuse_models.append(AFF(channels=width)) + + self.convs = nn.ModuleList(convs) + self.bns = nn.ModuleList(bns) + self.fuse_models = nn.ModuleList(fuse_models) + self.relu = ReLU(inplace=True) + + self.conv3 = conv1x1(width * scale, planes * self.expansion) + self.bn3 = nn.BatchNorm2d(planes * self.expansion) + self.shortcut = nn.Sequential() + if stride != 1 or in_planes != self.expansion * planes: + self.shortcut = nn.Sequential( + nn.Conv2d( + in_planes, + self.expansion * planes, + kernel_size=1, + stride=stride, + bias=False, + ), + nn.BatchNorm2d(self.expansion * planes), + ) + self.stride = stride + self.width = width + self.scale = scale + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + spx = torch.split(out, self.width, 1) + for i in range(self.nums): + if i == 0: + sp = spx[i] + else: + sp = self.fuse_models[i - 1](sp, spx[i]) + + sp = self.convs[i](sp) + sp = self.relu(self.bns[i](sp)) + if i == 0: + out = sp + else: + out = torch.cat((out, sp), 1) + + out = self.conv3(out) + out = self.bn3(out) + + residual = self.shortcut(x) + out += residual + out = self.relu(out) + + return out + + +class ERes2Net(nn.Module): + def __init__( + self, + block=BasicBlockERes2Net, + block_fuse=BasicBlockERes2Net_diff_AFF, + num_blocks=[3, 4, 6, 3], + m_channels=32, + feat_dim=80, + embedding_size=192, + pooling_func="TSTP", + two_emb_layer=False, + ): + super(ERes2Net, self).__init__() + self.in_planes = m_channels + self.feat_dim = feat_dim + self.embedding_size = embedding_size + self.stats_dim = int(feat_dim / 8) * m_channels * 8 + self.two_emb_layer = two_emb_layer + + self.conv1 = nn.Conv2d( + 1, m_channels, kernel_size=3, stride=1, padding=1, bias=False + ) + self.bn1 = nn.BatchNorm2d(m_channels) + self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=1) + self.layer2 = self._make_layer(block, m_channels * 2, num_blocks[1], stride=2) + self.layer3 = self._make_layer( + block_fuse, m_channels * 4, num_blocks[2], stride=2 + ) + self.layer4 = self._make_layer( + block_fuse, m_channels * 8, num_blocks[3], stride=2 + ) + + # Downsampling module for each layer + self.layer1_downsample = nn.Conv2d( + m_channels * 2, + m_channels * 4, + kernel_size=3, + stride=2, + padding=1, + bias=False, + ) + self.layer2_downsample = nn.Conv2d( + m_channels * 4, + m_channels * 8, + kernel_size=3, + padding=1, + stride=2, + bias=False, + ) + self.layer3_downsample = nn.Conv2d( + m_channels * 8, + m_channels * 16, + kernel_size=3, + padding=1, + stride=2, + bias=False, + ) + + # Bottom-up fusion module + self.fuse_mode12 = AFF(channels=m_channels * 4) + self.fuse_mode123 = AFF(channels=m_channels * 8) + self.fuse_mode1234 = AFF(channels=m_channels * 16) + + self.n_stats = 1 if pooling_func == "TAP" or pooling_func == "TSDP" else 2 + self.pool = getattr(pooling_layers, pooling_func)( + in_dim=self.stats_dim * block.expansion + ) + self.seg_1 = nn.Linear( + self.stats_dim * block.expansion * self.n_stats, embedding_size + ) + if self.two_emb_layer: + self.seg_bn_1 = nn.BatchNorm1d(embedding_size, affine=False) + self.seg_2 = nn.Linear(embedding_size, embedding_size) + else: + self.seg_bn_1 = nn.Identity() + self.seg_2 = nn.Identity() + + def _make_layer(self, block, planes, num_blocks, stride): + strides = [stride] + [1] * (num_blocks - 1) + layers = [] + for stride in strides: + layers.append(block(self.in_planes, planes, stride)) + self.in_planes = planes * block.expansion + return nn.Sequential(*layers) + + def forward(self, x): + x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T) + x = x.unsqueeze_(1) + out = F.relu(self.bn1(self.conv1(x))) + out1 = self.layer1(out) + out2 = self.layer2(out1) + out1_downsample = self.layer1_downsample(out1) + fuse_out12 = self.fuse_mode12(out2, out1_downsample) + out3 = self.layer3(out2) + fuse_out12_downsample = self.layer2_downsample(fuse_out12) + fuse_out123 = self.fuse_mode123(out3, fuse_out12_downsample) + out4 = self.layer4(out3) + fuse_out123_downsample = self.layer3_downsample(fuse_out123) + fuse_out1234 = self.fuse_mode1234(out4, fuse_out123_downsample) + stats = self.pool(fuse_out1234) + + embed_a = self.seg_1(stats) + if self.two_emb_layer: + out = F.relu(embed_a) + out = self.seg_bn_1(out) + embed_b = self.seg_2(out) + return embed_b + else: + return embed_a + + +class BasicBlockRes2Net(nn.Module): + expansion = 2 + + def __init__(self, in_planes, planes, stride=1, baseWidth=32, scale=2): + super(BasicBlockRes2Net, self).__init__() + width = int(math.floor(planes * (baseWidth / 64.0))) + self.conv1 = conv1x1(in_planes, width * scale, stride) + self.bn1 = nn.BatchNorm2d(width * scale) + self.nums = scale - 1 + convs = [] + bns = [] + for i in range(self.nums): + convs.append(conv3x3(width, width)) + bns.append(nn.BatchNorm2d(width)) + self.convs = nn.ModuleList(convs) + self.bns = nn.ModuleList(bns) + self.relu = ReLU(inplace=True) + + self.conv3 = conv1x1(width * scale, planes * self.expansion) + self.bn3 = nn.BatchNorm2d(planes * self.expansion) + self.shortcut = nn.Sequential() + if stride != 1 or in_planes != self.expansion * planes: + self.shortcut = nn.Sequential( + nn.Conv2d( + in_planes, + self.expansion * planes, + kernel_size=1, + stride=stride, + bias=False, + ), + nn.BatchNorm2d(self.expansion * planes), + ) + self.stride = stride + self.width = width + self.scale = scale + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + spx = torch.split(out, self.width, 1) + for i in range(self.nums): + if i == 0: + sp = spx[i] + else: + sp = sp + spx[i] + sp = self.convs[i](sp) + sp = self.relu(self.bns[i](sp)) + if i == 0: + out = sp + else: + out = torch.cat((out, sp), 1) + + out = torch.cat((out, spx[self.nums]), 1) + + out = self.conv3(out) + out = self.bn3(out) + + residual = self.shortcut(x) + out += residual + out = self.relu(out) + + return out + + +class Res2Net(nn.Module): + def __init__( + self, + block=BasicBlockRes2Net, + num_blocks=[3, 4, 6, 3], + m_channels=32, + feat_dim=80, + embedding_size=192, + pooling_func="TSTP", + two_emb_layer=False, + ): + super(Res2Net, self).__init__() + self.in_planes = m_channels + self.feat_dim = feat_dim + self.embedding_size = embedding_size + self.stats_dim = int(feat_dim / 8) * m_channels * 8 + self.two_emb_layer = two_emb_layer + + self.conv1 = nn.Conv2d( + 1, m_channels, kernel_size=3, stride=1, padding=1, bias=False + ) + self.bn1 = nn.BatchNorm2d(m_channels) + self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=1) + self.layer2 = self._make_layer(block, m_channels * 2, num_blocks[1], stride=2) + self.layer3 = self._make_layer(block, m_channels * 4, num_blocks[2], stride=2) + self.layer4 = self._make_layer(block, m_channels * 8, num_blocks[3], stride=2) + + self.n_stats = 1 if pooling_func == "TAP" or pooling_func == "TSDP" else 2 + self.pool = getattr(pooling_layers, pooling_func)( + in_dim=self.stats_dim * block.expansion + ) + self.seg_1 = nn.Linear( + self.stats_dim * block.expansion * self.n_stats, embedding_size + ) + if self.two_emb_layer: + self.seg_bn_1 = nn.BatchNorm1d(embedding_size, affine=False) + self.seg_2 = nn.Linear(embedding_size, embedding_size) + else: + self.seg_bn_1 = nn.Identity() + self.seg_2 = nn.Identity() + + def _make_layer(self, block, planes, num_blocks, stride): + strides = [stride] + [1] * (num_blocks - 1) + layers = [] + for stride in strides: + layers.append(block(self.in_planes, planes, stride)) + self.in_planes = planes * block.expansion + return nn.Sequential(*layers) + + def forward(self, x): + x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T) + + x = x.unsqueeze_(1) + out = F.relu(self.bn1(self.conv1(x))) + out = self.layer1(out) + out = self.layer2(out) + out = self.layer3(out) + out = self.layer4(out) + + stats = self.pool(out) + + embed_a = self.seg_1(stats) + if self.two_emb_layer: + out = F.relu(embed_a) + out = self.seg_bn_1(out) + embed_b = self.seg_2(out) + return embed_b + else: + return embed_a diff --git a/almeval/models/stepaudio/funasr_detach/models/eres2net/eres2net_aug.py b/almeval/models/stepaudio/funasr_detach/models/eres2net/eres2net_aug.py new file mode 100644 index 0000000000000000000000000000000000000000..1ad19e7e2ff1c29ce78f81abc9b1ef2531683f9e --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/eres2net/eres2net_aug.py @@ -0,0 +1,292 @@ +# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Res2Net implementation is adapted from https://github.com/wenet-e2e/wespeaker. +ERes2Net incorporates both local and global feature fusion techniques to improve the performance. +The local feature fusion (LFF) fuses the features within one single residual block to extract the local signal. +The global feature fusion (GFF) takes acoustic features of different scales as input to aggregate global signal. +ERes2Net-Large is an upgraded version of ERes2Net that uses a larger number of parameters to achieve better +recognition performance. Parameters expansion, baseWidth, and scale can be modified to obtain optimal performance. +""" + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import funasr_detach.models.sond.pooling.pooling_layers as pooling_layers + +from funasr_detach.models.eres2net.fusion import AFF + + +class ReLU(nn.Hardtanh): + + def __init__(self, inplace=False): + super(ReLU, self).__init__(0, 20, inplace) + + def __repr__(self): + inplace_str = "inplace" if self.inplace else "" + return self.__class__.__name__ + " (" + inplace_str + ")" + + +def conv1x1(in_planes, out_planes, stride=1): + "1x1 convolution without padding" + return nn.Conv2d( + in_planes, out_planes, kernel_size=1, stride=stride, padding=0, bias=False + ) + + +def conv3x3(in_planes, out_planes, stride=1): + "3x3 convolution with padding" + return nn.Conv2d( + in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False + ) + + +class BasicBlockERes2Net(nn.Module): + expansion = 4 + + def __init__(self, in_planes, planes, stride=1, baseWidth=24, scale=3): + super(BasicBlockERes2Net, self).__init__() + width = int(math.floor(planes * (baseWidth / 64.0))) + self.conv1 = conv1x1(in_planes, width * scale, stride) + self.bn1 = nn.BatchNorm2d(width * scale) + self.nums = scale + + convs = [] + bns = [] + for i in range(self.nums): + convs.append(conv3x3(width, width)) + bns.append(nn.BatchNorm2d(width)) + self.convs = nn.ModuleList(convs) + self.bns = nn.ModuleList(bns) + self.relu = ReLU(inplace=True) + + self.conv3 = conv1x1(width * scale, planes * self.expansion) + self.bn3 = nn.BatchNorm2d(planes * self.expansion) + self.shortcut = nn.Sequential() + if stride != 1 or in_planes != self.expansion * planes: + self.shortcut = nn.Sequential( + nn.Conv2d( + in_planes, + self.expansion * planes, + kernel_size=1, + stride=stride, + bias=False, + ), + nn.BatchNorm2d(self.expansion * planes), + ) + self.stride = stride + self.width = width + self.scale = scale + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + spx = torch.split(out, self.width, 1) + for i in range(self.nums): + if i == 0: + sp = spx[i] + else: + sp = sp + spx[i] + sp = self.convs[i](sp) + sp = self.relu(self.bns[i](sp)) + if i == 0: + out = sp + else: + out = torch.cat((out, sp), 1) + + out = self.conv3(out) + out = self.bn3(out) + + residual = self.shortcut(x) + out += residual + out = self.relu(out) + + return out + + +class BasicBlockERes2Net_diff_AFF(nn.Module): + expansion = 4 + + def __init__(self, in_planes, planes, stride=1, baseWidth=24, scale=3): + super(BasicBlockERes2Net_diff_AFF, self).__init__() + width = int(math.floor(planes * (baseWidth / 64.0))) + self.conv1 = conv1x1(in_planes, width * scale, stride) + self.bn1 = nn.BatchNorm2d(width * scale) + + self.nums = scale + + convs = [] + fuse_models = [] + bns = [] + for i in range(self.nums): + convs.append(conv3x3(width, width)) + bns.append(nn.BatchNorm2d(width)) + for j in range(self.nums - 1): + fuse_models.append(AFF(channels=width)) + + self.convs = nn.ModuleList(convs) + self.bns = nn.ModuleList(bns) + self.fuse_models = nn.ModuleList(fuse_models) + self.relu = ReLU(inplace=True) + + self.conv3 = conv1x1(width * scale, planes * self.expansion) + self.bn3 = nn.BatchNorm2d(planes * self.expansion) + self.shortcut = nn.Sequential() + if stride != 1 or in_planes != self.expansion * planes: + self.shortcut = nn.Sequential( + nn.Conv2d( + in_planes, + self.expansion * planes, + kernel_size=1, + stride=stride, + bias=False, + ), + nn.BatchNorm2d(self.expansion * planes), + ) + self.stride = stride + self.width = width + self.scale = scale + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + spx = torch.split(out, self.width, 1) + for i in range(self.nums): + if i == 0: + sp = spx[i] + else: + sp = self.fuse_models[i - 1](sp, spx[i]) + + sp = self.convs[i](sp) + sp = self.relu(self.bns[i](sp)) + if i == 0: + out = sp + else: + out = torch.cat((out, sp), 1) + + out = self.conv3(out) + out = self.bn3(out) + + residual = self.shortcut(x) + out += residual + out = self.relu(out) + + return out + + +class ERes2NetAug(nn.Module): + def __init__( + self, + block=BasicBlockERes2Net, + block_fuse=BasicBlockERes2Net_diff_AFF, + num_blocks=[3, 4, 6, 3], + m_channels=64, + feat_dim=80, + embedding_size=192, + pooling_func="TSTP", + two_emb_layer=False, + ): + super(ERes2NetAug, self).__init__() + self.in_planes = m_channels + self.feat_dim = feat_dim + self.embedding_size = embedding_size + self.stats_dim = int(feat_dim / 8) * m_channels * 8 + self.two_emb_layer = two_emb_layer + + self.conv1 = nn.Conv2d( + 1, m_channels, kernel_size=3, stride=1, padding=1, bias=False + ) + self.bn1 = nn.BatchNorm2d(m_channels) + self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=1) + self.layer2 = self._make_layer(block, m_channels * 2, num_blocks[1], stride=2) + self.layer3 = self._make_layer( + block_fuse, m_channels * 4, num_blocks[2], stride=2 + ) + self.layer4 = self._make_layer( + block_fuse, m_channels * 8, num_blocks[3], stride=2 + ) + + self.layer1_downsample = nn.Conv2d( + m_channels * 4, + m_channels * 8, + kernel_size=3, + padding=1, + stride=2, + bias=False, + ) + self.layer2_downsample = nn.Conv2d( + m_channels * 8, + m_channels * 16, + kernel_size=3, + padding=1, + stride=2, + bias=False, + ) + self.layer3_downsample = nn.Conv2d( + m_channels * 16, + m_channels * 32, + kernel_size=3, + padding=1, + stride=2, + bias=False, + ) + self.fuse_mode12 = AFF(channels=m_channels * 8) + self.fuse_mode123 = AFF(channels=m_channels * 16) + self.fuse_mode1234 = AFF(channels=m_channels * 32) + + self.n_stats = 1 if pooling_func == "TAP" or pooling_func == "TSDP" else 2 + self.pool = getattr(pooling_layers, pooling_func)( + in_dim=self.stats_dim * block.expansion + ) + self.seg_1 = nn.Linear( + self.stats_dim * block.expansion * self.n_stats, embedding_size + ) + if self.two_emb_layer: + self.seg_bn_1 = nn.BatchNorm1d(embedding_size, affine=False) + self.seg_2 = nn.Linear(embedding_size, embedding_size) + else: + self.seg_bn_1 = nn.Identity() + self.seg_2 = nn.Identity() + + def _make_layer(self, block, planes, num_blocks, stride): + strides = [stride] + [1] * (num_blocks - 1) + layers = [] + for stride in strides: + layers.append(block(self.in_planes, planes, stride)) + self.in_planes = planes * block.expansion + return nn.Sequential(*layers) + + def forward(self, x): + x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T) + + x = x.unsqueeze_(1) + out = F.relu(self.bn1(self.conv1(x))) + out1 = self.layer1(out) + out2 = self.layer2(out1) + out1_downsample = self.layer1_downsample(out1) + fuse_out12 = self.fuse_mode12(out2, out1_downsample) + out3 = self.layer3(out2) + fuse_out12_downsample = self.layer2_downsample(fuse_out12) + fuse_out123 = self.fuse_mode123(out3, fuse_out12_downsample) + out4 = self.layer4(out3) + fuse_out123_downsample = self.layer3_downsample(fuse_out123) + fuse_out1234 = self.fuse_mode1234(out4, fuse_out123_downsample) + stats = self.pool(fuse_out1234) + + embed_a = self.seg_1(stats) + if self.two_emb_layer: + out = F.relu(embed_a) + out = self.seg_bn_1(out) + embed_b = self.seg_2(out) + return embed_b + else: + return embed_a diff --git a/almeval/models/stepaudio/funasr_detach/models/eres2net/fusion.py b/almeval/models/stepaudio/funasr_detach/models/eres2net/fusion.py new file mode 100644 index 0000000000000000000000000000000000000000..fbe699efbaa33380594df2c5068b78e0ecbe7303 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/eres2net/fusion.py @@ -0,0 +1,28 @@ +# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +import torch +import torch.nn as nn + + +class AFF(nn.Module): + + def __init__(self, channels=64, r=4): + super(AFF, self).__init__() + inter_channels = int(channels // r) + + self.local_att = nn.Sequential( + nn.Conv2d(channels * 2, inter_channels, kernel_size=1, stride=1, padding=0), + nn.BatchNorm2d(inter_channels), + nn.SiLU(inplace=True), + nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0), + nn.BatchNorm2d(channels), + ) + + def forward(self, x, ds_y): + xa = torch.cat((x, ds_y), dim=1) + x_att = self.local_att(xa) + x_att = 1.0 + torch.tanh(x_att) + xo = torch.mul(x, x_att) + torch.mul(ds_y, 2.0 - x_att) + + return xo diff --git a/almeval/models/stepaudio/funasr_detach/models/fsmn_vad_streaming/__init__.py b/almeval/models/stepaudio/funasr_detach/models/fsmn_vad_streaming/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/almeval/models/stepaudio/funasr_detach/models/fsmn_vad_streaming/encoder.py b/almeval/models/stepaudio/funasr_detach/models/fsmn_vad_streaming/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..60169c08430ba45526caeef7f3b868b2bc5ef4c0 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/fsmn_vad_streaming/encoder.py @@ -0,0 +1,338 @@ +from typing import Tuple, Dict +import copy + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from funasr_detach.register import tables + + +class LinearTransform(nn.Module): + + def __init__(self, input_dim, output_dim): + super(LinearTransform, self).__init__() + self.input_dim = input_dim + self.output_dim = output_dim + self.linear = nn.Linear(input_dim, output_dim, bias=False) + + def forward(self, input): + output = self.linear(input) + + return output + + +class AffineTransform(nn.Module): + + def __init__(self, input_dim, output_dim): + super(AffineTransform, self).__init__() + self.input_dim = input_dim + self.output_dim = output_dim + self.linear = nn.Linear(input_dim, output_dim) + + def forward(self, input): + output = self.linear(input) + + return output + + +class RectifiedLinear(nn.Module): + + def __init__(self, input_dim, output_dim): + super(RectifiedLinear, self).__init__() + self.dim = input_dim + self.relu = nn.ReLU() + self.dropout = nn.Dropout(0.1) + + def forward(self, input): + out = self.relu(input) + return out + + +class FSMNBlock(nn.Module): + + def __init__( + self, + input_dim: int, + output_dim: int, + lorder=None, + rorder=None, + lstride=1, + rstride=1, + ): + super(FSMNBlock, self).__init__() + + self.dim = input_dim + + if lorder is None: + return + + self.lorder = lorder + self.rorder = rorder + self.lstride = lstride + self.rstride = rstride + + self.conv_left = nn.Conv2d( + self.dim, + self.dim, + [lorder, 1], + dilation=[lstride, 1], + groups=self.dim, + bias=False, + ) + + if self.rorder > 0: + self.conv_right = nn.Conv2d( + self.dim, + self.dim, + [rorder, 1], + dilation=[rstride, 1], + groups=self.dim, + bias=False, + ) + else: + self.conv_right = None + + def forward(self, input: torch.Tensor, cache: torch.Tensor): + x = torch.unsqueeze(input, 1) + x_per = x.permute(0, 3, 2, 1) # B D T C + + cache = cache.to(x_per.device) + y_left = torch.cat((cache, x_per), dim=2) + cache = y_left[:, :, -(self.lorder - 1) * self.lstride :, :] + y_left = self.conv_left(y_left) + out = x_per + y_left + + if self.conv_right is not None: + # maybe need to check + y_right = F.pad(x_per, [0, 0, 0, self.rorder * self.rstride]) + y_right = y_right[:, :, self.rstride :, :] + y_right = self.conv_right(y_right) + out += y_right + + out_per = out.permute(0, 3, 2, 1) + output = out_per.squeeze(1) + + return output, cache + + +class BasicBlock(nn.Module): + def __init__( + self, + linear_dim: int, + proj_dim: int, + lorder: int, + rorder: int, + lstride: int, + rstride: int, + stack_layer: int, + ): + super(BasicBlock, self).__init__() + self.lorder = lorder + self.rorder = rorder + self.lstride = lstride + self.rstride = rstride + self.stack_layer = stack_layer + self.linear = LinearTransform(linear_dim, proj_dim) + self.fsmn_block = FSMNBlock( + proj_dim, proj_dim, lorder, rorder, lstride, rstride + ) + self.affine = AffineTransform(proj_dim, linear_dim) + self.relu = RectifiedLinear(linear_dim, linear_dim) + + def forward(self, input: torch.Tensor, cache: Dict[str, torch.Tensor]): + x1 = self.linear(input) # B T D + cache_layer_name = "cache_layer_{}".format(self.stack_layer) + if cache_layer_name not in cache: + cache[cache_layer_name] = torch.zeros( + x1.shape[0], x1.shape[-1], (self.lorder - 1) * self.lstride, 1 + ) + x2, cache[cache_layer_name] = self.fsmn_block(x1, cache[cache_layer_name]) + x3 = self.affine(x2) + x4 = self.relu(x3) + return x4 + + +class FsmnStack(nn.Sequential): + def __init__(self, *args): + super(FsmnStack, self).__init__(*args) + + def forward(self, input: torch.Tensor, cache: Dict[str, torch.Tensor]): + x = input + for module in self._modules.values(): + x = module(x, cache) + return x + + +""" +FSMN net for keyword spotting +input_dim: input dimension +linear_dim: fsmn input dimensionll +proj_dim: fsmn projection dimension +lorder: fsmn left order +rorder: fsmn right order +num_syn: output dimension +fsmn_layers: no. of sequential fsmn layers +""" + + +@tables.register("encoder_classes", "FSMN") +class FSMN(nn.Module): + def __init__( + self, + input_dim: int, + input_affine_dim: int, + fsmn_layers: int, + linear_dim: int, + proj_dim: int, + lorder: int, + rorder: int, + lstride: int, + rstride: int, + output_affine_dim: int, + output_dim: int, + ): + super(FSMN, self).__init__() + + self.input_dim = input_dim + self.input_affine_dim = input_affine_dim + self.fsmn_layers = fsmn_layers + self.linear_dim = linear_dim + self.proj_dim = proj_dim + self.output_affine_dim = output_affine_dim + self.output_dim = output_dim + + self.in_linear1 = AffineTransform(input_dim, input_affine_dim) + self.in_linear2 = AffineTransform(input_affine_dim, linear_dim) + self.relu = RectifiedLinear(linear_dim, linear_dim) + self.fsmn = FsmnStack( + *[ + BasicBlock(linear_dim, proj_dim, lorder, rorder, lstride, rstride, i) + for i in range(fsmn_layers) + ] + ) + self.out_linear1 = AffineTransform(linear_dim, output_affine_dim) + self.out_linear2 = AffineTransform(output_affine_dim, output_dim) + self.softmax = nn.Softmax(dim=-1) + + def fuse_modules(self): + pass + + def forward( + self, input: torch.Tensor, cache: Dict[str, torch.Tensor] + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: + """ + Args: + input (torch.Tensor): Input tensor (B, T, D) + cache: when cache is not None, the forward is in streaming. The type of cache is a dict, egs, + {'cache_layer_1': torch.Tensor(B, T1, D)}, T1 is equal to self.lorder. It is {} for the 1st frame + """ + + x1 = self.in_linear1(input) + x2 = self.in_linear2(x1) + x3 = self.relu(x2) + x4 = self.fsmn(x3, cache) # self.cache will update automatically in self.fsmn + x5 = self.out_linear1(x4) + x6 = self.out_linear2(x5) + x7 = self.softmax(x6) + + return x7 + + +""" +one deep fsmn layer +dimproj: projection dimension, input and output dimension of memory blocks +dimlinear: dimension of mapping layer +lorder: left order +rorder: right order +lstride: left stride +rstride: right stride +""" + + +@tables.register("encoder_classes", "DFSMN") +class DFSMN(nn.Module): + + def __init__( + self, dimproj=64, dimlinear=128, lorder=20, rorder=1, lstride=1, rstride=1 + ): + super(DFSMN, self).__init__() + + self.lorder = lorder + self.rorder = rorder + self.lstride = lstride + self.rstride = rstride + + self.expand = AffineTransform(dimproj, dimlinear) + self.shrink = LinearTransform(dimlinear, dimproj) + + self.conv_left = nn.Conv2d( + dimproj, + dimproj, + [lorder, 1], + dilation=[lstride, 1], + groups=dimproj, + bias=False, + ) + + if rorder > 0: + self.conv_right = nn.Conv2d( + dimproj, + dimproj, + [rorder, 1], + dilation=[rstride, 1], + groups=dimproj, + bias=False, + ) + else: + self.conv_right = None + + def forward(self, input): + f1 = F.relu(self.expand(input)) + p1 = self.shrink(f1) + + x = torch.unsqueeze(p1, 1) + x_per = x.permute(0, 3, 2, 1) + + y_left = F.pad(x_per, [0, 0, (self.lorder - 1) * self.lstride, 0]) + + if self.conv_right is not None: + y_right = F.pad(x_per, [0, 0, 0, (self.rorder) * self.rstride]) + y_right = y_right[:, :, self.rstride :, :] + out = x_per + self.conv_left(y_left) + self.conv_right(y_right) + else: + out = x_per + self.conv_left(y_left) + + out1 = out.permute(0, 3, 2, 1) + output = input + out1.squeeze(1) + + return output + + +""" +build stacked dfsmn layers +""" + + +def buildDFSMNRepeats(linear_dim=128, proj_dim=64, lorder=20, rorder=1, fsmn_layers=6): + repeats = [ + nn.Sequential(DFSMN(proj_dim, linear_dim, lorder, rorder, 1, 1)) + for i in range(fsmn_layers) + ] + + return nn.Sequential(*repeats) + + +if __name__ == "__main__": + fsmn = FSMN(400, 140, 4, 250, 128, 10, 2, 1, 1, 140, 2599) + print(fsmn) + + num_params = sum(p.numel() for p in fsmn.parameters()) + print("the number of model params: {}".format(num_params)) + x = torch.zeros(128, 200, 400) # batch-size * time * dim + y, _ = fsmn(x) # batch-size * time * dim + print("input shape: {}".format(x.shape)) + print("output shape: {}".format(y.shape)) + + print(fsmn.to_kaldi_net()) diff --git a/almeval/models/stepaudio/funasr_detach/models/fsmn_vad_streaming/model.py b/almeval/models/stepaudio/funasr_detach/models/fsmn_vad_streaming/model.py new file mode 100644 index 0000000000000000000000000000000000000000..77ca870e2720205d97d596f86e4a44bdf42a904c --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/fsmn_vad_streaming/model.py @@ -0,0 +1,1048 @@ +#!/usr/bin/env python3 +# -*- encoding: utf-8 -*- +# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved. +# MIT License (https://opensource.org/licenses/MIT) + +import os +import json +import time +import math +import torch +from torch import nn +from enum import Enum +from dataclasses import dataclass +from funasr_detach.register import tables +from typing import List, Tuple, Dict, Any, Optional + +from funasr_detach.utils.datadir_writer import DatadirWriter +from funasr_detach.utils.load_utils import load_audio_text_image_video, extract_fbank + + +class VadStateMachine(Enum): + kVadInStateStartPointNotDetected = 1 + kVadInStateInSpeechSegment = 2 + kVadInStateEndPointDetected = 3 + + +class FrameState(Enum): + kFrameStateInvalid = -1 + kFrameStateSpeech = 1 + kFrameStateSil = 0 + + +# final voice/unvoice state per frame +class AudioChangeState(Enum): + kChangeStateSpeech2Speech = 0 + kChangeStateSpeech2Sil = 1 + kChangeStateSil2Sil = 2 + kChangeStateSil2Speech = 3 + kChangeStateNoBegin = 4 + kChangeStateInvalid = 5 + + +class VadDetectMode(Enum): + kVadSingleUtteranceDetectMode = 0 + kVadMutipleUtteranceDetectMode = 1 + + +class VADXOptions: + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Deep-FSMN for Large Vocabulary Continuous Speech Recognition + https://arxiv.org/abs/1803.05030 + """ + + def __init__( + self, + sample_rate: int = 16000, + detect_mode: int = VadDetectMode.kVadMutipleUtteranceDetectMode.value, + snr_mode: int = 0, + max_end_silence_time: int = 800, + max_start_silence_time: int = 3000, + do_start_point_detection: bool = True, + do_end_point_detection: bool = True, + window_size_ms: int = 200, + sil_to_speech_time_thres: int = 150, + speech_to_sil_time_thres: int = 150, + speech_2_noise_ratio: float = 1.0, + do_extend: int = 1, + lookback_time_start_point: int = 200, + lookahead_time_end_point: int = 100, + max_single_segment_time: int = 60000, + nn_eval_block_size: int = 8, + dcd_block_size: int = 4, + snr_thres: int = -100.0, + noise_frame_num_used_for_snr: int = 100, + decibel_thres: int = -100.0, + speech_noise_thres: float = 0.6, + fe_prior_thres: float = 1e-4, + silence_pdf_num: int = 1, + sil_pdf_ids: List[int] = [0], + speech_noise_thresh_low: float = -0.1, + speech_noise_thresh_high: float = 0.3, + output_frame_probs: bool = False, + frame_in_ms: int = 10, + frame_length_ms: int = 25, + **kwargs, + ): + self.sample_rate = sample_rate + self.detect_mode = detect_mode + self.snr_mode = snr_mode + self.max_end_silence_time = max_end_silence_time + self.max_start_silence_time = max_start_silence_time + self.do_start_point_detection = do_start_point_detection + self.do_end_point_detection = do_end_point_detection + self.window_size_ms = window_size_ms + self.sil_to_speech_time_thres = sil_to_speech_time_thres + self.speech_to_sil_time_thres = speech_to_sil_time_thres + self.speech_2_noise_ratio = speech_2_noise_ratio + self.do_extend = do_extend + self.lookback_time_start_point = lookback_time_start_point + self.lookahead_time_end_point = lookahead_time_end_point + self.max_single_segment_time = max_single_segment_time + self.nn_eval_block_size = nn_eval_block_size + self.dcd_block_size = dcd_block_size + self.snr_thres = snr_thres + self.noise_frame_num_used_for_snr = noise_frame_num_used_for_snr + self.decibel_thres = decibel_thres + self.speech_noise_thres = speech_noise_thres + self.fe_prior_thres = fe_prior_thres + self.silence_pdf_num = silence_pdf_num + self.sil_pdf_ids = sil_pdf_ids + self.speech_noise_thresh_low = speech_noise_thresh_low + self.speech_noise_thresh_high = speech_noise_thresh_high + self.output_frame_probs = output_frame_probs + self.frame_in_ms = frame_in_ms + self.frame_length_ms = frame_length_ms + + +class E2EVadSpeechBufWithDoa(object): + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Deep-FSMN for Large Vocabulary Continuous Speech Recognition + https://arxiv.org/abs/1803.05030 + """ + + def __init__(self): + self.start_ms = 0 + self.end_ms = 0 + self.buffer = [] + self.contain_seg_start_point = False + self.contain_seg_end_point = False + self.doa = 0 + + def Reset(self): + self.start_ms = 0 + self.end_ms = 0 + self.buffer = [] + self.contain_seg_start_point = False + self.contain_seg_end_point = False + self.doa = 0 + + +class E2EVadFrameProb(object): + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Deep-FSMN for Large Vocabulary Continuous Speech Recognition + https://arxiv.org/abs/1803.05030 + """ + + def __init__(self): + self.noise_prob = 0.0 + self.speech_prob = 0.0 + self.score = 0.0 + self.frame_id = 0 + self.frm_state = 0 + + +class WindowDetector(object): + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Deep-FSMN for Large Vocabulary Continuous Speech Recognition + https://arxiv.org/abs/1803.05030 + """ + + def __init__( + self, + window_size_ms: int, + sil_to_speech_time: int, + speech_to_sil_time: int, + frame_size_ms: int, + ): + self.window_size_ms = window_size_ms + self.sil_to_speech_time = sil_to_speech_time + self.speech_to_sil_time = speech_to_sil_time + self.frame_size_ms = frame_size_ms + + self.win_size_frame = int(window_size_ms / frame_size_ms) + self.win_sum = 0 + self.win_state = [0] * self.win_size_frame # 初始化窗 + + self.cur_win_pos = 0 + self.pre_frame_state = FrameState.kFrameStateSil + self.cur_frame_state = FrameState.kFrameStateSil + self.sil_to_speech_frmcnt_thres = int(sil_to_speech_time / frame_size_ms) + self.speech_to_sil_frmcnt_thres = int(speech_to_sil_time / frame_size_ms) + + self.voice_last_frame_count = 0 + self.noise_last_frame_count = 0 + self.hydre_frame_count = 0 + + def Reset(self) -> None: + self.cur_win_pos = 0 + self.win_sum = 0 + self.win_state = [0] * self.win_size_frame + self.pre_frame_state = FrameState.kFrameStateSil + self.cur_frame_state = FrameState.kFrameStateSil + self.voice_last_frame_count = 0 + self.noise_last_frame_count = 0 + self.hydre_frame_count = 0 + + def GetWinSize(self) -> int: + return int(self.win_size_frame) + + def DetectOneFrame( + self, frameState: FrameState, frame_count: int, cache: dict = {} + ) -> AudioChangeState: + cur_frame_state = FrameState.kFrameStateSil + if frameState == FrameState.kFrameStateSpeech: + cur_frame_state = 1 + elif frameState == FrameState.kFrameStateSil: + cur_frame_state = 0 + else: + return AudioChangeState.kChangeStateInvalid + self.win_sum -= self.win_state[self.cur_win_pos] + self.win_sum += cur_frame_state + self.win_state[self.cur_win_pos] = cur_frame_state + self.cur_win_pos = (self.cur_win_pos + 1) % self.win_size_frame + + if ( + self.pre_frame_state == FrameState.kFrameStateSil + and self.win_sum >= self.sil_to_speech_frmcnt_thres + ): + self.pre_frame_state = FrameState.kFrameStateSpeech + return AudioChangeState.kChangeStateSil2Speech + + if ( + self.pre_frame_state == FrameState.kFrameStateSpeech + and self.win_sum <= self.speech_to_sil_frmcnt_thres + ): + self.pre_frame_state = FrameState.kFrameStateSil + return AudioChangeState.kChangeStateSpeech2Sil + + if self.pre_frame_state == FrameState.kFrameStateSil: + return AudioChangeState.kChangeStateSil2Sil + if self.pre_frame_state == FrameState.kFrameStateSpeech: + return AudioChangeState.kChangeStateSpeech2Speech + return AudioChangeState.kChangeStateInvalid + + def FrameSizeMs(self) -> int: + return int(self.frame_size_ms) + + +class Stats(object): + def __init__( + self, + sil_pdf_ids, + max_end_sil_frame_cnt_thresh, + speech_noise_thres, + ): + self.data_buf_start_frame = 0 + self.frm_cnt = 0 + self.latest_confirmed_speech_frame = 0 + self.lastest_confirmed_silence_frame = -1 + self.continous_silence_frame_count = 0 + self.vad_state_machine = VadStateMachine.kVadInStateStartPointNotDetected + self.confirmed_start_frame = -1 + self.confirmed_end_frame = -1 + self.number_end_time_detected = 0 + self.sil_frame = 0 + self.sil_pdf_ids = sil_pdf_ids + self.noise_average_decibel = -100.0 + self.pre_end_silence_detected = False + self.next_seg = True + + self.output_data_buf = [] + self.output_data_buf_offset = 0 + self.frame_probs = [] + self.max_end_sil_frame_cnt_thresh = max_end_sil_frame_cnt_thresh + self.speech_noise_thres = speech_noise_thres + self.scores = None + self.max_time_out = False + self.decibel = [] + self.data_buf = None + self.data_buf_all = None + self.waveform = None + self.last_drop_frames = 0 + + +@tables.register("model_classes", "FsmnVADStreaming") +class FsmnVADStreaming(nn.Module): + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Deep-FSMN for Large Vocabulary Continuous Speech Recognition + https://arxiv.org/abs/1803.05030 + """ + + def __init__( + self, + encoder: str = None, + encoder_conf: Optional[Dict] = None, + vad_post_args: Dict[str, Any] = None, + **kwargs, + ): + super().__init__() + self.vad_opts = VADXOptions(**kwargs) + + encoder_class = tables.encoder_classes.get(encoder) + encoder = encoder_class(**encoder_conf) + self.encoder = encoder + + def ResetDetection(self, cache: dict = {}): + cache["stats"].continous_silence_frame_count = 0 + cache["stats"].latest_confirmed_speech_frame = 0 + cache["stats"].lastest_confirmed_silence_frame = -1 + cache["stats"].confirmed_start_frame = -1 + cache["stats"].confirmed_end_frame = -1 + cache["stats"].vad_state_machine = ( + VadStateMachine.kVadInStateStartPointNotDetected + ) + cache["windows_detector"].Reset() + cache["stats"].sil_frame = 0 + cache["stats"].frame_probs = [] + + if cache["stats"].output_data_buf: + assert cache["stats"].output_data_buf[-1].contain_seg_end_point == True + drop_frames = int( + cache["stats"].output_data_buf[-1].end_ms / self.vad_opts.frame_in_ms + ) + real_drop_frames = drop_frames - cache["stats"].last_drop_frames + cache["stats"].last_drop_frames = drop_frames + cache["stats"].data_buf_all = cache["stats"].data_buf_all[ + real_drop_frames + * int(self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000) : + ] + cache["stats"].decibel = cache["stats"].decibel[real_drop_frames:] + cache["stats"].scores = cache["stats"].scores[:, real_drop_frames:, :] + + def ComputeDecibel(self, cache: dict = {}) -> None: + frame_sample_length = int( + self.vad_opts.frame_length_ms * self.vad_opts.sample_rate / 1000 + ) + frame_shift_length = int( + self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000 + ) + if cache["stats"].data_buf_all is None: + cache["stats"].data_buf_all = cache["stats"].waveform[ + 0 + ] # cache["stats"].data_buf is pointed to cache["stats"].waveform[0] + cache["stats"].data_buf = cache["stats"].data_buf_all + else: + cache["stats"].data_buf_all = torch.cat( + (cache["stats"].data_buf_all, cache["stats"].waveform[0]) + ) + for offset in range( + 0, + cache["stats"].waveform.shape[1] - frame_sample_length + 1, + frame_shift_length, + ): + cache["stats"].decibel.append( + 10 + * math.log10( + (cache["stats"].waveform[0][offset : offset + frame_sample_length]) + .square() + .sum() + + 0.000001 + ) + ) + + def ComputeScores(self, feats: torch.Tensor, cache: dict = {}) -> None: + scores = self.encoder(feats, cache=cache["encoder"]).to( + "cpu" + ) # return B * T * D + assert ( + scores.shape[1] == feats.shape[1] + ), "The shape between feats and scores does not match" + self.vad_opts.nn_eval_block_size = scores.shape[1] + cache["stats"].frm_cnt += scores.shape[1] # count total frames + if cache["stats"].scores is None: + cache["stats"].scores = scores # the first calculation + else: + cache["stats"].scores = torch.cat((cache["stats"].scores, scores), dim=1) + + def PopDataBufTillFrame( + self, frame_idx: int, cache: dict = {} + ) -> None: # need check again + while cache["stats"].data_buf_start_frame < frame_idx: + if len(cache["stats"].data_buf) >= int( + self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000 + ): + cache["stats"].data_buf_start_frame += 1 + cache["stats"].data_buf = cache["stats"].data_buf_all[ + ( + cache["stats"].data_buf_start_frame + - cache["stats"].last_drop_frames + ) + * int( + self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000 + ) : + ] + + def PopDataToOutputBuf( + self, + start_frm: int, + frm_cnt: int, + first_frm_is_start_point: bool, + last_frm_is_end_point: bool, + end_point_is_sent_end: bool, + cache: dict = {}, + ) -> None: + self.PopDataBufTillFrame(start_frm, cache=cache) + expected_sample_number = int( + frm_cnt * self.vad_opts.sample_rate * self.vad_opts.frame_in_ms / 1000 + ) + if last_frm_is_end_point: + extra_sample = max( + 0, + int( + self.vad_opts.frame_length_ms * self.vad_opts.sample_rate / 1000 + - self.vad_opts.sample_rate * self.vad_opts.frame_in_ms / 1000 + ), + ) + expected_sample_number += int(extra_sample) + if end_point_is_sent_end: + expected_sample_number = max( + expected_sample_number, len(cache["stats"].data_buf) + ) + if len(cache["stats"].data_buf) < expected_sample_number: + print("error in calling pop data_buf\n") + + if len(cache["stats"].output_data_buf) == 0 or first_frm_is_start_point: + cache["stats"].output_data_buf.append(E2EVadSpeechBufWithDoa()) + cache["stats"].output_data_buf[-1].Reset() + cache["stats"].output_data_buf[-1].start_ms = ( + start_frm * self.vad_opts.frame_in_ms + ) + cache["stats"].output_data_buf[-1].end_ms = ( + cache["stats"].output_data_buf[-1].start_ms + ) + cache["stats"].output_data_buf[-1].doa = 0 + cur_seg = cache["stats"].output_data_buf[-1] + if cur_seg.end_ms != start_frm * self.vad_opts.frame_in_ms: + print("warning\n") + out_pos = len(cur_seg.buffer) # cur_seg.buff现在没做任何操作 + data_to_pop = 0 + if end_point_is_sent_end: + data_to_pop = expected_sample_number + else: + data_to_pop = int( + frm_cnt * self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000 + ) + if data_to_pop > len(cache["stats"].data_buf): + print('VAD data_to_pop is bigger than cache["stats"].data_buf.size()!!!\n') + data_to_pop = len(cache["stats"].data_buf) + expected_sample_number = len(cache["stats"].data_buf) + + cur_seg.doa = 0 + for sample_cpy_out in range(0, data_to_pop): + # cur_seg.buffer[out_pos ++] = data_buf_.back(); + out_pos += 1 + for sample_cpy_out in range(data_to_pop, expected_sample_number): + # cur_seg.buffer[out_pos++] = data_buf_.back() + out_pos += 1 + if cur_seg.end_ms != start_frm * self.vad_opts.frame_in_ms: + print("Something wrong with the VAD algorithm\n") + cache["stats"].data_buf_start_frame += frm_cnt + cur_seg.end_ms = (start_frm + frm_cnt) * self.vad_opts.frame_in_ms + if first_frm_is_start_point: + cur_seg.contain_seg_start_point = True + if last_frm_is_end_point: + cur_seg.contain_seg_end_point = True + + def OnSilenceDetected(self, valid_frame: int, cache: dict = {}): + cache["stats"].lastest_confirmed_silence_frame = valid_frame + if ( + cache["stats"].vad_state_machine + == VadStateMachine.kVadInStateStartPointNotDetected + ): + self.PopDataBufTillFrame(valid_frame, cache=cache) + + # silence_detected_callback_ + # pass + + def OnVoiceDetected(self, valid_frame: int, cache: dict = {}) -> None: + cache["stats"].latest_confirmed_speech_frame = valid_frame + self.PopDataToOutputBuf(valid_frame, 1, False, False, False, cache=cache) + + def OnVoiceStart( + self, start_frame: int, fake_result: bool = False, cache: dict = {} + ) -> None: + if self.vad_opts.do_start_point_detection: + pass + if cache["stats"].confirmed_start_frame != -1: + print("not reset vad properly\n") + else: + cache["stats"].confirmed_start_frame = start_frame + + if ( + not fake_result + and cache["stats"].vad_state_machine + == VadStateMachine.kVadInStateStartPointNotDetected + ): + self.PopDataToOutputBuf( + cache["stats"].confirmed_start_frame, 1, True, False, False, cache=cache + ) + + def OnVoiceEnd( + self, end_frame: int, fake_result: bool, is_last_frame: bool, cache: dict = {} + ) -> None: + for t in range(cache["stats"].latest_confirmed_speech_frame + 1, end_frame): + self.OnVoiceDetected(t, cache=cache) + if self.vad_opts.do_end_point_detection: + pass + if cache["stats"].confirmed_end_frame != -1: + print("not reset vad properly\n") + else: + cache["stats"].confirmed_end_frame = end_frame + if not fake_result: + cache["stats"].sil_frame = 0 + self.PopDataToOutputBuf( + cache["stats"].confirmed_end_frame, + 1, + False, + True, + is_last_frame, + cache=cache, + ) + cache["stats"].number_end_time_detected += 1 + + def MaybeOnVoiceEndIfLastFrame( + self, is_final_frame: bool, cur_frm_idx: int, cache: dict = {} + ) -> None: + if is_final_frame: + self.OnVoiceEnd(cur_frm_idx, False, True, cache=cache) + cache["stats"].vad_state_machine = ( + VadStateMachine.kVadInStateEndPointDetected + ) + + def GetLatency(self, cache: dict = {}) -> int: + return int( + self.LatencyFrmNumAtStartPoint(cache=cache) * self.vad_opts.frame_in_ms + ) + + def LatencyFrmNumAtStartPoint(self, cache: dict = {}) -> int: + vad_latency = cache["windows_detector"].GetWinSize() + if self.vad_opts.do_extend: + vad_latency += int( + self.vad_opts.lookback_time_start_point / self.vad_opts.frame_in_ms + ) + return vad_latency + + def GetFrameState(self, t: int, cache: dict = {}): + frame_state = FrameState.kFrameStateInvalid + cur_decibel = cache["stats"].decibel[t] + cur_snr = cur_decibel - cache["stats"].noise_average_decibel + # for each frame, calc log posterior probability of each state + if cur_decibel < self.vad_opts.decibel_thres: + frame_state = FrameState.kFrameStateSil + self.DetectOneFrame(frame_state, t, False, cache=cache) + return frame_state + + sum_score = 0.0 + noise_prob = 0.0 + assert len(cache["stats"].sil_pdf_ids) == self.vad_opts.silence_pdf_num + if len(cache["stats"].sil_pdf_ids) > 0: + assert len(cache["stats"].scores) == 1 # 只支持batch_size = 1的测试 + sil_pdf_scores = [ + cache["stats"].scores[0][t][sil_pdf_id] + for sil_pdf_id in cache["stats"].sil_pdf_ids + ] + sum_score = sum(sil_pdf_scores) + noise_prob = math.log(sum_score) * self.vad_opts.speech_2_noise_ratio + total_score = 1.0 + sum_score = total_score - sum_score + speech_prob = math.log(sum_score) + if self.vad_opts.output_frame_probs: + frame_prob = E2EVadFrameProb() + frame_prob.noise_prob = noise_prob + frame_prob.speech_prob = speech_prob + frame_prob.score = sum_score + frame_prob.frame_id = t + cache["stats"].frame_probs.append(frame_prob) + if ( + math.exp(speech_prob) + >= math.exp(noise_prob) + cache["stats"].speech_noise_thres + ): + if ( + cur_snr >= self.vad_opts.snr_thres + and cur_decibel >= self.vad_opts.decibel_thres + ): + frame_state = FrameState.kFrameStateSpeech + else: + frame_state = FrameState.kFrameStateSil + else: + frame_state = FrameState.kFrameStateSil + if cache["stats"].noise_average_decibel < -99.9: + cache["stats"].noise_average_decibel = cur_decibel + else: + cache["stats"].noise_average_decibel = ( + cur_decibel + + cache["stats"].noise_average_decibel + * (self.vad_opts.noise_frame_num_used_for_snr - 1) + ) / self.vad_opts.noise_frame_num_used_for_snr + + return frame_state + + def forward( + self, + feats: torch.Tensor, + waveform: torch.tensor, + cache: dict = {}, + is_final: bool = False, + **kwargs, + ): + # if len(cache) == 0: + # self.AllResetDetection() + # self.waveform = waveform # compute decibel for each frame + cache["stats"].waveform = waveform + is_streaming_input = kwargs.get("is_streaming_input", True) + self.ComputeDecibel(cache=cache) + self.ComputeScores(feats, cache=cache) + if not is_final: + self.DetectCommonFrames(cache=cache) + else: + self.DetectLastFrames(cache=cache) + segments = [] + for batch_num in range(0, feats.shape[0]): # only support batch_size = 1 now + segment_batch = [] + if len(cache["stats"].output_data_buf) > 0: + for i in range( + cache["stats"].output_data_buf_offset, + len(cache["stats"].output_data_buf), + ): + if ( + is_streaming_input + ): # in this case, return [beg, -1], [], [-1, end], [beg, end] + if ( + not cache["stats"] + .output_data_buf[i] + .contain_seg_start_point + ): + continue + if ( + not cache["stats"].next_seg + and not cache["stats"] + .output_data_buf[i] + .contain_seg_end_point + ): + continue + start_ms = ( + cache["stats"].output_data_buf[i].start_ms + if cache["stats"].next_seg + else -1 + ) + if cache["stats"].output_data_buf[i].contain_seg_end_point: + end_ms = cache["stats"].output_data_buf[i].end_ms + cache["stats"].next_seg = True + cache["stats"].output_data_buf_offset += 1 + else: + end_ms = -1 + cache["stats"].next_seg = False + segment = [start_ms, end_ms] + + else: # in this case, return [beg, end] + + if not is_final and ( + not cache["stats"] + .output_data_buf[i] + .contain_seg_start_point + or not cache["stats"] + .output_data_buf[i] + .contain_seg_end_point + ): + continue + segment = [ + cache["stats"].output_data_buf[i].start_ms, + cache["stats"].output_data_buf[i].end_ms, + ] + cache[ + "stats" + ].output_data_buf_offset += 1 # need update this parameter + + segment_batch.append(segment) + + if segment_batch: + segments.append(segment_batch) + # if is_final: + # # reset class variables and clear the dict for the next query + # self.AllResetDetection() + return segments + + def init_cache(self, cache: dict = {}, **kwargs): + + cache["frontend"] = {} + cache["prev_samples"] = torch.empty(0) + cache["encoder"] = {} + windows_detector = WindowDetector( + self.vad_opts.window_size_ms, + self.vad_opts.sil_to_speech_time_thres, + self.vad_opts.speech_to_sil_time_thres, + self.vad_opts.frame_in_ms, + ) + windows_detector.Reset() + + stats = Stats( + sil_pdf_ids=self.vad_opts.sil_pdf_ids, + max_end_sil_frame_cnt_thresh=self.vad_opts.max_end_silence_time + - self.vad_opts.speech_to_sil_time_thres, + speech_noise_thres=self.vad_opts.speech_noise_thres, + ) + cache["windows_detector"] = windows_detector + cache["stats"] = stats + return cache + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + cache: dict = {}, + **kwargs, + ): + + if len(cache) == 0: + self.init_cache(cache, **kwargs) + + meta_data = {} + chunk_size = kwargs.get("chunk_size", 60000) # 50ms + chunk_stride_samples = int(chunk_size * frontend.fs / 1000) + + time1 = time.perf_counter() + is_streaming_input = ( + kwargs.get("is_streaming_input", False) + if chunk_size >= 15000 + else kwargs.get("is_streaming_input", True) + ) + is_final = ( + kwargs.get("is_final", False) + if is_streaming_input + else kwargs.get("is_final", True) + ) + cfg = {"is_final": is_final, "is_streaming_input": is_streaming_input} + audio_sample_list = load_audio_text_image_video( + data_in, + fs=frontend.fs, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + cache=cfg, + ) + _is_final = cfg["is_final"] # if data_in is a file or url, set is_final=True + is_streaming_input = cfg["is_streaming_input"] + time2 = time.perf_counter() + meta_data["load_data"] = f"{time2 - time1:0.3f}" + assert len(audio_sample_list) == 1, "batch_size must be set 1" + + audio_sample = torch.cat((cache["prev_samples"], audio_sample_list[0])) + + n = int(len(audio_sample) // chunk_stride_samples + int(_is_final)) + m = int(len(audio_sample) % chunk_stride_samples * (1 - int(_is_final))) + segments = [] + for i in range(n): + kwargs["is_final"] = _is_final and i == n - 1 + audio_sample_i = audio_sample[ + i * chunk_stride_samples : (i + 1) * chunk_stride_samples + ] + + # extract fbank feats + speech, speech_lengths = extract_fbank( + [audio_sample_i], + data_type=kwargs.get("data_type", "sound"), + frontend=frontend, + cache=cache["frontend"], + is_final=kwargs["is_final"], + ) + time3 = time.perf_counter() + meta_data["extract_feat"] = f"{time3 - time2:0.3f}" + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() + * frontend.frame_shift + * frontend.lfr_n + / 1000 + ) + speech = speech.to(device=kwargs["device"]) + speech_lengths = speech_lengths.to(device=kwargs["device"]) + + batch = { + "feats": speech, + "waveform": cache["frontend"]["waveforms"], + "is_final": kwargs["is_final"], + "cache": cache, + "is_streaming_input": is_streaming_input, + } + segments_i = self.forward(**batch) + if len(segments_i) > 0: + segments.extend(*segments_i) + + cache["prev_samples"] = audio_sample[:-m] + if _is_final: + self.init_cache(cache) + + ibest_writer = None + if kwargs.get("output_dir") is not None: + if not hasattr(self, "writer"): + self.writer = DatadirWriter(kwargs.get("output_dir")) + ibest_writer = self.writer[f"{1}best_recog"] + + results = [] + result_i = {"key": key[0], "value": segments} + if ( + "MODELSCOPE_ENVIRONMENT" in os.environ + and os.environ["MODELSCOPE_ENVIRONMENT"] == "eas" + ): + result_i = json.dumps(result_i) + + results.append(result_i) + + if ibest_writer is not None: + ibest_writer["text"][key[0]] = segments + + return results, meta_data + + def DetectCommonFrames(self, cache: dict = {}) -> int: + if ( + cache["stats"].vad_state_machine + == VadStateMachine.kVadInStateEndPointDetected + ): + return 0 + for i in range(self.vad_opts.nn_eval_block_size - 1, -1, -1): + frame_state = FrameState.kFrameStateInvalid + frame_state = self.GetFrameState( + cache["stats"].frm_cnt - 1 - i - cache["stats"].last_drop_frames, + cache=cache, + ) + self.DetectOneFrame( + frame_state, cache["stats"].frm_cnt - 1 - i, False, cache=cache + ) + + return 0 + + def DetectLastFrames(self, cache: dict = {}) -> int: + if ( + cache["stats"].vad_state_machine + == VadStateMachine.kVadInStateEndPointDetected + ): + return 0 + for i in range(self.vad_opts.nn_eval_block_size - 1, -1, -1): + frame_state = FrameState.kFrameStateInvalid + frame_state = self.GetFrameState( + cache["stats"].frm_cnt - 1 - i - cache["stats"].last_drop_frames, + cache=cache, + ) + if i != 0: + self.DetectOneFrame( + frame_state, cache["stats"].frm_cnt - 1 - i, False, cache=cache + ) + else: + self.DetectOneFrame( + frame_state, cache["stats"].frm_cnt - 1, True, cache=cache + ) + + return 0 + + def DetectOneFrame( + self, + cur_frm_state: FrameState, + cur_frm_idx: int, + is_final_frame: bool, + cache: dict = {}, + ) -> None: + tmp_cur_frm_state = FrameState.kFrameStateInvalid + if cur_frm_state == FrameState.kFrameStateSpeech: + if math.fabs(1.0) > self.vad_opts.fe_prior_thres: + tmp_cur_frm_state = FrameState.kFrameStateSpeech + else: + tmp_cur_frm_state = FrameState.kFrameStateSil + elif cur_frm_state == FrameState.kFrameStateSil: + tmp_cur_frm_state = FrameState.kFrameStateSil + state_change = cache["windows_detector"].DetectOneFrame( + tmp_cur_frm_state, cur_frm_idx, cache=cache + ) + frm_shift_in_ms = self.vad_opts.frame_in_ms + if AudioChangeState.kChangeStateSil2Speech == state_change: + silence_frame_count = cache["stats"].continous_silence_frame_count + cache["stats"].continous_silence_frame_count = 0 + cache["stats"].pre_end_silence_detected = False + start_frame = 0 + if ( + cache["stats"].vad_state_machine + == VadStateMachine.kVadInStateStartPointNotDetected + ): + start_frame = max( + cache["stats"].data_buf_start_frame, + cur_frm_idx - self.LatencyFrmNumAtStartPoint(cache=cache), + ) + self.OnVoiceStart(start_frame, cache=cache) + cache["stats"].vad_state_machine = ( + VadStateMachine.kVadInStateInSpeechSegment + ) + for t in range(start_frame + 1, cur_frm_idx + 1): + self.OnVoiceDetected(t, cache=cache) + elif ( + cache["stats"].vad_state_machine + == VadStateMachine.kVadInStateInSpeechSegment + ): + for t in range( + cache["stats"].latest_confirmed_speech_frame + 1, cur_frm_idx + ): + self.OnVoiceDetected(t, cache=cache) + if ( + cur_frm_idx - cache["stats"].confirmed_start_frame + 1 + > self.vad_opts.max_single_segment_time / frm_shift_in_ms + ): + self.OnVoiceEnd(cur_frm_idx, False, False, cache=cache) + cache["stats"].vad_state_machine = ( + VadStateMachine.kVadInStateEndPointDetected + ) + elif not is_final_frame: + self.OnVoiceDetected(cur_frm_idx, cache=cache) + else: + self.MaybeOnVoiceEndIfLastFrame( + is_final_frame, cur_frm_idx, cache=cache + ) + else: + pass + elif AudioChangeState.kChangeStateSpeech2Sil == state_change: + cache["stats"].continous_silence_frame_count = 0 + if ( + cache["stats"].vad_state_machine + == VadStateMachine.kVadInStateStartPointNotDetected + ): + pass + elif ( + cache["stats"].vad_state_machine + == VadStateMachine.kVadInStateInSpeechSegment + ): + if ( + cur_frm_idx - cache["stats"].confirmed_start_frame + 1 + > self.vad_opts.max_single_segment_time / frm_shift_in_ms + ): + self.OnVoiceEnd(cur_frm_idx, False, False, cache=cache) + cache["stats"].vad_state_machine = ( + VadStateMachine.kVadInStateEndPointDetected + ) + elif not is_final_frame: + self.OnVoiceDetected(cur_frm_idx, cache=cache) + else: + self.MaybeOnVoiceEndIfLastFrame( + is_final_frame, cur_frm_idx, cache=cache + ) + else: + pass + elif AudioChangeState.kChangeStateSpeech2Speech == state_change: + cache["stats"].continous_silence_frame_count = 0 + if ( + cache["stats"].vad_state_machine + == VadStateMachine.kVadInStateInSpeechSegment + ): + if ( + cur_frm_idx - cache["stats"].confirmed_start_frame + 1 + > self.vad_opts.max_single_segment_time / frm_shift_in_ms + ): + cache["stats"].max_time_out = True + self.OnVoiceEnd(cur_frm_idx, False, False, cache=cache) + cache["stats"].vad_state_machine = ( + VadStateMachine.kVadInStateEndPointDetected + ) + elif not is_final_frame: + self.OnVoiceDetected(cur_frm_idx, cache=cache) + else: + self.MaybeOnVoiceEndIfLastFrame( + is_final_frame, cur_frm_idx, cache=cache + ) + else: + pass + elif AudioChangeState.kChangeStateSil2Sil == state_change: + cache["stats"].continous_silence_frame_count += 1 + if ( + cache["stats"].vad_state_machine + == VadStateMachine.kVadInStateStartPointNotDetected + ): + # silence timeout, return zero length decision + if ( + ( + self.vad_opts.detect_mode + == VadDetectMode.kVadSingleUtteranceDetectMode.value + ) + and ( + cache["stats"].continous_silence_frame_count * frm_shift_in_ms + > self.vad_opts.max_start_silence_time + ) + ) or (is_final_frame and cache["stats"].number_end_time_detected == 0): + for t in range( + cache["stats"].lastest_confirmed_silence_frame + 1, cur_frm_idx + ): + self.OnSilenceDetected(t, cache=cache) + self.OnVoiceStart(0, True, cache=cache) + self.OnVoiceEnd(0, True, False, cache=cache) + cache["stats"].vad_state_machine = ( + VadStateMachine.kVadInStateEndPointDetected + ) + else: + if cur_frm_idx >= self.LatencyFrmNumAtStartPoint(cache=cache): + self.OnSilenceDetected( + cur_frm_idx - self.LatencyFrmNumAtStartPoint(cache=cache), + cache=cache, + ) + elif ( + cache["stats"].vad_state_machine + == VadStateMachine.kVadInStateInSpeechSegment + ): + if ( + cache["stats"].continous_silence_frame_count * frm_shift_in_ms + >= cache["stats"].max_end_sil_frame_cnt_thresh + ): + lookback_frame = int( + cache["stats"].max_end_sil_frame_cnt_thresh / frm_shift_in_ms + ) + if self.vad_opts.do_extend: + lookback_frame -= int( + self.vad_opts.lookahead_time_end_point / frm_shift_in_ms + ) + lookback_frame -= 1 + lookback_frame = max(0, lookback_frame) + self.OnVoiceEnd( + cur_frm_idx - lookback_frame, False, False, cache=cache + ) + cache["stats"].vad_state_machine = ( + VadStateMachine.kVadInStateEndPointDetected + ) + elif ( + cur_frm_idx - cache["stats"].confirmed_start_frame + 1 + > self.vad_opts.max_single_segment_time / frm_shift_in_ms + ): + self.OnVoiceEnd(cur_frm_idx, False, False, cache=cache) + cache["stats"].vad_state_machine = ( + VadStateMachine.kVadInStateEndPointDetected + ) + elif self.vad_opts.do_extend and not is_final_frame: + if cache["stats"].continous_silence_frame_count <= int( + self.vad_opts.lookahead_time_end_point / frm_shift_in_ms + ): + self.OnVoiceDetected(cur_frm_idx, cache=cache) + else: + self.MaybeOnVoiceEndIfLastFrame( + is_final_frame, cur_frm_idx, cache=cache + ) + else: + pass + + if ( + cache["stats"].vad_state_machine + == VadStateMachine.kVadInStateEndPointDetected + and self.vad_opts.detect_mode + == VadDetectMode.kVadMutipleUtteranceDetectMode.value + ): + self.ResetDetection(cache=cache) diff --git a/almeval/models/stepaudio/funasr_detach/models/fsmn_vad_streaming/template.yaml b/almeval/models/stepaudio/funasr_detach/models/fsmn_vad_streaming/template.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e8a3a4f30d823305e9b698a44b695cbcb33feea3 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/fsmn_vad_streaming/template.yaml @@ -0,0 +1,62 @@ +# This is an example that demonstrates how to configure a model file. +# You can modify the configuration according to your own requirements. + +# to print the register_table: +# from funasr.register import tables +# tables.print() + +# network architecture +model: FsmnVADStreaming +model_conf: + sample_rate: 16000 + detect_mode: 1 + snr_mode: 0 + max_end_silence_time: 800 + max_start_silence_time: 3000 + do_start_point_detection: True + do_end_point_detection: True + window_size_ms: 200 + sil_to_speech_time_thres: 150 + speech_to_sil_time_thres: 150 + speech_2_noise_ratio: 1.0 + do_extend: 1 + lookback_time_start_point: 200 + lookahead_time_end_point: 100 + max_single_segment_time: 60000 + snr_thres: -100.0 + noise_frame_num_used_for_snr: 100 + decibel_thres: -100.0 + speech_noise_thres: 0.6 + fe_prior_thres: 0.0001 + silence_pdf_num: 1 + sil_pdf_ids: [0] + speech_noise_thresh_low: -0.1 + speech_noise_thresh_high: 0.3 + output_frame_probs: False + frame_in_ms: 10 + frame_length_ms: 25 + +encoder: FSMN +encoder_conf: + input_dim: 400 + input_affine_dim: 140 + fsmn_layers: 4 + linear_dim: 250 + proj_dim: 128 + lorder: 20 + rorder: 0 + lstride: 1 + rstride: 0 + output_affine_dim: 140 + output_dim: 248 + +frontend: WavFrontend +frontend_conf: + fs: 16000 + window: hamming + n_mels: 80 + frame_length: 25 + frame_shift: 10 + dither: 0.0 + lfr_m: 5 + lfr_n: 1 diff --git a/almeval/models/stepaudio/funasr_detach/models/mfcca/__init__.py b/almeval/models/stepaudio/funasr_detach/models/mfcca/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/almeval/models/stepaudio/funasr_detach/models/mfcca/e2e_asr_mfcca.py b/almeval/models/stepaudio/funasr_detach/models/mfcca/e2e_asr_mfcca.py new file mode 100644 index 0000000000000000000000000000000000000000..fba061fb56d2283c1740c63cd7049f167f5c9d63 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/mfcca/e2e_asr_mfcca.py @@ -0,0 +1,331 @@ +from contextlib import contextmanager +from distutils.version import LooseVersion +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union +import logging +import torch + +from funasr_detach.metrics import ErrorCalculator +from funasr_detach.metrics.compute_acc import th_accuracy +from funasr_detach.models.transformer.utils.add_sos_eos import add_sos_eos +from funasr_detach.losses.label_smoothing_loss import ( + LabelSmoothingLoss, # noqa: H301 +) +from funasr_detach.models.ctc import CTC +from funasr_detach.models.decoder.abs_decoder import AbsDecoder +from funasr_detach.models.encoder.abs_encoder import AbsEncoder +from funasr_detach.frontends.abs_frontend import AbsFrontend +from funasr_detach.models.preencoder.abs_preencoder import AbsPreEncoder +from funasr_detach.models.specaug.abs_specaug import AbsSpecAug +from funasr_detach.layers.abs_normalize import AbsNormalize +from funasr_detach.train_utils.device_funcs import force_gatherable +from funasr_detach.models.base_model import FunASRModel + +if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): + from torch.cuda.amp import autocast +else: + # Nothing to do if torch<1.6.0 + @contextmanager + def autocast(enabled=True): + yield + + +import pdb +import random +import math + + +class MFCCA(FunASRModel): + """ + Author: Audio, Speech and Language Processing Group (ASLP@NPU), Northwestern Polytechnical University + MFCCA:Multi-Frame Cross-Channel attention for multi-speaker ASR in Multi-party meeting scenario + https://arxiv.org/abs/2210.05265 + """ + + def __init__( + self, + vocab_size: int, + token_list: Union[Tuple[str, ...], List[str]], + frontend: Optional[AbsFrontend], + specaug: Optional[AbsSpecAug], + normalize: Optional[AbsNormalize], + encoder: AbsEncoder, + decoder: AbsDecoder, + ctc: CTC, + rnnt_decoder: None = None, + ctc_weight: float = 0.5, + ignore_id: int = -1, + lsm_weight: float = 0.0, + mask_ratio: float = 0.0, + length_normalized_loss: bool = False, + report_cer: bool = True, + report_wer: bool = True, + sym_space: str = "", + sym_blank: str = "", + preencoder: Optional[AbsPreEncoder] = None, + ): + assert 0.0 <= ctc_weight <= 1.0, ctc_weight + assert rnnt_decoder is None, "Not implemented" + + super().__init__() + # note that eos is the same as sos (equivalent ID) + self.sos = vocab_size - 1 + self.eos = vocab_size - 1 + self.vocab_size = vocab_size + self.ignore_id = ignore_id + self.ctc_weight = ctc_weight + self.token_list = token_list.copy() + + self.mask_ratio = mask_ratio + + self.frontend = frontend + self.specaug = specaug + self.normalize = normalize + self.preencoder = preencoder + self.encoder = encoder + # we set self.decoder = None in the CTC mode since + # self.decoder parameters were never used and PyTorch complained + # and threw an Exception in the multi-GPU experiment. + # thanks Jeff Farris for pointing out the issue. + if ctc_weight == 1.0: + self.decoder = None + else: + self.decoder = decoder + if ctc_weight == 0.0: + self.ctc = None + else: + self.ctc = ctc + self.rnnt_decoder = rnnt_decoder + self.criterion_att = LabelSmoothingLoss( + size=vocab_size, + padding_idx=ignore_id, + smoothing=lsm_weight, + normalize_length=length_normalized_loss, + ) + + if report_cer or report_wer: + self.error_calculator = ErrorCalculator( + token_list, sym_space, sym_blank, report_cer, report_wer + ) + else: + self.error_calculator = None + + def forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: + """Frontend + Encoder + Decoder + Calc loss + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + text: (Batch, Length) + text_lengths: (Batch,) + """ + assert text_lengths.dim() == 1, text_lengths.shape + # Check that batch_size is unified + assert ( + speech.shape[0] + == speech_lengths.shape[0] + == text.shape[0] + == text_lengths.shape[0] + ), (speech.shape, speech_lengths.shape, text.shape, text_lengths.shape) + # pdb.set_trace() + if speech.dim() == 3 and speech.size(2) == 8 and self.mask_ratio != 0: + rate_num = random.random() + # rate_num = 0.1 + if rate_num <= self.mask_ratio: + retain_channel = math.ceil(random.random() * 8) + if retain_channel > 1: + speech = speech[ + :, :, torch.randperm(8)[0:retain_channel].sort().values + ] + else: + speech = speech[:, :, torch.randperm(8)[0]] + # pdb.set_trace() + batch_size = speech.shape[0] + # for data-parallel + text = text[:, : text_lengths.max()] + + # 1. Encoder + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + + # 2a. Attention-decoder branch + if self.ctc_weight == 1.0: + loss_att, acc_att, cer_att, wer_att = None, None, None, None + else: + loss_att, acc_att, cer_att, wer_att = self._calc_att_loss( + encoder_out, encoder_out_lens, text, text_lengths + ) + + # 2b. CTC branch + if self.ctc_weight == 0.0: + loss_ctc, cer_ctc = None, None + else: + loss_ctc, cer_ctc = self._calc_ctc_loss( + encoder_out, encoder_out_lens, text, text_lengths + ) + + # 2c. RNN-T branch + if self.rnnt_decoder is not None: + _ = self._calc_rnnt_loss(encoder_out, encoder_out_lens, text, text_lengths) + + if self.ctc_weight == 0.0: + loss = loss_att + elif self.ctc_weight == 1.0: + loss = loss_ctc + else: + loss = self.ctc_weight * loss_ctc + (1 - self.ctc_weight) * loss_att + + stats = dict( + loss=loss.detach(), + loss_att=loss_att.detach() if loss_att is not None else None, + loss_ctc=loss_ctc.detach() if loss_ctc is not None else None, + acc=acc_att, + cer=cer_att, + wer=wer_att, + cer_ctc=cer_ctc, + ) + + # force_gatherable: to-device and to-tensor if scalar for DataParallel + loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) + return loss, stats, weight + + def collect_feats( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, + ) -> Dict[str, torch.Tensor]: + feats, feats_lengths, channel_size = self._extract_feats(speech, speech_lengths) + return {"feats": feats, "feats_lengths": feats_lengths} + + def encode( + self, speech: torch.Tensor, speech_lengths: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Frontend + Encoder. Note that this method is used by asr_inference.py + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + """ + with autocast(False): + # 1. Extract feats + feats, feats_lengths, channel_size = self._extract_feats( + speech, speech_lengths + ) + # 2. Data augmentation + if self.specaug is not None and self.training: + feats, feats_lengths = self.specaug(feats, feats_lengths) + + # 3. Normalization for feature: e.g. Global-CMVN, Utterance-CMVN + if self.normalize is not None: + feats, feats_lengths = self.normalize(feats, feats_lengths) + + # Pre-encoder, e.g. used for raw input data + if self.preencoder is not None: + feats, feats_lengths = self.preencoder(feats, feats_lengths) + # pdb.set_trace() + encoder_out, encoder_out_lens, _ = self.encoder( + feats, feats_lengths, channel_size + ) + + assert encoder_out.size(0) == speech.size(0), ( + encoder_out.size(), + speech.size(0), + ) + if encoder_out.dim() == 4: + assert encoder_out.size(2) <= encoder_out_lens.max(), ( + encoder_out.size(), + encoder_out_lens.max(), + ) + else: + assert encoder_out.size(1) <= encoder_out_lens.max(), ( + encoder_out.size(), + encoder_out_lens.max(), + ) + + return encoder_out, encoder_out_lens + + def _extract_feats( + self, speech: torch.Tensor, speech_lengths: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + assert speech_lengths.dim() == 1, speech_lengths.shape + # for data-parallel + speech = speech[:, : speech_lengths.max()] + if self.frontend is not None: + # Frontend + # e.g. STFT and Feature extract + # data_loader may send time-domain signal in this case + # speech (Batch, NSamples) -> feats: (Batch, NFrames, Dim) + feats, feats_lengths, channel_size = self.frontend(speech, speech_lengths) + else: + # No frontend and no feature extract + feats, feats_lengths = speech, speech_lengths + channel_size = 1 + return feats, feats_lengths, channel_size + + def _calc_att_loss( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + ): + ys_in_pad, ys_out_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id) + ys_in_lens = ys_pad_lens + 1 + + # 1. Forward decoder + decoder_out, _ = self.decoder( + encoder_out, encoder_out_lens, ys_in_pad, ys_in_lens + ) + + # 2. Compute attention loss + loss_att = self.criterion_att(decoder_out, ys_out_pad) + acc_att = th_accuracy( + decoder_out.view(-1, self.vocab_size), + ys_out_pad, + ignore_label=self.ignore_id, + ) + + # Compute cer/wer using attention-decoder + if self.training or self.error_calculator is None: + cer_att, wer_att = None, None + else: + ys_hat = decoder_out.argmax(dim=-1) + cer_att, wer_att = self.error_calculator(ys_hat.cpu(), ys_pad.cpu()) + + return loss_att, acc_att, cer_att, wer_att + + def _calc_ctc_loss( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + ): + # Calc CTC loss + if encoder_out.dim() == 4: + encoder_out = encoder_out.mean(1) + loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens) + + # Calc CER using CTC + cer_ctc = None + if not self.training and self.error_calculator is not None: + ys_hat = self.ctc.argmax(encoder_out).data + cer_ctc = self.error_calculator(ys_hat.cpu(), ys_pad.cpu(), is_ctc=True) + return loss_ctc, cer_ctc + + def _calc_rnnt_loss( + self, + encoder_out: torch.Tensor, + encoder_out_lens: torch.Tensor, + ys_pad: torch.Tensor, + ys_pad_lens: torch.Tensor, + ): + raise NotImplementedError diff --git a/almeval/models/stepaudio/funasr_detach/models/mfcca/encoder_layer_mfcca.py b/almeval/models/stepaudio/funasr_detach/models/mfcca/encoder_layer_mfcca.py new file mode 100644 index 0000000000000000000000000000000000000000..09567f07a7218a3004cbaa3c07084a9c6d8cd2f8 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/mfcca/encoder_layer_mfcca.py @@ -0,0 +1,275 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +# Copyright 2020 Johns Hopkins University (Shinji Watanabe) +# Northwestern Polytechnical University (Pengcheng Guo) +# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +"""Encoder self-attention layer definition.""" + +import torch + +from torch import nn + +from funasr_detach.models.transformer.layer_norm import LayerNorm +from torch.autograd import Variable + + +class Encoder_Conformer_Layer(nn.Module): + """Encoder layer module. + + Args: + size (int): Input dimension. + self_attn (torch.nn.Module): Self-attention module instance. + `MultiHeadedAttention` or `RelPositionMultiHeadedAttention` instance + can be used as the argument. + feed_forward (torch.nn.Module): Feed-forward module instance. + `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance + can be used as the argument. + feed_forward_macaron (torch.nn.Module): Additional feed-forward module instance. + `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance + can be used as the argument. + conv_module (torch.nn.Module): Convolution module instance. + `ConvlutionModule` instance can be used as the argument. + dropout_rate (float): Dropout rate. + normalize_before (bool): Whether to use layer_norm before the first block. + concat_after (bool): Whether to concat attention layer's input and output. + if True, additional linear will be applied. + i.e. x -> x + linear(concat(x, att(x))) + if False, no additional linear will be applied. i.e. x -> x + att(x) + + """ + + def __init__( + self, + size, + self_attn, + feed_forward, + feed_forward_macaron, + conv_module, + dropout_rate, + normalize_before=True, + concat_after=False, + cca_pos=0, + ): + """Construct an Encoder_Conformer_Layer object.""" + super(Encoder_Conformer_Layer, self).__init__() + self.self_attn = self_attn + self.feed_forward = feed_forward + self.feed_forward_macaron = feed_forward_macaron + self.conv_module = conv_module + self.norm_ff = LayerNorm(size) # for the FNN module + self.norm_mha = LayerNorm(size) # for the MHA module + if feed_forward_macaron is not None: + self.norm_ff_macaron = LayerNorm(size) + self.ff_scale = 0.5 + else: + self.ff_scale = 1.0 + if self.conv_module is not None: + self.norm_conv = LayerNorm(size) # for the CNN module + self.norm_final = LayerNorm(size) # for the final output of the block + self.dropout = nn.Dropout(dropout_rate) + self.size = size + self.normalize_before = normalize_before + self.concat_after = concat_after + self.cca_pos = cca_pos + + if self.concat_after: + self.concat_linear = nn.Linear(size + size, size) + + def forward(self, x_input, mask, cache=None): + """Compute encoded features. + + Args: + x_input (Union[Tuple, torch.Tensor]): Input tensor w/ or w/o pos emb. + - w/ pos emb: Tuple of tensors [(#batch, time, size), (1, time, size)]. + - w/o pos emb: Tensor (#batch, time, size). + mask (torch.Tensor): Mask tensor for the input (#batch, time). + cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size). + + Returns: + torch.Tensor: Output tensor (#batch, time, size). + torch.Tensor: Mask tensor (#batch, time). + + """ + if isinstance(x_input, tuple): + x, pos_emb = x_input[0], x_input[1] + else: + x, pos_emb = x_input, None + # whether to use macaron style + if self.feed_forward_macaron is not None: + residual = x + if self.normalize_before: + x = self.norm_ff_macaron(x) + x = residual + self.ff_scale * self.dropout(self.feed_forward_macaron(x)) + if not self.normalize_before: + x = self.norm_ff_macaron(x) + + # multi-headed self-attention module + residual = x + if self.normalize_before: + x = self.norm_mha(x) + + if cache is None: + x_q = x + else: + assert cache.shape == (x.shape[0], x.shape[1] - 1, self.size) + x_q = x[:, -1:, :] + residual = residual[:, -1:, :] + mask = None if mask is None else mask[:, -1:, :] + + if self.cca_pos < 2: + if pos_emb is not None: + x_att = self.self_attn(x_q, x, x, pos_emb, mask) + else: + x_att = self.self_attn(x_q, x, x, mask) + else: + x_att = self.self_attn(x_q, x, x, mask) + + if self.concat_after: + x_concat = torch.cat((x, x_att), dim=-1) + x = residual + self.concat_linear(x_concat) + else: + x = residual + self.dropout(x_att) + if not self.normalize_before: + x = self.norm_mha(x) + + # convolution module + if self.conv_module is not None: + residual = x + if self.normalize_before: + x = self.norm_conv(x) + x = residual + self.dropout(self.conv_module(x)) + if not self.normalize_before: + x = self.norm_conv(x) + + # feed forward module + residual = x + if self.normalize_before: + x = self.norm_ff(x) + x = residual + self.ff_scale * self.dropout(self.feed_forward(x)) + if not self.normalize_before: + x = self.norm_ff(x) + + if self.conv_module is not None: + x = self.norm_final(x) + + if cache is not None: + x = torch.cat([cache, x], dim=1) + + if pos_emb is not None: + return (x, pos_emb), mask + + return x, mask + + +class EncoderLayer(nn.Module): + """Encoder layer module. + + Args: + size (int): Input dimension. + self_attn (torch.nn.Module): Self-attention module instance. + `MultiHeadedAttention` or `RelPositionMultiHeadedAttention` instance + can be used as the argument. + feed_forward (torch.nn.Module): Feed-forward module instance. + `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance + can be used as the argument. + feed_forward_macaron (torch.nn.Module): Additional feed-forward module instance. + `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance + can be used as the argument. + conv_module (torch.nn.Module): Convolution module instance. + `ConvlutionModule` instance can be used as the argument. + dropout_rate (float): Dropout rate. + normalize_before (bool): Whether to use layer_norm before the first block. + concat_after (bool): Whether to concat attention layer's input and output. + if True, additional linear will be applied. + i.e. x -> x + linear(concat(x, att(x))) + if False, no additional linear will be applied. i.e. x -> x + att(x) + + """ + + def __init__( + self, + size, + self_attn_cros_channel, + self_attn_conformer, + feed_forward_csa, + feed_forward_macaron_csa, + conv_module_csa, + dropout_rate, + normalize_before=True, + concat_after=False, + ): + """Construct an EncoderLayer object.""" + super(EncoderLayer, self).__init__() + + self.encoder_cros_channel_atten = self_attn_cros_channel + self.encoder_csa = Encoder_Conformer_Layer( + size, + self_attn_conformer, + feed_forward_csa, + feed_forward_macaron_csa, + conv_module_csa, + dropout_rate, + normalize_before, + concat_after, + cca_pos=0, + ) + self.norm_mha = LayerNorm(size) # for the MHA module + self.dropout = nn.Dropout(dropout_rate) + + def forward(self, x_input, mask, channel_size, cache=None): + """Compute encoded features. + + Args: + x_input (Union[Tuple, torch.Tensor]): Input tensor w/ or w/o pos emb. + - w/ pos emb: Tuple of tensors [(#batch, time, size), (1, time, size)]. + - w/o pos emb: Tensor (#batch, time, size). + mask (torch.Tensor): Mask tensor for the input (#batch, time). + cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size). + + Returns: + torch.Tensor: Output tensor (#batch, time, size). + torch.Tensor: Mask tensor (#batch, time). + + """ + if isinstance(x_input, tuple): + x, pos_emb = x_input[0], x_input[1] + else: + x, pos_emb = x_input, None + residual = x + x = self.norm_mha(x) + t_leng = x.size(1) + d_dim = x.size(2) + x_new = x.reshape(-1, channel_size, t_leng, d_dim).transpose( + 1, 2 + ) # x_new B*T * C * D + x_k_v = x_new.new(x_new.size(0), x_new.size(1), 5, x_new.size(2), x_new.size(3)) + pad_before = Variable( + torch.zeros(x_new.size(0), 2, x_new.size(2), x_new.size(3)) + ).type(x_new.type()) + pad_after = Variable( + torch.zeros(x_new.size(0), 2, x_new.size(2), x_new.size(3)) + ).type(x_new.type()) + x_pad = torch.cat([pad_before, x_new, pad_after], 1) + x_k_v[:, :, 0, :, :] = x_pad[:, 0:-4, :, :] + x_k_v[:, :, 1, :, :] = x_pad[:, 1:-3, :, :] + x_k_v[:, :, 2, :, :] = x_pad[:, 2:-2, :, :] + x_k_v[:, :, 3, :, :] = x_pad[:, 3:-1, :, :] + x_k_v[:, :, 4, :, :] = x_pad[:, 4:, :, :] + x_new = x_new.reshape(-1, channel_size, d_dim) + x_k_v = x_k_v.reshape(-1, 5 * channel_size, d_dim) + x_att = self.encoder_cros_channel_atten(x_new, x_k_v, x_k_v, None) + x_att = ( + x_att.reshape(-1, t_leng, channel_size, d_dim) + .transpose(1, 2) + .reshape(-1, t_leng, d_dim) + ) + x = residual + self.dropout(x_att) + if pos_emb is not None: + x_input = (x, pos_emb) + else: + x_input = x + x_input, mask = self.encoder_csa(x_input, mask) + + return x_input, mask, channel_size diff --git a/almeval/models/stepaudio/funasr_detach/models/mfcca/mfcca_encoder.py b/almeval/models/stepaudio/funasr_detach/models/mfcca/mfcca_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..8358cd3927d33ee95a0eb822c1fd37e9089c1abd --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/mfcca/mfcca_encoder.py @@ -0,0 +1,436 @@ +from typing import Optional +from typing import Tuple + +import logging +import torch +from torch import nn + + +from funasr_detach.models.encoder.encoder_layer_mfcca import EncoderLayer +from funasr_detach.models.transformer.utils.nets_utils import get_activation +from funasr_detach.models.transformer.utils.nets_utils import make_pad_mask +from funasr_detach.models.transformer.attention import ( + MultiHeadedAttention, # noqa: H301 + RelPositionMultiHeadedAttention, # noqa: H301 + LegacyRelPositionMultiHeadedAttention, # noqa: H301 +) +from funasr_detach.models.transformer.embedding import ( + PositionalEncoding, # noqa: H301 + ScaledPositionalEncoding, # noqa: H301 + RelPositionalEncoding, # noqa: H301 + LegacyRelPositionalEncoding, # noqa: H301 +) +from funasr_detach.models.transformer.layer_norm import LayerNorm +from funasr_detach.models.transformer.utils.multi_layer_conv import Conv1dLinear +from funasr_detach.models.transformer.utils.multi_layer_conv import MultiLayeredConv1d +from funasr_detach.models.transformer.positionwise_feed_forward import ( + PositionwiseFeedForward, # noqa: H301 +) +from funasr_detach.models.transformer.utils.repeat import repeat +from funasr_detach.models.transformer.utils.subsampling import Conv2dSubsampling +from funasr_detach.models.transformer.utils.subsampling import Conv2dSubsampling2 +from funasr_detach.models.transformer.utils.subsampling import Conv2dSubsampling6 +from funasr_detach.models.transformer.utils.subsampling import Conv2dSubsampling8 +from funasr_detach.models.transformer.utils.subsampling import TooShortUttError +from funasr_detach.models.transformer.utils.subsampling import check_short_utt +from funasr_detach.models.encoder.abs_encoder import AbsEncoder +import pdb +import math + + +class ConvolutionModule(nn.Module): + """ConvolutionModule in Conformer model. + Args: + channels (int): The number of channels of conv layers. + kernel_size (int): Kernerl size of conv layers. + """ + + def __init__(self, channels, kernel_size, activation=nn.ReLU(), bias=True): + """Construct an ConvolutionModule object.""" + super(ConvolutionModule, self).__init__() + # kernerl_size should be a odd number for 'SAME' padding + assert (kernel_size - 1) % 2 == 0 + + self.pointwise_conv1 = nn.Conv1d( + channels, + 2 * channels, + kernel_size=1, + stride=1, + padding=0, + bias=bias, + ) + self.depthwise_conv = nn.Conv1d( + channels, + channels, + kernel_size, + stride=1, + padding=(kernel_size - 1) // 2, + groups=channels, + bias=bias, + ) + self.norm = nn.BatchNorm1d(channels) + self.pointwise_conv2 = nn.Conv1d( + channels, + channels, + kernel_size=1, + stride=1, + padding=0, + bias=bias, + ) + self.activation = activation + + def forward(self, x): + """Compute convolution module. + Args: + x (torch.Tensor): Input tensor (#batch, time, channels). + Returns: + torch.Tensor: Output tensor (#batch, time, channels). + """ + # exchange the temporal dimension and the feature dimension + x = x.transpose(1, 2) + + # GLU mechanism + x = self.pointwise_conv1(x) # (batch, 2*channel, dim) + x = nn.functional.glu(x, dim=1) # (batch, channel, dim) + + # 1D Depthwise Conv + x = self.depthwise_conv(x) + x = self.activation(self.norm(x)) + + x = self.pointwise_conv2(x) + + return x.transpose(1, 2) + + +class MFCCAEncoder(AbsEncoder): + """Conformer encoder module. + Args: + input_size (int): Input dimension. + output_size (int): Dimention of attention. + attention_heads (int): The number of heads of multi head attention. + linear_units (int): The number of units of position-wise feed forward. + num_blocks (int): The number of decoder blocks. + dropout_rate (float): Dropout rate. + attention_dropout_rate (float): Dropout rate in attention. + positional_dropout_rate (float): Dropout rate after adding positional encoding. + input_layer (Union[str, torch.nn.Module]): Input layer type. + normalize_before (bool): Whether to use layer_norm before the first block. + concat_after (bool): Whether to concat attention layer's input and output. + If True, additional linear will be applied. + i.e. x -> x + linear(concat(x, att(x))) + If False, no additional linear will be applied. i.e. x -> x + att(x) + positionwise_layer_type (str): "linear", "conv1d", or "conv1d-linear". + positionwise_conv_kernel_size (int): Kernel size of positionwise conv1d layer. + rel_pos_type (str): Whether to use the latest relative positional encoding or + the legacy one. The legacy relative positional encoding will be deprecated + in the future. More Details can be found in + https://github.com/espnet/espnet/pull/2816. + encoder_pos_enc_layer_type (str): Encoder positional encoding layer type. + encoder_attn_layer_type (str): Encoder attention layer type. + activation_type (str): Encoder activation function type. + macaron_style (bool): Whether to use macaron style for positionwise layer. + use_cnn_module (bool): Whether to use convolution module. + zero_triu (bool): Whether to zero the upper triangular part of attention matrix. + cnn_module_kernel (int): Kernerl size of convolution module. + padding_idx (int): Padding idx for input_layer=embed. + """ + + def __init__( + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.0, + input_layer: str = "conv2d", + normalize_before: bool = True, + concat_after: bool = False, + positionwise_layer_type: str = "linear", + positionwise_conv_kernel_size: int = 3, + macaron_style: bool = False, + rel_pos_type: str = "legacy", + pos_enc_layer_type: str = "rel_pos", + selfattention_layer_type: str = "rel_selfattn", + activation_type: str = "swish", + use_cnn_module: bool = True, + zero_triu: bool = False, + cnn_module_kernel: int = 31, + padding_idx: int = -1, + ): + super().__init__() + self._output_size = output_size + + if rel_pos_type == "legacy": + if pos_enc_layer_type == "rel_pos": + pos_enc_layer_type = "legacy_rel_pos" + if selfattention_layer_type == "rel_selfattn": + selfattention_layer_type = "legacy_rel_selfattn" + elif rel_pos_type == "latest": + assert selfattention_layer_type != "legacy_rel_selfattn" + assert pos_enc_layer_type != "legacy_rel_pos" + else: + raise ValueError("unknown rel_pos_type: " + rel_pos_type) + + activation = get_activation(activation_type) + if pos_enc_layer_type == "abs_pos": + pos_enc_class = PositionalEncoding + elif pos_enc_layer_type == "scaled_abs_pos": + pos_enc_class = ScaledPositionalEncoding + elif pos_enc_layer_type == "rel_pos": + assert selfattention_layer_type == "rel_selfattn" + pos_enc_class = RelPositionalEncoding + elif pos_enc_layer_type == "legacy_rel_pos": + assert selfattention_layer_type == "legacy_rel_selfattn" + pos_enc_class = LegacyRelPositionalEncoding + logging.warning( + "Using legacy_rel_pos and it will be deprecated in the future." + ) + else: + raise ValueError("unknown pos_enc_layer: " + pos_enc_layer_type) + + if input_layer == "linear": + self.embed = torch.nn.Sequential( + torch.nn.Linear(input_size, output_size), + torch.nn.LayerNorm(output_size), + torch.nn.Dropout(dropout_rate), + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer == "conv2d": + self.embed = Conv2dSubsampling( + input_size, + output_size, + dropout_rate, + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer == "conv2d6": + self.embed = Conv2dSubsampling6( + input_size, + output_size, + dropout_rate, + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer == "conv2d8": + self.embed = Conv2dSubsampling8( + input_size, + output_size, + dropout_rate, + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer == "embed": + self.embed = torch.nn.Sequential( + torch.nn.Embedding(input_size, output_size, padding_idx=padding_idx), + pos_enc_class(output_size, positional_dropout_rate), + ) + elif isinstance(input_layer, torch.nn.Module): + self.embed = torch.nn.Sequential( + input_layer, + pos_enc_class(output_size, positional_dropout_rate), + ) + elif input_layer is None: + self.embed = torch.nn.Sequential( + pos_enc_class(output_size, positional_dropout_rate) + ) + else: + raise ValueError("unknown input_layer: " + input_layer) + self.normalize_before = normalize_before + if positionwise_layer_type == "linear": + positionwise_layer = PositionwiseFeedForward + positionwise_layer_args = ( + output_size, + linear_units, + dropout_rate, + activation, + ) + elif positionwise_layer_type == "conv1d": + positionwise_layer = MultiLayeredConv1d + positionwise_layer_args = ( + output_size, + linear_units, + positionwise_conv_kernel_size, + dropout_rate, + ) + elif positionwise_layer_type == "conv1d-linear": + positionwise_layer = Conv1dLinear + positionwise_layer_args = ( + output_size, + linear_units, + positionwise_conv_kernel_size, + dropout_rate, + ) + else: + raise NotImplementedError("Support only linear or conv1d.") + + if selfattention_layer_type == "selfattn": + encoder_selfattn_layer = MultiHeadedAttention + encoder_selfattn_layer_args = ( + attention_heads, + output_size, + attention_dropout_rate, + ) + elif selfattention_layer_type == "legacy_rel_selfattn": + assert pos_enc_layer_type == "legacy_rel_pos" + encoder_selfattn_layer = LegacyRelPositionMultiHeadedAttention + encoder_selfattn_layer_args = ( + attention_heads, + output_size, + attention_dropout_rate, + ) + logging.warning( + "Using legacy_rel_selfattn and it will be deprecated in the future." + ) + elif selfattention_layer_type == "rel_selfattn": + assert pos_enc_layer_type == "rel_pos" + encoder_selfattn_layer = RelPositionMultiHeadedAttention + encoder_selfattn_layer_args = ( + attention_heads, + output_size, + attention_dropout_rate, + zero_triu, + ) + else: + raise ValueError("unknown encoder_attn_layer: " + selfattention_layer_type) + + convolution_layer = ConvolutionModule + convolution_layer_args = (output_size, cnn_module_kernel, activation) + encoder_selfattn_layer_raw = MultiHeadedAttention + encoder_selfattn_layer_args_raw = ( + attention_heads, + output_size, + attention_dropout_rate, + ) + self.encoders = repeat( + num_blocks, + lambda lnum: EncoderLayer( + output_size, + encoder_selfattn_layer_raw(*encoder_selfattn_layer_args_raw), + encoder_selfattn_layer(*encoder_selfattn_layer_args), + positionwise_layer(*positionwise_layer_args), + positionwise_layer(*positionwise_layer_args) if macaron_style else None, + convolution_layer(*convolution_layer_args) if use_cnn_module else None, + dropout_rate, + normalize_before, + concat_after, + ), + ) + if self.normalize_before: + self.after_norm = LayerNorm(output_size) + self.conv1 = torch.nn.Conv2d(8, 16, [5, 7], stride=[1, 1], padding=(2, 3)) + + self.conv2 = torch.nn.Conv2d(16, 32, [5, 7], stride=[1, 1], padding=(2, 3)) + + self.conv3 = torch.nn.Conv2d(32, 16, [5, 7], stride=[1, 1], padding=(2, 3)) + + self.conv4 = torch.nn.Conv2d(16, 1, [5, 7], stride=[1, 1], padding=(2, 3)) + + def output_size(self) -> int: + return self._output_size + + def forward( + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + channel_size: torch.Tensor, + prev_states: torch.Tensor = None, + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Calculate forward propagation. + Args: + xs_pad (torch.Tensor): Input tensor (#batch, L, input_size). + ilens (torch.Tensor): Input length (#batch). + prev_states (torch.Tensor): Not to be used now. + Returns: + torch.Tensor: Output tensor (#batch, L, output_size). + torch.Tensor: Output length (#batch). + torch.Tensor: Not to be used now. + """ + masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device) + if ( + isinstance(self.embed, Conv2dSubsampling) + or isinstance(self.embed, Conv2dSubsampling6) + or isinstance(self.embed, Conv2dSubsampling8) + ): + short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1)) + if short_status: + raise TooShortUttError( + f"has {xs_pad.size(1)} frames and is too short for subsampling " + + f"(it needs more than {limit_size} frames), return empty results", + xs_pad.size(1), + limit_size, + ) + xs_pad, masks = self.embed(xs_pad, masks) + else: + xs_pad = self.embed(xs_pad) + xs_pad, masks, channel_size = self.encoders(xs_pad, masks, channel_size) + if isinstance(xs_pad, tuple): + xs_pad = xs_pad[0] + + t_leng = xs_pad.size(1) + d_dim = xs_pad.size(2) + xs_pad = xs_pad.reshape(-1, channel_size, t_leng, d_dim) + # pdb.set_trace() + if channel_size < 8: + repeat_num = math.ceil(8 / channel_size) + xs_pad = xs_pad.repeat(1, repeat_num, 1, 1)[:, 0:8, :, :] + xs_pad = self.conv1(xs_pad) + xs_pad = self.conv2(xs_pad) + xs_pad = self.conv3(xs_pad) + xs_pad = self.conv4(xs_pad) + xs_pad = xs_pad.squeeze().reshape(-1, t_leng, d_dim) + mask_tmp = masks.size(1) + masks = masks.reshape(-1, channel_size, mask_tmp, t_leng)[:, 0, :, :] + + if self.normalize_before: + xs_pad = self.after_norm(xs_pad) + + olens = masks.squeeze(1).sum(1) + return xs_pad, olens, None + + def forward_hidden( + self, + xs_pad: torch.Tensor, + ilens: torch.Tensor, + prev_states: torch.Tensor = None, + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Calculate forward propagation. + Args: + xs_pad (torch.Tensor): Input tensor (#batch, L, input_size). + ilens (torch.Tensor): Input length (#batch). + prev_states (torch.Tensor): Not to be used now. + Returns: + torch.Tensor: Output tensor (#batch, L, output_size). + torch.Tensor: Output length (#batch). + torch.Tensor: Not to be used now. + """ + masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device) + if ( + isinstance(self.embed, Conv2dSubsampling) + or isinstance(self.embed, Conv2dSubsampling6) + or isinstance(self.embed, Conv2dSubsampling8) + ): + short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1)) + if short_status: + raise TooShortUttError( + f"has {xs_pad.size(1)} frames and is too short for subsampling " + + f"(it needs more than {limit_size} frames), return empty results", + xs_pad.size(1), + limit_size, + ) + xs_pad, masks = self.embed(xs_pad, masks) + else: + xs_pad = self.embed(xs_pad) + num_layer = len(self.encoders) + for idx, encoder in enumerate(self.encoders): + xs_pad, masks = encoder(xs_pad, masks) + if idx == num_layer // 2 - 1: + hidden_feature = xs_pad + if isinstance(xs_pad, tuple): + xs_pad = xs_pad[0] + hidden_feature = hidden_feature[0] + if self.normalize_before: + xs_pad = self.after_norm(xs_pad) + self.hidden_feature = self.after_norm(hidden_feature) + + olens = masks.squeeze(1).sum(1) + return xs_pad, olens, None diff --git a/almeval/models/stepaudio/funasr_detach/models/monotonic_aligner/__init__.py b/almeval/models/stepaudio/funasr_detach/models/monotonic_aligner/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/almeval/models/stepaudio/funasr_detach/models/monotonic_aligner/model.py b/almeval/models/stepaudio/funasr_detach/models/monotonic_aligner/model.py new file mode 100644 index 0000000000000000000000000000000000000000..6c4ed9f024db40f97b1dda9e07456db3f9a490f3 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/monotonic_aligner/model.py @@ -0,0 +1,238 @@ +#!/usr/bin/env python3 +# -*- encoding: utf-8 -*- +# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved. +# MIT License (https://opensource.org/licenses/MIT) + +import time +import copy +import torch +from torch.cuda.amp import autocast +from typing import Union, Dict, List, Tuple, Optional + +from funasr_detach.register import tables +from funasr_detach.models.ctc.ctc import CTC +from funasr_detach.utils import postprocess_utils +from funasr_detach.utils.datadir_writer import DatadirWriter +from funasr_detach.models.paraformer.cif_predictor import mae_loss +from funasr_detach.train_utils.device_funcs import force_gatherable +from funasr_detach.models.transformer.utils.add_sos_eos import add_sos_eos +from funasr_detach.models.transformer.utils.nets_utils import make_pad_mask +from funasr_detach.utils.timestamp_tools import ts_prediction_lfr6_standard +from funasr_detach.utils.load_utils import load_audio_text_image_video, extract_fbank + + +@tables.register("model_classes", "MonotonicAligner") +class MonotonicAligner(torch.nn.Module): + """ + Author: Speech Lab of DAMO Academy, Alibaba Group + Achieving timestamp prediction while recognizing with non-autoregressive end-to-end ASR model + https://arxiv.org/abs/2301.12343 + """ + + def __init__( + self, + input_size: int = 80, + specaug: Optional[str] = None, + specaug_conf: Optional[Dict] = None, + normalize: str = None, + normalize_conf: Optional[Dict] = None, + encoder: str = None, + encoder_conf: Optional[Dict] = None, + predictor: str = None, + predictor_conf: Optional[Dict] = None, + predictor_bias: int = 0, + length_normalized_loss: bool = False, + **kwargs, + ): + super().__init__() + + if specaug is not None: + specaug_class = tables.specaug_classes.get(specaug) + specaug = specaug_class(**specaug_conf) + if normalize is not None: + normalize_class = tables.normalize_classes.get(normalize) + normalize = normalize_class(**normalize_conf) + encoder_class = tables.encoder_classes.get(encoder) + encoder = encoder_class(input_size=input_size, **encoder_conf) + encoder_output_size = encoder.output_size() + predictor_class = tables.predictor_classes.get(predictor) + predictor = predictor_class(**predictor_conf) + self.specaug = specaug + self.normalize = normalize + self.encoder = encoder + self.predictor = predictor + self.criterion_pre = mae_loss(normalize_length=length_normalized_loss) + self.predictor_bias = predictor_bias + + def forward( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + text: torch.Tensor, + text_lengths: torch.Tensor, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: + """Frontend + Encoder + Decoder + Calc loss + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + text: (Batch, Length) + text_lengths: (Batch,) + """ + assert text_lengths.dim() == 1, text_lengths.shape + # Check that batch_size is unified + assert ( + speech.shape[0] + == speech_lengths.shape[0] + == text.shape[0] + == text_lengths.shape[0] + ), (speech.shape, speech_lengths.shape, text.shape, text_lengths.shape) + batch_size = speech.shape[0] + # for data-parallel + text = text[:, : text_lengths.max()] + speech = speech[:, : speech_lengths.max()] + + # 1. Encoder + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) + if self.predictor_bias == 1: + _, text = add_sos_eos(text, 1, 2, -1) + text_lengths = text_lengths + self.predictor_bias + _, _, _, _, pre_token_length2 = self.predictor( + encoder_out, text, encoder_out_mask, ignore_id=-1 + ) + + # loss_pre = self.criterion_pre(ys_pad_lens.type_as(pre_token_length), pre_token_length) + loss_pre = self.criterion_pre( + text_lengths.type_as(pre_token_length2), pre_token_length2 + ) + + loss = loss_pre + stats = dict() + + # Collect Attn branch stats + stats["loss_pre"] = loss_pre.detach().cpu() if loss_pre is not None else None + stats["loss"] = torch.clone(loss.detach()) + + # force_gatherable: to-device and to-tensor if scalar for DataParallel + loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) + return loss, stats, weight + + def calc_predictor_timestamp(self, encoder_out, encoder_out_lens, token_num): + encoder_out_mask = ( + ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :] + ).to(encoder_out.device) + ds_alphas, ds_cif_peak, us_alphas, us_peaks = ( + self.predictor.get_upsample_timestamp( + encoder_out, encoder_out_mask, token_num + ) + ) + return ds_alphas, ds_cif_peak, us_alphas, us_peaks + + def encode( + self, + speech: torch.Tensor, + speech_lengths: torch.Tensor, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Encoder. Note that this method is used by asr_inference.py + Args: + speech: (Batch, Length, ...) + speech_lengths: (Batch, ) + ind: int + """ + with autocast(False): + + # Data augmentation + if self.specaug is not None and self.training: + speech, speech_lengths = self.specaug(speech, speech_lengths) + + # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN + if self.normalize is not None: + speech, speech_lengths = self.normalize(speech, speech_lengths) + + # Forward encoder + encoder_out, encoder_out_lens, _ = self.encoder(speech, speech_lengths) + if isinstance(encoder_out, tuple): + encoder_out = encoder_out[0] + + return encoder_out, encoder_out_lens + + def inference( + self, + data_in, + data_lengths=None, + key: list = None, + tokenizer=None, + frontend=None, + **kwargs, + ): + meta_data = {} + # extract fbank feats + time1 = time.perf_counter() + audio_list, text_token_int_list = load_audio_text_image_video( + data_in, + fs=frontend.fs, + audio_fs=kwargs.get("fs", 16000), + data_type=kwargs.get("data_type", "sound"), + tokenizer=tokenizer, + ) + time2 = time.perf_counter() + meta_data["load_data"] = f"{time2 - time1:0.3f}" + speech, speech_lengths = extract_fbank( + audio_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend + ) + time3 = time.perf_counter() + meta_data["extract_feat"] = f"{time3 - time2:0.3f}" + meta_data["batch_data_time"] = ( + speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000 + ) + + speech = speech.to(device=kwargs["device"]) + speech_lengths = speech_lengths.to(device=kwargs["device"]) + + # Encoder + encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) + + # predictor + text_lengths = torch.tensor([len(i) + 1 for i in text_token_int_list]).to( + encoder_out.device + ) + _, _, us_alphas, us_peaks = self.calc_predictor_timestamp( + encoder_out, encoder_out_lens, token_num=text_lengths + ) + + results = [] + ibest_writer = None + if kwargs.get("output_dir") is not None: + if not hasattr(self, "writer"): + self.writer = DatadirWriter(kwargs.get("output_dir")) + ibest_writer = self.writer["tp_res"] + + for i, (us_alpha, us_peak, token_int) in enumerate( + zip(us_alphas, us_peaks, text_token_int_list) + ): + token = tokenizer.ids2tokens(token_int) + timestamp_str, timestamp = ts_prediction_lfr6_standard( + us_alpha[: encoder_out_lens[i] * 3], + us_peak[: encoder_out_lens[i] * 3], + copy.copy(token), + ) + text_postprocessed, time_stamp_postprocessed, _ = ( + postprocess_utils.sentence_postprocess(token, timestamp) + ) + result_i = { + "key": key[i], + "text": text_postprocessed, + "timestamp": time_stamp_postprocessed, + } + results.append(result_i) + + if ibest_writer: + # ibest_writer["token"][key[i]] = " ".join(token) + ibest_writer["timestamp_list"][key[i]] = time_stamp_postprocessed + ibest_writer["timestamp_str"][key[i]] = timestamp_str + + return results, meta_data diff --git a/almeval/models/stepaudio/funasr_detach/models/monotonic_aligner/template.yaml b/almeval/models/stepaudio/funasr_detach/models/monotonic_aligner/template.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b1379de787be8ccefb848244b7186bf55cb6856b --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/monotonic_aligner/template.yaml @@ -0,0 +1,115 @@ +# This is an example that demonstrates how to configure a model file. +# You can modify the configuration according to your own requirements. + +# to print the register_table: +# from funasr.register import tables +# tables.print() + +# network architecture +model: MonotonicAligner +model_conf: + length_normalized_loss: False + predictor_bias: 1 + +# encoder +encoder: SANMEncoder +encoder_conf: + output_size: 320 + attention_heads: 4 + linear_units: 1280 + num_blocks: 30 + dropout_rate: 0.1 + positional_dropout_rate: 0.1 + attention_dropout_rate: 0.1 + input_layer: pe + pos_enc_class: SinusoidalPositionEncoder + normalize_before: true + kernel_size: 11 + sanm_shfit: 0 + selfattention_layer_type: sanm + +predictor: CifPredictorV3 +predictor_conf: + idim: 320 + threshold: 1.0 + l_order: 1 + r_order: 1 + tail_threshold: 0.45 + smooth_factor2: 0.25 + noise_threshold2: 0.01 + upsample_times: 3 + use_cif1_cnn: false + upsample_type: cnn_blstm + +# frontend related +frontend: WavFrontend +frontend_conf: + fs: 16000 + window: hamming + n_mels: 80 + frame_length: 25 + frame_shift: 10 + lfr_m: 7 + lfr_n: 6 + +specaug: SpecAugLFR +specaug_conf: + apply_time_warp: false + time_warp_window: 5 + time_warp_mode: bicubic + apply_freq_mask: true + freq_mask_width_range: + - 0 + - 30 + lfr_rate: 6 + num_freq_mask: 1 + apply_time_mask: true + time_mask_width_range: + - 0 + - 12 + num_time_mask: 1 + +train_conf: + accum_grad: 1 + grad_clip: 5 + max_epoch: 150 + val_scheduler_criterion: + - valid + - acc + best_model_criterion: + - - valid + - acc + - max + keep_nbest_models: 10 + log_interval: 50 + +optim: adam +optim_conf: + lr: 0.0005 +scheduler: warmuplr +scheduler_conf: + warmup_steps: 30000 + +dataset: AudioDataset +dataset_conf: + index_ds: IndexDSJsonl + batch_sampler: DynamicBatchLocalShuffleSampler + batch_type: example # example or length + batch_size: 1 # if batch_type is example, batch_size is the numbers of samples; if length, batch_size is source_token_len+target_token_len; + max_token_length: 2048 # filter samples if source_token_len+target_token_len > max_token_length, + buffer_size: 500 + shuffle: True + num_workers: 0 + +tokenizer: CharTokenizer +tokenizer_conf: + unk_symbol: + split_with_space: true + +ctc_conf: + dropout_rate: 0.0 + ctc_type: builtin + reduce: true + ignore_nan_grad: true + +normalize: null diff --git a/almeval/models/stepaudio/funasr_detach/models/mossformer/mossformer.py b/almeval/models/stepaudio/funasr_detach/models/mossformer/mossformer.py new file mode 100644 index 0000000000000000000000000000000000000000..be352895c478406224b4706e509752cfa04b3e52 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/models/mossformer/mossformer.py @@ -0,0 +1,324 @@ +import torch +import torch.nn.functional as F +from torch import nn, einsum +from einops import rearrange + + +def identity(t, *args, **kwargs): + return t + + +def append_dims(x, num_dims): + if num_dims <= 0: + return x + return x.view(*x.shape, *((1,) * num_dims)) + + +def exists(val): + return val is not None + + +def default(val, d): + return val if exists(val) else d + + +def padding_to_multiple_of(n, mult): + remainder = n % mult + if remainder == 0: + return 0 + return mult - remainder + + +class Transpose(nn.Module): + """Wrapper class of torch.transpose() for Sequential module.""" + + def __init__(self, shape: tuple): + super(Transpose, self).__init__() + self.shape = shape + + def forward(self, x): + return x.transpose(*self.shape) + + +class DepthwiseConv1d(nn.Module): + """ + When groups == in_channels and out_channels == K * in_channels, where K is a positive integer, + this operation is termed in literature as depthwise convolution. + Args: + in_channels (int): Number of channels in the input + out_channels (int): Number of channels produced by the convolution + kernel_size (int or tuple): Size of the convolving kernel + stride (int, optional): Stride of the convolution. Default: 1 + padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0 + bias (bool, optional): If True, adds a learnable bias to the output. Default: True + Inputs: inputs + - **inputs** (batch, in_channels, time): Tensor containing input vector + Returns: outputs + - **outputs** (batch, out_channels, time): Tensor produces by depthwise 1-D convolution. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + padding: int = 0, + bias: bool = False, + ) -> None: + super(DepthwiseConv1d, self).__init__() + assert ( + out_channels % in_channels == 0 + ), "out_channels should be constant multiple of in_channels" + self.conv = nn.Conv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + groups=in_channels, + stride=stride, + padding=padding, + bias=bias, + ) + + def forward(self, inputs): + return self.conv(inputs) + + +class ConvModule(nn.Module): + """ + Conformer convolution module starts with a pointwise convolution and a gated linear unit (GLU). + This is followed by a single 1-D depthwise convolution layer. Batchnorm is deployed just after the convolution + to aid training deep models. + Args: + in_channels (int): Number of channels in the input + kernel_size (int or tuple, optional): Size of the convolving kernel Default: 31 + dropout_p (float, optional): probability of dropout + Inputs: inputs + inputs (batch, time, dim): Tensor contains input sequences + Outputs: outputs + outputs (batch, time, dim): Tensor produces by conformer convolution module. + """ + + def __init__( + self, + in_channels: int, + kernel_size: int = 17, + expansion_factor: int = 2, + dropout_p: float = 0.1, + ) -> None: + super(ConvModule, self).__init__() + assert ( + kernel_size - 1 + ) % 2 == 0, "kernel_size should be a odd number for 'SAME' padding" + assert expansion_factor == 2, "Currently, Only Supports expansion_factor 2" + + self.sequential = nn.Sequential( + Transpose(shape=(1, 2)), + DepthwiseConv1d( + in_channels, + in_channels, + kernel_size, + stride=1, + padding=(kernel_size - 1) // 2, + ), + ) + + def forward(self, inputs): + return inputs + self.sequential(inputs).transpose(1, 2) + + +class OffsetScale(nn.Module): + def __init__(self, dim, heads=1): + super().__init__() + self.gamma = nn.Parameter(torch.ones(heads, dim)) + self.beta = nn.Parameter(torch.zeros(heads, dim)) + nn.init.normal_(self.gamma, std=0.02) + + def forward(self, x): + out = einsum("... d, h d -> ... h d", x, self.gamma) + self.beta + return out.unbind(dim=-2) + + +class FFConvM(nn.Module): + def __init__(self, dim_in, dim_out, norm_klass=nn.LayerNorm, dropout=0.1): + super().__init__() + self.mdl = nn.Sequential( + norm_klass(dim_in), + nn.Linear(dim_in, dim_out), + nn.SiLU(), + ConvModule(dim_out), + nn.Dropout(dropout), + ) + + def forward( + self, + x, + ): + output = self.mdl(x) + return output + + +class FLASH_ShareA_FFConvM(nn.Module): + def __init__( + self, + *, + dim, + group_size=256, + query_key_dim=128, + expansion_factor=1.0, + causal=False, + dropout=0.1, + rotary_pos_emb=None, + norm_klass=nn.LayerNorm, + shift_tokens=True + ): + super().__init__() + hidden_dim = int(dim * expansion_factor) + self.group_size = group_size + self.causal = causal + self.shift_tokens = shift_tokens + + # positional embeddings + self.rotary_pos_emb = rotary_pos_emb + # norm + self.dropout = nn.Dropout(dropout) + # projections + + self.to_hidden = FFConvM( + dim_in=dim, + dim_out=hidden_dim, + norm_klass=norm_klass, + dropout=dropout, + ) + self.to_qk = FFConvM( + dim_in=dim, + dim_out=query_key_dim, + norm_klass=norm_klass, + dropout=dropout, + ) + + self.qk_offset_scale = OffsetScale(query_key_dim, heads=4) + + self.to_out = FFConvM( + dim_in=dim * 2, + dim_out=dim, + norm_klass=norm_klass, + dropout=dropout, + ) + + self.gateActivate = nn.Sigmoid() + + def forward(self, x, *, mask=None): + """ + b - batch + n - sequence length (within groups) + g - group dimension + d - feature dimension (keys) + e - feature dimension (values) + i - sequence dimension (source) + j - sequence dimension (target) + """ + + normed_x = x + + # do token shift - a great, costless trick from an independent AI researcher in Shenzhen + residual = x + + if self.shift_tokens: + x_shift, x_pass = normed_x.chunk(2, dim=-1) + x_shift = F.pad(x_shift, (0, 0, 1, -1), value=0.0) + normed_x = torch.cat((x_shift, x_pass), dim=-1) + + # initial projections + + v, u = self.to_hidden(normed_x).chunk(2, dim=-1) + qk = self.to_qk(normed_x) + + # offset and scale + quad_q, lin_q, quad_k, lin_k = self.qk_offset_scale(qk) + att_v, att_u = self.cal_attention(x, quad_q, lin_q, quad_k, lin_k, v, u) + out = (att_u * v) * self.gateActivate(att_v * u) + x = x + self.to_out(out) + return x + + def cal_attention(self, x, quad_q, lin_q, quad_k, lin_k, v, u, mask=None): + b, n, device, g = x.shape[0], x.shape[-2], x.device, self.group_size + + if exists(mask): + lin_mask = rearrange(mask, "... -> ... 1") + lin_k = lin_k.masked_fill(~lin_mask, 0.0) + + # rotate queries and keys + + if exists(self.rotary_pos_emb): + quad_q, lin_q, quad_k, lin_k = map( + self.rotary_pos_emb.rotate_queries_or_keys, + (quad_q, lin_q, quad_k, lin_k), + ) + + # padding for groups + + padding = padding_to_multiple_of(n, g) + + if padding > 0: + quad_q, quad_k, lin_q, lin_k, v, u = map( + lambda t: F.pad(t, (0, 0, 0, padding), value=0.0), + (quad_q, quad_k, lin_q, lin_k, v, u), + ) + + mask = default(mask, torch.ones((b, n), device=device, dtype=torch.bool)) + mask = F.pad(mask, (0, padding), value=False) + + # group along sequence + + quad_q, quad_k, lin_q, lin_k, v, u = map( + lambda t: rearrange(t, "b (g n) d -> b g n d", n=self.group_size), + (quad_q, quad_k, lin_q, lin_k, v, u), + ) + + if exists(mask): + mask = rearrange(mask, "b (g j) -> b g 1 j", j=g) + + # calculate quadratic attention output + + sim = einsum("... i d, ... j d -> ... i j", quad_q, quad_k) / g + + attn = F.relu(sim) ** 2 + attn = self.dropout(attn) + + if exists(mask): + attn = attn.masked_fill(~mask, 0.0) + + if self.causal: + causal_mask = torch.ones((g, g), dtype=torch.bool, device=device).triu(1) + attn = attn.masked_fill(causal_mask, 0.0) + + quad_out_v = einsum("... i j, ... j d -> ... i d", attn, v) + quad_out_u = einsum("... i j, ... j d -> ... i d", attn, u) + + # calculate linear attention output + + if self.causal: + lin_kv = einsum("b g n d, b g n e -> b g d e", lin_k, v) / g + # exclusive cumulative sum along group dimension + lin_kv = lin_kv.cumsum(dim=1) + lin_kv = F.pad(lin_kv, (0, 0, 0, 0, 1, -1), value=0.0) + lin_out_v = einsum("b g d e, b g n d -> b g n e", lin_kv, lin_q) + + lin_ku = einsum("b g n d, b g n e -> b g d e", lin_k, u) / g + # exclusive cumulative sum along group dimension + lin_ku = lin_ku.cumsum(dim=1) + lin_ku = F.pad(lin_ku, (0, 0, 0, 0, 1, -1), value=0.0) + lin_out_u = einsum("b g d e, b g n d -> b g n e", lin_ku, lin_q) + else: + lin_kv = einsum("b g n d, b g n e -> b d e", lin_k, v) / n + lin_out_v = einsum("b g n d, b d e -> b g n e", lin_q, lin_kv) + + lin_ku = einsum("b g n d, b g n e -> b d e", lin_k, u) / n + lin_out_u = einsum("b g n d, b d e -> b g n e", lin_q, lin_ku) + + # fold back groups into full sequence, and excise out padding + return map( + lambda t: rearrange(t, "b g n d -> b (g n) d")[:, :n], + (quad_out_v + lin_out_v, quad_out_u + lin_out_u), + ) diff --git a/almeval/models/stepaudio/funasr_detach/tokenizer/__init__.py b/almeval/models/stepaudio/funasr_detach/tokenizer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/almeval/models/stepaudio/funasr_detach/tokenizer/abs_tokenizer.py b/almeval/models/stepaudio/funasr_detach/tokenizer/abs_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..69cce662807443d8bcbacbafb505b828f6012029 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/tokenizer/abs_tokenizer.py @@ -0,0 +1,95 @@ +import json +import numpy as np +from abc import ABC +from pathlib import Path +from abc import abstractmethod +from typing import Union, Iterable, List, Dict + + +class AbsTokenizer(ABC): + @abstractmethod + def text2tokens(self, line: str) -> List[str]: + raise NotImplementedError + + @abstractmethod + def tokens2text(self, tokens: Iterable[str]) -> str: + raise NotImplementedError + + +class BaseTokenizer(ABC): + def __init__( + self, + token_list: Union[Path, str, Iterable[str]] = None, + unk_symbol: str = "", + **kwargs, + ): + + if token_list is not None: + if isinstance(token_list, (Path, str)) and token_list.endswith(".txt"): + token_list = Path(token_list) + self.token_list_repr = str(token_list) + self.token_list: List[str] = [] + + with token_list.open("r", encoding="utf-8") as f: + for idx, line in enumerate(f): + line = line.rstrip() + self.token_list.append(line) + elif isinstance(token_list, (Path, str)) and token_list.endswith(".json"): + token_list = Path(token_list) + self.token_list_repr = str(token_list) + self.token_list: List[str] = [] + + with open(token_list, "r", encoding="utf-8") as f: + self.token_list = json.load(f) + + else: + self.token_list: List[str] = list(token_list) + self.token_list_repr = "" + for i, t in enumerate(self.token_list): + if i == 3: + break + self.token_list_repr += f"{t}, " + self.token_list_repr += f"... (NVocab={(len(self.token_list))})" + + self.token2id: Dict[str, int] = {} + for i, t in enumerate(self.token_list): + if t in self.token2id: + raise RuntimeError(f'Symbol "{t}" is duplicated') + self.token2id[t] = i + + self.unk_symbol = unk_symbol + if self.unk_symbol not in self.token2id: + raise RuntimeError( + f"Unknown symbol '{unk_symbol}' doesn't exist in the token_list" + ) + self.unk_id = self.token2id[self.unk_symbol] + + def encode(self, text): + tokens = self.text2tokens(text) + text_ints = self.tokens2ids(tokens) + + return text_ints + + def decode(self, text_ints): + token = self.ids2tokens(text_ints) + text = self.tokens2text(token) + return text + + def get_num_vocabulary_size(self) -> int: + return len(self.token_list) + + def ids2tokens(self, integers: Union[np.ndarray, Iterable[int]]) -> List[str]: + if isinstance(integers, np.ndarray) and integers.ndim != 1: + raise ValueError(f"Must be 1 dim ndarray, but got {integers.ndim}") + return [self.token_list[i] for i in integers] + + def tokens2ids(self, tokens: Iterable[str]) -> List[int]: + return [self.token2id.get(i, self.unk_id) for i in tokens] + + @abstractmethod + def text2tokens(self, line: str) -> List[str]: + raise NotImplementedError + + @abstractmethod + def tokens2text(self, tokens: Iterable[str]) -> str: + raise NotImplementedError diff --git a/almeval/models/stepaudio/funasr_detach/tokenizer/build_tokenizer.py b/almeval/models/stepaudio/funasr_detach/tokenizer/build_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..a4e717f62f79516fe714f7ab97a1f15c9f640edf --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/tokenizer/build_tokenizer.py @@ -0,0 +1,61 @@ +from pathlib import Path +from typing import Iterable +from typing import Union + + +from funasr_detach.tokenizer.abs_tokenizer import AbsTokenizer +from funasr_detach.tokenizer.char_tokenizer import CharTokenizer +from funasr_detach.tokenizer.phoneme_tokenizer import PhonemeTokenizer +from funasr_detach.tokenizer.sentencepiece_tokenizer import SentencepiecesTokenizer +from funasr_detach.tokenizer.word_tokenizer import WordTokenizer + + +def build_tokenizer( + token_type: str, + bpemodel: Union[Path, str, Iterable[str]] = None, + non_linguistic_symbols: Union[Path, str, Iterable[str]] = None, + remove_non_linguistic_symbols: bool = False, + space_symbol: str = "", + delimiter: str = None, + g2p_type: str = None, +) -> AbsTokenizer: + """A helper function to instantiate Tokenizer""" + if token_type == "bpe": + if bpemodel is None: + raise ValueError('bpemodel is required if token_type = "bpe"') + + if remove_non_linguistic_symbols: + raise RuntimeError( + "remove_non_linguistic_symbols is not implemented for token_type=bpe" + ) + return SentencepiecesTokenizer(bpemodel) + + elif token_type == "word": + if remove_non_linguistic_symbols and non_linguistic_symbols is not None: + return WordTokenizer( + delimiter=delimiter, + non_linguistic_symbols=non_linguistic_symbols, + remove_non_linguistic_symbols=True, + ) + else: + return WordTokenizer(delimiter=delimiter) + + elif token_type == "char": + return CharTokenizer( + non_linguistic_symbols=non_linguistic_symbols, + space_symbol=space_symbol, + remove_non_linguistic_symbols=remove_non_linguistic_symbols, + ) + + elif token_type == "phn": + return PhonemeTokenizer( + g2p_type=g2p_type, + non_linguistic_symbols=non_linguistic_symbols, + space_symbol=space_symbol, + remove_non_linguistic_symbols=remove_non_linguistic_symbols, + ) + + else: + raise ValueError( + f"token_mode must be one of bpe, word, char or phn: " f"{token_type}" + ) diff --git a/almeval/models/stepaudio/funasr_detach/tokenizer/char_tokenizer.py b/almeval/models/stepaudio/funasr_detach/tokenizer/char_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..916429080fa2964ad93160793c2d4427092a9ad8 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/tokenizer/char_tokenizer.py @@ -0,0 +1,111 @@ +from pathlib import Path +from typing import Iterable +from typing import List +from typing import Union +import warnings +import re + +from funasr_detach.tokenizer.abs_tokenizer import BaseTokenizer +from funasr_detach.register import tables + + +@tables.register("tokenizer_classes", "CharTokenizer") +class CharTokenizer(BaseTokenizer): + def __init__( + self, + non_linguistic_symbols: Union[Path, str, Iterable[str]] = None, + space_symbol: str = "", + remove_non_linguistic_symbols: bool = False, + split_with_space: bool = False, + seg_dict: str = None, + **kwargs, + ): + super().__init__(**kwargs) + self.space_symbol = space_symbol + if non_linguistic_symbols is None: + self.non_linguistic_symbols = set() + elif isinstance(non_linguistic_symbols, (Path, str)): + non_linguistic_symbols = Path(non_linguistic_symbols) + try: + with non_linguistic_symbols.open("r", encoding="utf-8") as f: + self.non_linguistic_symbols = set(line.rstrip() for line in f) + except FileNotFoundError: + warnings.warn(f"{non_linguistic_symbols} doesn't exist.") + self.non_linguistic_symbols = set() + else: + self.non_linguistic_symbols = set(non_linguistic_symbols) + self.remove_non_linguistic_symbols = remove_non_linguistic_symbols + self.split_with_space = split_with_space + self.seg_dict = None + if seg_dict is not None: + self.seg_dict = load_seg_dict(seg_dict) + + def __repr__(self): + return ( + f"{self.__class__.__name__}(" + f'space_symbol="{self.space_symbol}"' + f'non_linguistic_symbols="{self.non_linguistic_symbols}"' + f")" + ) + + def text2tokens(self, line: Union[str, list]) -> List[str]: + + # if self.split_with_space: + + if self.seg_dict is not None: + tokens = line.strip().split(" ") + tokens = seg_tokenize(tokens, self.seg_dict) + else: + tokens = [] + while len(line) != 0: + for w in self.non_linguistic_symbols: + if line.startswith(w): + if not self.remove_non_linguistic_symbols: + tokens.append(line[: len(w)]) + line = line[len(w) :] + break + else: + t = line[0] + if t == " ": + # t = "" + line = line[1:] + continue + tokens.append(t) + line = line[1:] + return tokens + + def tokens2text(self, tokens: Iterable[str]) -> str: + tokens = [t if t != self.space_symbol else " " for t in tokens] + return "".join(tokens) + + +def load_seg_dict(seg_dict_file): + seg_dict = {} + assert isinstance(seg_dict_file, str) + with open(seg_dict_file, "r", encoding="utf8") as f: + lines = f.readlines() + for line in lines: + s = line.strip().split() + key = s[0] + value = s[1:] + seg_dict[key] = " ".join(value) + return seg_dict + + +def seg_tokenize(txt, seg_dict): + pattern = re.compile(r"^[\u4E00-\u9FA50-9]+$") + out_txt = "" + for word in txt: + word = word.lower() + if word in seg_dict: + out_txt += seg_dict[word] + " " + else: + if pattern.match(word): + for char in word: + if char in seg_dict: + out_txt += seg_dict[char] + " " + else: + out_txt += "" + " " + else: + out_txt += "" + " " + return out_txt.strip().split() diff --git a/almeval/models/stepaudio/funasr_detach/tokenizer/cleaner.py b/almeval/models/stepaudio/funasr_detach/tokenizer/cleaner.py new file mode 100644 index 0000000000000000000000000000000000000000..882db5c970c40350547a75d706f547ff12799430 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/tokenizer/cleaner.py @@ -0,0 +1,48 @@ +from typing import Collection + +from jaconv import jaconv + +# import tacotron_cleaner.cleaners + +try: + from vietnamese_cleaner import vietnamese_cleaners +except ImportError: + vietnamese_cleaners = None + + +class TextCleaner: + """Text cleaner. + + Examples: + >>> cleaner = TextCleaner("tacotron") + >>> cleaner("(Hello-World); & jr. & dr.") + 'HELLO WORLD, AND JUNIOR AND DOCTOR' + + """ + + def __init__(self, cleaner_types: Collection[str] = None): + + if cleaner_types is None: + self.cleaner_types = [] + elif isinstance(cleaner_types, str): + self.cleaner_types = [cleaner_types] + else: + self.cleaner_types = list(cleaner_types) + + def __call__(self, text: str) -> str: + for t in self.cleaner_types: + if t == "tacotron": + # text = tacotron_cleaner.cleaners.custom_english_cleaners(text) + pass + elif t == "jaconv": + text = jaconv.normalize(text) + elif t == "vietnamese": + if vietnamese_cleaners is None: + raise RuntimeError("Please install underthesea") + text = vietnamese_cleaners.vietnamese_cleaner(text) + elif t == "korean_cleaner": + text = KoreanCleaner.normalize_text(text) + else: + raise RuntimeError(f"Not supported: type={t}") + + return text diff --git a/almeval/models/stepaudio/funasr_detach/tokenizer/korean_cleaner.py b/almeval/models/stepaudio/funasr_detach/tokenizer/korean_cleaner.py new file mode 100644 index 0000000000000000000000000000000000000000..e98d0b81921894712cf02959f2df344e7bda047a --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/tokenizer/korean_cleaner.py @@ -0,0 +1,79 @@ +# Referenced from https://github.com/hccho2/Tacotron-Wavenet-Vocoder-Korean + +import re + + +class KoreanCleaner: + @classmethod + def _normalize_numbers(cls, text): + number_to_kor = { + "0": "영", + "1": "일", + "2": "이", + "3": "삼", + "4": "사", + "5": "오", + "6": "육", + "7": "칠", + "8": "팔", + "9": "구", + } + new_text = "".join( + number_to_kor[char] if char in number_to_kor.keys() else char + for char in text + ) + return new_text + + @classmethod + def _normalize_english_text(cls, text): + upper_alphabet_to_kor = { + "A": "에이", + "B": "비", + "C": "씨", + "D": "디", + "E": "이", + "F": "에프", + "G": "지", + "H": "에이치", + "I": "아이", + "J": "제이", + "K": "케이", + "L": "엘", + "M": "엠", + "N": "엔", + "O": "오", + "P": "피", + "Q": "큐", + "R": "알", + "S": "에스", + "T": "티", + "U": "유", + "V": "브이", + "W": "더블유", + "X": "엑스", + "Y": "와이", + "Z": "지", + } + new_text = re.sub("[a-z]+", lambda x: str.upper(x.group()), text) + new_text = "".join( + ( + upper_alphabet_to_kor[char] + if char in upper_alphabet_to_kor.keys() + else char + ) + for char in new_text + ) + + return new_text + + @classmethod + def normalize_text(cls, text): + # stage 0 : text strip + text = text.strip() + + # stage 1 : normalize numbers + text = cls._normalize_numbers(text) + + # stage 2 : normalize english text + text = cls._normalize_english_text(text) + return text diff --git a/almeval/models/stepaudio/funasr_detach/tokenizer/phoneme_tokenizer.py b/almeval/models/stepaudio/funasr_detach/tokenizer/phoneme_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..2ce6eaedb14d0c1bc467cde3c644c97ced3dd0b3 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/tokenizer/phoneme_tokenizer.py @@ -0,0 +1,526 @@ +import logging +from pathlib import Path +import re +from typing import Iterable +from typing import List +from typing import Optional +from typing import Union +import warnings + +# import g2p_en +import jamo + +from funasr_detach.tokenizer.abs_tokenizer import AbsTokenizer + + +g2p_classes = [ + None, + "g2p_en", + "g2p_en_no_space", + "pyopenjtalk", + "pyopenjtalk_kana", + "pyopenjtalk_accent", + "pyopenjtalk_accent_with_pause", + "pyopenjtalk_prosody", + "pypinyin_g2p", + "pypinyin_g2p_phone", + "espeak_ng_arabic", + "espeak_ng_german", + "espeak_ng_french", + "espeak_ng_spanish", + "espeak_ng_russian", + "espeak_ng_greek", + "espeak_ng_finnish", + "espeak_ng_hungarian", + "espeak_ng_dutch", + "espeak_ng_english_us_vits", + "espeak_ng_hindi", + "g2pk", + "g2pk_no_space", + "korean_jaso", + "korean_jaso_no_space", +] + + +def split_by_space(text) -> List[str]: + if " " in text: + text = text.replace(" ", " ") + return [c.replace("", " ") for c in text.split(" ")] + else: + return text.split(" ") + + +def pyopenjtalk_g2p(text) -> List[str]: + import pyopenjtalk + + # phones is a str object separated by space + phones = pyopenjtalk.g2p(text, kana=False) + phones = phones.split(" ") + return phones + + +def pyopenjtalk_g2p_accent(text) -> List[str]: + import pyopenjtalk + import re + + phones = [] + for labels in pyopenjtalk.run_frontend(text)[1]: + p = re.findall(r"\-(.*?)\+.*?\/A:([0-9\-]+).*?\/F:.*?_([0-9]+)", labels) + if len(p) == 1: + phones += [p[0][0], p[0][2], p[0][1]] + return phones + + +def pyopenjtalk_g2p_accent_with_pause(text) -> List[str]: + import pyopenjtalk + import re + + phones = [] + for labels in pyopenjtalk.run_frontend(text)[1]: + if labels.split("-")[1].split("+")[0] == "pau": + phones += ["pau"] + continue + p = re.findall(r"\-(.*?)\+.*?\/A:([0-9\-]+).*?\/F:.*?_([0-9]+)", labels) + if len(p) == 1: + phones += [p[0][0], p[0][2], p[0][1]] + return phones + + +def pyopenjtalk_g2p_kana(text) -> List[str]: + import pyopenjtalk + + kanas = pyopenjtalk.g2p(text, kana=True) + return list(kanas) + + +def pyopenjtalk_g2p_prosody(text: str, drop_unvoiced_vowels: bool = True) -> List[str]: + """Extract phoneme + prosoody symbol sequence from input full-context labels. + + The algorithm is based on `Prosodic features control by symbols as input of + sequence-to-sequence acoustic modeling for neural TTS`_ with some r9y9's tweaks. + + Args: + text (str): Input text. + drop_unvoiced_vowels (bool): whether to drop unvoiced vowels. + + Returns: + List[str]: List of phoneme + prosody symbols. + + Examples: + >>> from funasr_detach.tokenizer.phoneme_tokenizer import pyopenjtalk_g2p_prosody + >>> pyopenjtalk_g2p_prosody("こんにちは。") + ['^', 'k', 'o', '[', 'N', 'n', 'i', 'ch', 'i', 'w', 'a', '$'] + + .. _`Prosodic features control by symbols as input of sequence-to-sequence acoustic + modeling for neural TTS`: https://doi.org/10.1587/transinf.2020EDP7104 + + """ + import pyopenjtalk + + labels = pyopenjtalk.run_frontend(text)[1] + N = len(labels) + + phones = [] + for n in range(N): + lab_curr = labels[n] + + # current phoneme + p3 = re.search(r"\-(.*?)\+", lab_curr).group(1) + + # deal unvoiced vowels as normal vowels + if drop_unvoiced_vowels and p3 in "AEIOU": + p3 = p3.lower() + + # deal with sil at the beginning and the end of text + if p3 == "sil": + assert n == 0 or n == N - 1 + if n == 0: + phones.append("^") + elif n == N - 1: + # check question form or not + e3 = _numeric_feature_by_regex(r"!(\d+)_", lab_curr) + if e3 == 0: + phones.append("$") + elif e3 == 1: + phones.append("?") + continue + elif p3 == "pau": + phones.append("_") + continue + else: + phones.append(p3) + + # accent type and position info (forward or backward) + a1 = _numeric_feature_by_regex(r"/A:([0-9\-]+)\+", lab_curr) + a2 = _numeric_feature_by_regex(r"\+(\d+)\+", lab_curr) + a3 = _numeric_feature_by_regex(r"\+(\d+)/", lab_curr) + + # number of mora in accent phrase + f1 = _numeric_feature_by_regex(r"/F:(\d+)_", lab_curr) + + a2_next = _numeric_feature_by_regex(r"\+(\d+)\+", labels[n + 1]) + # accent phrase border + if a3 == 1 and a2_next == 1 and p3 in "aeiouAEIOUNcl": + phones.append("#") + # pitch falling + elif a1 == 0 and a2_next == a2 + 1 and a2 != f1: + phones.append("]") + # pitch rising + elif a2 == 1 and a2_next == 2: + phones.append("[") + + return phones + + +def _numeric_feature_by_regex(regex, s): + match = re.search(regex, s) + if match is None: + return -50 + return int(match.group(1)) + + +def pypinyin_g2p(text) -> List[str]: + from pypinyin import pinyin + from pypinyin import Style + + phones = [phone[0] for phone in pinyin(text, style=Style.TONE3)] + return phones + + +def pypinyin_g2p_phone(text) -> List[str]: + from pypinyin import pinyin + from pypinyin import Style + from pypinyin.style._utils import get_finals + from pypinyin.style._utils import get_initials + + phones = [ + p + for phone in pinyin(text, style=Style.TONE3) + for p in [ + get_initials(phone[0], strict=True), + get_finals(phone[0], strict=True), + ] + if len(p) != 0 + ] + return phones + + +class G2p_en: + """On behalf of g2p_en.G2p. + + g2p_en.G2p isn't pickalable and it can't be copied to the other processes + via multiprocessing module. + As a workaround, g2p_en.G2p is instantiated upon calling this class. + + """ + + def __init__(self, no_space: bool = False): + self.no_space = no_space + self.g2p = None + + def __call__(self, text) -> List[str]: + if self.g2p is None: + self.g2p = g2p_en.G2p() + + phones = self.g2p(text) + if self.no_space: + # remove space which represents word serapater + phones = list(filter(lambda s: s != " ", phones)) + return phones + + +class G2pk: + """On behalf of g2pk.G2p. + + g2pk.G2p isn't pickalable and it can't be copied to the other processes + via multiprocessing module. + As a workaround, g2pk.G2p is instantiated upon calling this class. + + """ + + def __init__( + self, descritive=False, group_vowels=False, to_syl=False, no_space=False + ): + self.descritive = descritive + self.group_vowels = group_vowels + self.to_syl = to_syl + self.no_space = no_space + self.g2p = None + + def __call__(self, text) -> List[str]: + if self.g2p is None: + import g2pk + + self.g2p = g2pk.G2p() + + phones = list( + self.g2p( + text, + descriptive=self.descritive, + group_vowels=self.group_vowels, + to_syl=self.to_syl, + ) + ) + if self.no_space: + # remove space which represents word serapater + phones = list(filter(lambda s: s != " ", phones)) + return phones + + +class Jaso: + PUNC = "!'(),-.:;?" + SPACE = " " + + JAMO_LEADS = "".join([chr(_) for _ in range(0x1100, 0x1113)]) + JAMO_VOWELS = "".join([chr(_) for _ in range(0x1161, 0x1176)]) + JAMO_TAILS = "".join([chr(_) for _ in range(0x11A8, 0x11C3)]) + + VALID_CHARS = JAMO_LEADS + JAMO_VOWELS + JAMO_TAILS + PUNC + SPACE + + def __init__(self, space_symbol=" ", no_space=False): + self.space_symbol = space_symbol + self.no_space = no_space + + def _text_to_jaso(self, line: str) -> List[str]: + jasos = list(jamo.hangul_to_jamo(line)) + return jasos + + def _remove_non_korean_characters(self, tokens): + new_tokens = [token for token in tokens if token in self.VALID_CHARS] + return new_tokens + + def __call__(self, text) -> List[str]: + graphemes = [x for x in self._text_to_jaso(text)] + graphemes = self._remove_non_korean_characters(graphemes) + + if self.no_space: + graphemes = list(filter(lambda s: s != " ", graphemes)) + else: + graphemes = [x if x != " " else self.space_symbol for x in graphemes] + return graphemes + + +class Phonemizer: + """Phonemizer module for various languages. + + This is wrapper module of https://github.com/bootphon/phonemizer. + You can define various g2p modules by specifying options for phonemizer. + + See available options: + https://github.com/bootphon/phonemizer/blob/master/phonemizer/phonemize.py#L32 + + """ + + def __init__( + self, + backend, + word_separator: Optional[str] = None, + syllable_separator: Optional[str] = None, + phone_separator: Optional[str] = " ", + strip=False, + split_by_single_token: bool = False, + **phonemizer_kwargs, + ): + # delayed import + from phonemizer.backend import BACKENDS + from phonemizer.separator import Separator + + self.separator = Separator( + word=word_separator, + syllable=syllable_separator, + phone=phone_separator, + ) + + # define logger to suppress the warning in phonemizer + logger = logging.getLogger("phonemizer") + logger.setLevel(logging.ERROR) + self.phonemizer = BACKENDS[backend]( + **phonemizer_kwargs, + logger=logger, + ) + self.strip = strip + self.split_by_single_token = split_by_single_token + + def __call__(self, text) -> List[str]: + tokens = self.phonemizer.phonemize( + [text], + separator=self.separator, + strip=self.strip, + njobs=1, + )[0] + if not self.split_by_single_token: + return tokens.split() + else: + # "a: ab" -> ["a", ":", "", "a", "b"] + # TODO(kan-bayashi): space replacement should be dealt in PhonemeTokenizer + return [c.replace(" ", "") for c in tokens] + + +class PhonemeTokenizer(AbsTokenizer): + def __init__( + self, + g2p_type: Union[None, str], + non_linguistic_symbols: Union[Path, str, Iterable[str]] = None, + space_symbol: str = "", + remove_non_linguistic_symbols: bool = False, + ): + if g2p_type is None: + self.g2p = split_by_space + elif g2p_type == "g2p_en": + self.g2p = G2p_en(no_space=False) + elif g2p_type == "g2p_en_no_space": + self.g2p = G2p_en(no_space=True) + elif g2p_type == "pyopenjtalk": + self.g2p = pyopenjtalk_g2p + elif g2p_type == "pyopenjtalk_kana": + self.g2p = pyopenjtalk_g2p_kana + elif g2p_type == "pyopenjtalk_accent": + self.g2p = pyopenjtalk_g2p_accent + elif g2p_type == "pyopenjtalk_accent_with_pause": + self.g2p = pyopenjtalk_g2p_accent_with_pause + elif g2p_type == "pyopenjtalk_prosody": + self.g2p = pyopenjtalk_g2p_prosody + elif g2p_type == "pypinyin_g2p": + self.g2p = pypinyin_g2p + elif g2p_type == "pypinyin_g2p_phone": + self.g2p = pypinyin_g2p_phone + elif g2p_type == "espeak_ng_arabic": + self.g2p = Phonemizer( + language="ar", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_german": + self.g2p = Phonemizer( + language="de", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_french": + self.g2p = Phonemizer( + language="fr-fr", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_spanish": + self.g2p = Phonemizer( + language="es", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_russian": + self.g2p = Phonemizer( + language="ru", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_greek": + self.g2p = Phonemizer( + language="el", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_finnish": + self.g2p = Phonemizer( + language="fi", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_hungarian": + self.g2p = Phonemizer( + language="hu", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_dutch": + self.g2p = Phonemizer( + language="nl", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "espeak_ng_hindi": + self.g2p = Phonemizer( + language="hi", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + ) + elif g2p_type == "g2pk": + self.g2p = G2pk(no_space=False) + elif g2p_type == "g2pk_no_space": + self.g2p = G2pk(no_space=True) + elif g2p_type == "espeak_ng_english_us_vits": + # VITS official implementation-like processing + # Reference: https://github.com/jaywalnut310/vits + self.g2p = Phonemizer( + language="en-us", + backend="espeak", + with_stress=True, + preserve_punctuation=True, + strip=True, + word_separator=" ", + phone_separator="", + split_by_single_token=True, + ) + elif g2p_type == "korean_jaso": + self.g2p = Jaso(space_symbol=space_symbol, no_space=False) + elif g2p_type == "korean_jaso_no_space": + self.g2p = Jaso(no_space=True) + else: + raise NotImplementedError(f"Not supported: g2p_type={g2p_type}") + + self.g2p_type = g2p_type + self.space_symbol = space_symbol + if non_linguistic_symbols is None: + self.non_linguistic_symbols = set() + elif isinstance(non_linguistic_symbols, (Path, str)): + non_linguistic_symbols = Path(non_linguistic_symbols) + try: + with non_linguistic_symbols.open("r", encoding="utf-8") as f: + self.non_linguistic_symbols = set(line.rstrip() for line in f) + except FileNotFoundError: + warnings.warn(f"{non_linguistic_symbols} doesn't exist.") + self.non_linguistic_symbols = set() + else: + self.non_linguistic_symbols = set(non_linguistic_symbols) + self.remove_non_linguistic_symbols = remove_non_linguistic_symbols + + def __repr__(self): + return ( + f"{self.__class__.__name__}(" + f'g2p_type="{self.g2p_type}", ' + f'space_symbol="{self.space_symbol}", ' + f'non_linguistic_symbols="{self.non_linguistic_symbols}"' + ")" + ) + + def text2tokens(self, line: str) -> List[str]: + tokens = [] + while len(line) != 0: + for w in self.non_linguistic_symbols: + if line.startswith(w): + if not self.remove_non_linguistic_symbols: + tokens.append(line[: len(w)]) + line = line[len(w) :] + break + else: + t = line[0] + tokens.append(t) + line = line[1:] + + line = "".join(tokens) + tokens = self.g2p(line) + return tokens + + def tokens2text(self, tokens: Iterable[str]) -> str: + # phoneme type is not invertible + return "".join(tokens) diff --git a/almeval/models/stepaudio/funasr_detach/tokenizer/sentencepiece_tokenizer.py b/almeval/models/stepaudio/funasr_detach/tokenizer/sentencepiece_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..573947c35ac9b46af02afcbc5d0a6aa573860939 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/tokenizer/sentencepiece_tokenizer.py @@ -0,0 +1,47 @@ +from pathlib import Path +from typing import Iterable +from typing import List +from typing import Union + +import sentencepiece as spm + +from funasr_detach.tokenizer.abs_tokenizer import BaseTokenizer +from funasr_detach.register import tables + + +@tables.register("tokenizer_classes", "SentencepiecesTokenizer") +class SentencepiecesTokenizer(BaseTokenizer): + def __init__(self, bpemodel: Union[Path, str], **kwargs): + super().__init__(**kwargs) + self.bpemodel = str(bpemodel) + # NOTE(kamo): + # Don't build SentencePieceProcessor in __init__() + # because it's not picklable and it may cause following error, + # "TypeError: can't pickle SwigPyObject objects", + # when giving it as argument of "multiprocessing.Process()". + self.sp = None + + def __repr__(self): + return f'{self.__class__.__name__}(model="{self.bpemodel}")' + + def _build_sentence_piece_processor(self): + # Build SentencePieceProcessor lazily. + if self.sp is None: + self.sp = spm.SentencePieceProcessor() + self.sp.load(self.bpemodel) + + def text2tokens(self, line: str) -> List[str]: + self._build_sentence_piece_processor() + return self.sp.EncodeAsPieces(line) + + def tokens2text(self, tokens: Iterable[str]) -> str: + self._build_sentence_piece_processor() + return self.sp.DecodePieces(list(tokens)) + + def encode(self, line: str) -> List[int]: + self._build_sentence_piece_processor() + return self.sp.EncodeAsIds(line) + + def decode(self, line: List[int]): + self._build_sentence_piece_processor() + return self.sp.DecodeIds(line) diff --git a/almeval/models/stepaudio/funasr_detach/tokenizer/token_id_converter.py b/almeval/models/stepaudio/funasr_detach/tokenizer/token_id_converter.py new file mode 100644 index 0000000000000000000000000000000000000000..1888d750837b008e128b9ae89b8c109fc17a167e --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/tokenizer/token_id_converter.py @@ -0,0 +1,58 @@ +from pathlib import Path +from typing import Dict +from typing import Iterable +from typing import List +from typing import Union + +import numpy as np + + +class TokenIDConverter: + def __init__( + self, + token_list: Union[Path, str, Iterable[str]], + unk_symbol: str = "", + ): + + if isinstance(token_list, (Path, str)): + token_list = Path(token_list) + self.token_list_repr = str(token_list) + self.token_list: List[str] = [] + + with token_list.open("r", encoding="utf-8") as f: + for idx, line in enumerate(f): + line = line.rstrip() + self.token_list.append(line) + + else: + self.token_list: List[str] = list(token_list) + self.token_list_repr = "" + for i, t in enumerate(self.token_list): + if i == 3: + break + self.token_list_repr += f"{t}, " + self.token_list_repr += f"... (NVocab={(len(self.token_list))})" + + self.token2id: Dict[str, int] = {} + for i, t in enumerate(self.token_list): + if t in self.token2id: + raise RuntimeError(f'Symbol "{t}" is duplicated') + self.token2id[t] = i + + self.unk_symbol = unk_symbol + if self.unk_symbol not in self.token2id: + raise RuntimeError( + f"Unknown symbol '{unk_symbol}' doesn't exist in the token_list" + ) + self.unk_id = self.token2id[self.unk_symbol] + + def get_num_vocabulary_size(self) -> int: + return len(self.token_list) + + def ids2tokens(self, integers: Union[np.ndarray, Iterable[int]]) -> List[str]: + if isinstance(integers, np.ndarray) and integers.ndim != 1: + raise ValueError(f"Must be 1 dim ndarray, but got {integers.ndim}") + return [self.token_list[i] for i in integers] + + def tokens2ids(self, tokens: Iterable[str]) -> List[int]: + return [self.token2id.get(i, self.unk_id) for i in tokens] diff --git a/almeval/models/stepaudio/funasr_detach/tokenizer/word_tokenizer.py b/almeval/models/stepaudio/funasr_detach/tokenizer/word_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..8a9ed307669d9f57b53d534eb294f684c56aecd2 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/tokenizer/word_tokenizer.py @@ -0,0 +1,56 @@ +from pathlib import Path +from typing import Iterable +from typing import List +from typing import Union +import warnings + + +from funasr_detach.tokenizer.abs_tokenizer import AbsTokenizer + + +class WordTokenizer(AbsTokenizer): + def __init__( + self, + delimiter: str = None, + non_linguistic_symbols: Union[Path, str, Iterable[str]] = None, + remove_non_linguistic_symbols: bool = False, + ): + self.delimiter = delimiter + + if not remove_non_linguistic_symbols and non_linguistic_symbols is not None: + warnings.warn( + "non_linguistic_symbols is only used " + "when remove_non_linguistic_symbols = True" + ) + + if non_linguistic_symbols is None: + self.non_linguistic_symbols = set() + elif isinstance(non_linguistic_symbols, (Path, str)): + non_linguistic_symbols = Path(non_linguistic_symbols) + try: + with non_linguistic_symbols.open("r", encoding="utf-8") as f: + self.non_linguistic_symbols = set(line.rstrip() for line in f) + except FileNotFoundError: + warnings.warn(f"{non_linguistic_symbols} doesn't exist.") + self.non_linguistic_symbols = set() + else: + self.non_linguistic_symbols = set(non_linguistic_symbols) + self.remove_non_linguistic_symbols = remove_non_linguistic_symbols + + def __repr__(self): + return f'{self.__class__.__name__}(delimiter="{self.delimiter}")' + + def text2tokens(self, line: str) -> List[str]: + tokens = [] + for t in line.split(self.delimiter): + if self.remove_non_linguistic_symbols and t in self.non_linguistic_symbols: + continue + tokens.append(t) + return tokens + + def tokens2text(self, tokens: Iterable[str]) -> str: + if self.delimiter is None: + delimiter = " " + else: + delimiter = self.delimiter + return delimiter.join(tokens) diff --git a/almeval/models/stepaudio/funasr_detach/utils/__init__.py b/almeval/models/stepaudio/funasr_detach/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/almeval/models/stepaudio/funasr_detach/utils/datadir_writer.py b/almeval/models/stepaudio/funasr_detach/utils/datadir_writer.py new file mode 100644 index 0000000000000000000000000000000000000000..dd64018e9feffa40d4f8e29f21f5aaf4f3eb113e --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/utils/datadir_writer.py @@ -0,0 +1,71 @@ +from pathlib import Path +from typing import Union +import warnings + + +class DatadirWriter: + """Writer class to create kaldi like data directory. + + Examples: + >>> with DatadirWriter("output") as writer: + ... # output/sub.txt is created here + ... subwriter = writer["sub.txt"] + ... # Write "uttidA some/where/a.wav" + ... subwriter["uttidA"] = "some/where/a.wav" + ... subwriter["uttidB"] = "some/where/b.wav" + + """ + + def __init__(self, p: Union[Path, str]): + self.path = Path(p) + self.chilidren = {} + self.fd = None + self.has_children = False + self.keys = set() + + def __enter__(self): + return self + + def __getitem__(self, key: str) -> "DatadirWriter": + if self.fd is not None: + raise RuntimeError("This writer points out a file") + + if key not in self.chilidren: + w = DatadirWriter((self.path / key)) + self.chilidren[key] = w + self.has_children = True + + retval = self.chilidren[key] + return retval + + def __setitem__(self, key: str, value: str): + if self.has_children: + raise RuntimeError("This writer points out a directory") + if key in self.keys: + warnings.warn(f"Duplicated: {key}") + + if self.fd is None: + self.path.parent.mkdir(parents=True, exist_ok=True) + self.fd = self.path.open("w", encoding="utf-8") + + self.keys.add(key) + self.fd.write(f"{key} {value}\n") + self.fd.flush() + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + + def close(self): + if self.has_children: + prev_child = None + for child in self.chilidren.values(): + child.close() + if prev_child is not None and prev_child.keys != child.keys: + warnings.warn( + f"Ids are mismatching between " + f"{prev_child.path} and {child.path}" + ) + prev_child = child + + elif self.fd is not None: + self.fd.close() diff --git a/almeval/models/stepaudio/funasr_detach/utils/load_utils.py b/almeval/models/stepaudio/funasr_detach/utils/load_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..548f41a6a0a1f1a2365149cbeff99f55d94175d5 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/utils/load_utils.py @@ -0,0 +1,148 @@ +import os +import io +import torch +import numpy as np +import torchaudio +from torch.nn.utils.rnn import pad_sequence + +try: + from funasr_detach.download.file import download_from_url +except: + print("urllib is not installed, if you infer from url, please install it first.") + + +def load_audio_text_image_video( + data_or_path_or_list, + fs: int = 16000, + audio_fs: int = 16000, + data_type="sound", + tokenizer=None, + **kwargs +): + if isinstance(data_or_path_or_list, (list, tuple)): + if data_type is not None and isinstance(data_type, (list, tuple)): + + data_types = [data_type] * len(data_or_path_or_list) + data_or_path_or_list_ret = [[] for d in data_type] + for i, (data_type_i, data_or_path_or_list_i) in enumerate( + zip(data_types, data_or_path_or_list) + ): + + for j, (data_type_j, data_or_path_or_list_j) in enumerate( + zip(data_type_i, data_or_path_or_list_i) + ): + + data_or_path_or_list_j = load_audio_text_image_video( + data_or_path_or_list_j, + fs=fs, + audio_fs=audio_fs, + data_type=data_type_j, + tokenizer=tokenizer, + **kwargs + ) + data_or_path_or_list_ret[j].append(data_or_path_or_list_j) + + return data_or_path_or_list_ret + else: + return [ + load_audio_text_image_video( + audio, fs=fs, audio_fs=audio_fs, data_type=data_type, **kwargs + ) + for audio in data_or_path_or_list + ] + + if isinstance(data_or_path_or_list, str) and data_or_path_or_list.startswith( + "http" + ): # download url to local file + data_or_path_or_list = download_from_url(data_or_path_or_list) + + if isinstance(data_or_path_or_list, io.BytesIO): + data_or_path_or_list, audio_fs = torchaudio.load(data_or_path_or_list) + if kwargs.get("reduce_channels", True): + data_or_path_or_list = data_or_path_or_list.mean(0) + elif isinstance(data_or_path_or_list, str) and os.path.exists( + data_or_path_or_list + ): # local file + if data_type is None or data_type == "sound": + data_or_path_or_list, audio_fs = torchaudio.load(data_or_path_or_list) + if kwargs.get("reduce_channels", True): + data_or_path_or_list = data_or_path_or_list.mean(0) + elif data_type == "text" and tokenizer is not None: + data_or_path_or_list = tokenizer.encode(data_or_path_or_list) + elif data_type == "image": # undo + pass + elif data_type == "video": # undo + pass + + # if data_in is a file or url, set is_final=True + if "cache" in kwargs: + kwargs["cache"]["is_final"] = True + kwargs["cache"]["is_streaming_input"] = False + elif ( + isinstance(data_or_path_or_list, str) + and data_type == "text" + and tokenizer is not None + ): + data_or_path_or_list = tokenizer.encode(data_or_path_or_list) + elif isinstance(data_or_path_or_list, np.ndarray): # audio sample point + data_or_path_or_list = torch.from_numpy( + data_or_path_or_list + ).squeeze() # [n_samples,] + else: + pass + # print(f"unsupport data type: {data_or_path_or_list}, return raw data") + + if audio_fs != fs and data_type != "text": + resampler = torchaudio.transforms.Resample(audio_fs, fs) + data_or_path_or_list = resampler(data_or_path_or_list[None, :])[0, :] + return data_or_path_or_list + + +def load_bytes(input): + middle_data = np.frombuffer(input, dtype=np.int16) + middle_data = np.asarray(middle_data) + if middle_data.dtype.kind not in "iu": + raise TypeError("'middle_data' must be an array of integers") + dtype = np.dtype("float32") + if dtype.kind != "f": + raise TypeError("'dtype' must be a floating point type") + + i = np.iinfo(middle_data.dtype) + abs_max = 2 ** (i.bits - 1) + offset = i.min + abs_max + array = np.frombuffer( + (middle_data.astype(dtype) - offset) / abs_max, dtype=np.float32 + ) + return array + + +def extract_fbank( + data, data_len=None, data_type: str = "sound", frontend=None, **kwargs +): + # import pdb; + # pdb.set_trace() + if isinstance(data, np.ndarray): + data = torch.from_numpy(data) + if len(data.shape) < 2: + data = data[None, :] # data: [batch, N] + data_len = [data.shape[1]] if data_len is None else data_len + elif isinstance(data, torch.Tensor): + if len(data.shape) < 2: + data = data[None, :] # data: [batch, N] + data_len = [data.shape[1]] if data_len is None else data_len + elif isinstance(data, (list, tuple)): + data_list, data_len = [], [] + for data_i in data: + if isinstance(data_i, np.ndarray): + data_i = torch.from_numpy(data_i) + data_list.append(data_i) + data_len.append(data_i.shape[0]) + data = pad_sequence(data_list, batch_first=True) # data: [batch, N] + # import pdb; + # pdb.set_trace() + # if data_type == "sound": + data, data_len = frontend(data, data_len, **kwargs) + + if isinstance(data_len, (list, tuple)): + data_len = torch.tensor([data_len]) + return data.to(torch.float32), data_len.to(torch.int32) diff --git a/almeval/models/stepaudio/funasr_detach/utils/misc.py b/almeval/models/stepaudio/funasr_detach/utils/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..b8aaa75333be867754934ba85b27b2cfbcccb351 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/utils/misc.py @@ -0,0 +1,50 @@ +import io +from collections import OrderedDict +import numpy as np + + +def statistic_model_parameters(model, prefix=None): + var_dict = model.state_dict() + numel = 0 + for i, key in enumerate( + sorted(list([x for x in var_dict.keys() if "num_batches_tracked" not in x])) + ): + if prefix is None or key.startswith(prefix): + numel += var_dict[key].numel() + return numel + + +def int2vec(x, vec_dim=8, dtype=np.int32): + b = ("{:0" + str(vec_dim) + "b}").format(x) + # little-endian order: lower bit first + return (np.array(list(b)[::-1]) == "1").astype(dtype) + + +def seq2arr(seq, vec_dim=8): + return np.row_stack([int2vec(int(x), vec_dim) for x in seq]) + + +def load_scp_as_dict(scp_path, value_type="str", kv_sep=" "): + with io.open(scp_path, "r", encoding="utf-8") as f: + ret_dict = OrderedDict() + for one_line in f.readlines(): + one_line = one_line.strip() + pos = one_line.find(kv_sep) + key, value = one_line[:pos], one_line[pos + 1 :] + if value_type == "list": + value = value.split(" ") + ret_dict[key] = value + return ret_dict + + +def load_scp_as_list(scp_path, value_type="str", kv_sep=" "): + with io.open(scp_path, "r", encoding="utf8") as f: + ret_dict = [] + for one_line in f.readlines(): + one_line = one_line.strip() + pos = one_line.find(kv_sep) + key, value = one_line[:pos], one_line[pos + 1 :] + if value_type == "list": + value = value.split(" ") + ret_dict.append((key, value)) + return ret_dict diff --git a/almeval/models/stepaudio/funasr_detach/utils/postprocess_utils.py b/almeval/models/stepaudio/funasr_detach/utils/postprocess_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..34f933c967d4add7680ad371cd2565af708d461b --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/utils/postprocess_utils.py @@ -0,0 +1,301 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import string +import logging +from typing import Any, List, Union + + +def isChinese(ch: str): + if "\u4e00" <= ch <= "\u9fff" or "\u0030" <= ch <= "\u0039" or ch == "@": + return True + return False + + +def isAllChinese(word: Union[List[Any], str]): + word_lists = [] + for i in word: + cur = i.replace(" ", "") + cur = cur.replace("", "") + cur = cur.replace("", "") + cur = cur.replace("", "") + cur = cur.replace("", "") + word_lists.append(cur) + + if len(word_lists) == 0: + return False + + for ch in word_lists: + if isChinese(ch) is False: + return False + return True + + +def isAllAlpha(word: Union[List[Any], str]): + word_lists = [] + for i in word: + cur = i.replace(" ", "") + cur = cur.replace("", "") + cur = cur.replace("", "") + cur = cur.replace("", "") + cur = cur.replace("", "") + word_lists.append(cur) + + if len(word_lists) == 0: + return False + + for ch in word_lists: + if ch.isalpha() is False and ch != "'": + return False + elif ch.isalpha() is True and isChinese(ch) is True: + return False + + return True + + +# def abbr_dispose(words: List[Any]) -> List[Any]: +def abbr_dispose(words: List[Any], time_stamp: List[List] = None) -> List[Any]: + words_size = len(words) + word_lists = [] + abbr_begin = [] + abbr_end = [] + last_num = -1 + ts_lists = [] + ts_nums = [] + ts_index = 0 + for num in range(words_size): + if num <= last_num: + continue + + if len(words[num]) == 1 and words[num].encode("utf-8").isalpha(): + if ( + num + 1 < words_size + and words[num + 1] == " " + and num + 2 < words_size + and len(words[num + 2]) == 1 + and words[num + 2].encode("utf-8").isalpha() + ): + # found the begin of abbr + abbr_begin.append(num) + num += 2 + abbr_end.append(num) + # to find the end of abbr + while True: + num += 1 + if num < words_size and words[num] == " ": + num += 1 + if ( + num < words_size + and len(words[num]) == 1 + and words[num].encode("utf-8").isalpha() + ): + abbr_end.pop() + abbr_end.append(num) + last_num = num + else: + break + else: + break + + for num in range(words_size): + if words[num] == " ": + ts_nums.append(ts_index) + else: + ts_nums.append(ts_index) + ts_index += 1 + last_num = -1 + for num in range(words_size): + if num <= last_num: + continue + + if num in abbr_begin: + if time_stamp is not None: + begin = time_stamp[ts_nums[num]][0] + abbr_word = words[num].upper() + num += 1 + while num < words_size: + if num in abbr_end: + abbr_word += words[num].upper() + last_num = num + break + else: + if words[num].encode("utf-8").isalpha(): + abbr_word += words[num].upper() + num += 1 + word_lists.append(abbr_word) + if time_stamp is not None: + end = time_stamp[ts_nums[num]][1] + ts_lists.append([begin, end]) + else: + word_lists.append(words[num]) + if time_stamp is not None and words[num] != " ": + begin = time_stamp[ts_nums[num]][0] + end = time_stamp[ts_nums[num]][1] + ts_lists.append([begin, end]) + begin = end + + if time_stamp is not None: + return word_lists, ts_lists + else: + return word_lists + + +def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): + middle_lists = [] + word_lists = [] + word_item = "" + ts_lists = [] + + # wash words lists + for i in words: + word = "" + if isinstance(i, str): + word = i + else: + word = i.decode("utf-8") + + if word in ["", "", "", ""]: + continue + else: + middle_lists.append(word) + + # all chinese characters + if isAllChinese(middle_lists): + for i, ch in enumerate(middle_lists): + word_lists.append(ch.replace(" ", "")) + if time_stamp is not None: + ts_lists = time_stamp + + # all alpha characters + elif isAllAlpha(middle_lists): + ts_flag = True + for i, ch in enumerate(middle_lists): + if ts_flag and time_stamp is not None: + begin = time_stamp[i][0] + end = time_stamp[i][1] + word = "" + if "@@" in ch: + word = ch.replace("@@", "") + word_item += word + if time_stamp is not None: + ts_flag = False + end = time_stamp[i][1] + else: + word_item += ch + word_lists.append(word_item) + word_lists.append(" ") + word_item = "" + if time_stamp is not None: + ts_flag = True + end = time_stamp[i][1] + ts_lists.append([begin, end]) + begin = end + + # mix characters + else: + alpha_blank = False + ts_flag = True + begin = -1 + end = -1 + for i, ch in enumerate(middle_lists): + if ts_flag and time_stamp is not None: + begin = time_stamp[i][0] + end = time_stamp[i][1] + word = "" + if isAllChinese(ch): + if alpha_blank is True: + word_lists.pop() + word_lists.append(ch) + alpha_blank = False + if time_stamp is not None: + ts_flag = True + ts_lists.append([begin, end]) + begin = end + elif "@@" in ch: + word = ch.replace("@@", "") + word_item += word + alpha_blank = False + if time_stamp is not None: + ts_flag = False + end = time_stamp[i][1] + elif isAllAlpha(ch): + word_item += ch + word_lists.append(word_item) + word_lists.append(" ") + word_item = "" + alpha_blank = True + if time_stamp is not None: + ts_flag = True + end = time_stamp[i][1] + ts_lists.append([begin, end]) + begin = end + else: + word_lists.append(ch) + + if time_stamp is not None: + word_lists, ts_lists = abbr_dispose(word_lists, ts_lists) + real_word_lists = [] + for ch in word_lists: + if ch != " ": + real_word_lists.append(ch) + sentence = " ".join(real_word_lists).strip() + return sentence, ts_lists, real_word_lists + else: + word_lists = abbr_dispose(word_lists) + real_word_lists = [] + for ch in word_lists: + if ch != " ": + real_word_lists.append(ch) + sentence = "".join(word_lists).strip() + return sentence, real_word_lists + + +def sentence_postprocess_sentencepiece(words): + middle_lists = [] + word_lists = [] + word_item = "" + + # wash words lists + for i in words: + word = "" + if isinstance(i, str): + word = i + else: + word = i.decode("utf-8") + + if word in ["", "", "", ""]: + continue + else: + middle_lists.append(word) + + # all alpha characters + for i, ch in enumerate(middle_lists): + word = "" + if "\u2581" in ch and i == 0: + word_item = "" + word = ch.replace("\u2581", "") + word_item += word + elif "\u2581" in ch and i != 0: + word_lists.append(word_item) + word_lists.append(" ") + word_item = "" + word = ch.replace("\u2581", "") + word_item += word + else: + word_item += ch + if word_item is not None: + word_lists.append(word_item) + # word_lists = abbr_dispose(word_lists) + real_word_lists = [] + for ch in word_lists: + if ch != " ": + if ch == "i": + ch = ch.replace("i", "I") + elif ch == "i'm": + ch = ch.replace("i'm", "I'm") + elif ch == "i've": + ch = ch.replace("i've", "I've") + elif ch == "i'll": + ch = ch.replace("i'll", "I'll") + real_word_lists.append(ch) + sentence = "".join(word_lists) + return sentence, real_word_lists diff --git a/almeval/models/stepaudio/funasr_detach/utils/prepare_data.py b/almeval/models/stepaudio/funasr_detach/utils/prepare_data.py new file mode 100644 index 0000000000000000000000000000000000000000..356cd3651a1690cd9cb46409a97498c2f6ea972e --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/utils/prepare_data.py @@ -0,0 +1,308 @@ +import logging +import os +import shutil +from multiprocessing import Pool + +import kaldiio +import numpy as np +import librosa +import torch.distributed as dist +import torchaudio + + +def filter_wav_text(data_dir, dataset): + wav_file = os.path.join(data_dir, dataset, "wav.scp") + text_file = os.path.join(data_dir, dataset, "text") + with open(wav_file) as f_wav, open(text_file) as f_text: + wav_lines = f_wav.readlines() + text_lines = f_text.readlines() + os.rename(wav_file, "{}.bak".format(wav_file)) + os.rename(text_file, "{}.bak".format(text_file)) + wav_dict = {} + for line in wav_lines: + parts = line.strip().split() + if len(parts) < 2: + continue + wav_dict[parts[0]] = parts[1] + text_dict = {} + for line in text_lines: + parts = line.strip().split() + if len(parts) < 2: + continue + text_dict[parts[0]] = " ".join(parts[1:]) + filter_count = 0 + with open(wav_file, "w") as f_wav, open(text_file, "w") as f_text: + for sample_name, wav_path in wav_dict.items(): + if sample_name in text_dict.keys(): + f_wav.write(sample_name + " " + wav_path + "\n") + f_text.write(sample_name + " " + text_dict[sample_name] + "\n") + else: + filter_count += 1 + logging.info( + "{}/{} samples in {} are filtered because of the mismatch between wav.scp and text".format( + filter_count, len(wav_lines), dataset + ) + ) + + +def wav2num_frame(wav_path, frontend_conf): + try: + waveform, sampling_rate = torchaudio.load(wav_path) + except: + waveform, sampling_rate = librosa.load(wav_path) + waveform = np.expand_dims(waveform, axis=0) + n_frames = (waveform.shape[1] * 1000.0) / ( + sampling_rate * frontend_conf["frame_shift"] * frontend_conf["lfr_n"] + ) + feature_dim = frontend_conf["n_mels"] * frontend_conf["lfr_m"] + return n_frames, feature_dim + + +def calc_shape_core(root_path, args, idx): + file_name = args.data_file_names.split(",")[0] + data_name = args.dataset_conf.get("data_names", "speech,text").split(",")[0] + scp_file = os.path.join(root_path, "{}.{}".format(file_name, idx)) + shape_file = os.path.join(root_path, "{}_shape.{}".format(data_name, idx)) + with open(scp_file) as f: + lines = f.readlines() + data_type = args.dataset_conf.get("data_types", "sound,text").split(",")[0] + if data_type == "sound": + frontend_conf = args.frontend_conf + dataset_conf = args.dataset_conf + length_min = ( + dataset_conf.speech_length_min + if hasattr(dataset_conf, "{}_length_min".format(data_name)) + else -1 + ) + length_max = ( + dataset_conf.speech_length_max + if hasattr(dataset_conf, "{}_length_max".format(data_name)) + else -1 + ) + with open(shape_file, "w") as f: + for line in lines: + sample_name, wav_path = line.strip().split() + n_frames, feature_dim = wav2num_frame(wav_path, frontend_conf) + write_flag = True + if n_frames > 0 and length_min > 0: + write_flag = n_frames >= length_min + if n_frames > 0 and length_max > 0: + write_flag = n_frames <= length_max + if write_flag: + f.write( + "{} {},{}\n".format( + sample_name, + str(int(np.ceil(n_frames))), + str(int(feature_dim)), + ) + ) + f.flush() + elif data_type == "kaldi_ark": + dataset_conf = args.dataset_conf + length_min = ( + dataset_conf.speech_length_min + if hasattr(dataset_conf, "{}_length_min".format(data_name)) + else -1 + ) + length_max = ( + dataset_conf.speech_length_max + if hasattr(dataset_conf, "{}_length_max".format(data_name)) + else -1 + ) + with open(shape_file, "w") as f: + for line in lines: + sample_name, feature_path = line.strip().split() + feature = kaldiio.load_mat(feature_path) + n_frames, feature_dim = feature.shape + write_flag = True + if n_frames > 0 and length_min > 0: + write_flag = n_frames >= length_min + if n_frames > 0 and length_max > 0: + write_flag = n_frames <= length_max + if write_flag: + f.write( + "{} {},{}\n".format( + sample_name, + str(int(np.ceil(n_frames))), + str(int(feature_dim)), + ) + ) + f.flush() + elif data_type == "text": + with open(shape_file, "w") as f: + for line in lines: + sample_name, text = line.strip().split(maxsplit=1) + n_tokens = len(text.split()) + f.write("{} {}\n".format(sample_name, str(int(np.ceil(n_tokens))))) + f.flush() + else: + raise RuntimeError("Unsupported data_type: {}".format(data_type)) + + +def calc_shape(args, dataset, nj=64): + data_name = args.dataset_conf.get("data_names", "speech,text").split(",")[0] + shape_path = os.path.join(args.data_dir, dataset, "{}_shape".format(data_name)) + if os.path.exists(shape_path): + logging.info("Shape file for small dataset already exists.") + return + + split_shape_path = os.path.join( + args.data_dir, dataset, "{}_shape_files".format(data_name) + ) + if os.path.exists(split_shape_path): + shutil.rmtree(split_shape_path) + os.mkdir(split_shape_path) + + # split + file_name = args.data_file_names.split(",")[0] + scp_file = os.path.join(args.data_dir, dataset, file_name) + with open(scp_file) as f: + lines = f.readlines() + num_lines = len(lines) + num_job_lines = num_lines // nj + start = 0 + for i in range(nj): + end = start + num_job_lines + file = os.path.join(split_shape_path, "{}.{}".format(file_name, str(i + 1))) + with open(file, "w") as f: + if i == nj - 1: + f.writelines(lines[start:]) + else: + f.writelines(lines[start:end]) + start = end + + p = Pool(nj) + for i in range(nj): + p.apply_async(calc_shape_core, args=(split_shape_path, args, str(i + 1))) + logging.info("Generating shape files, please wait a few minutes...") + p.close() + p.join() + + # combine + with open(shape_path, "w") as f: + for i in range(nj): + job_file = os.path.join( + split_shape_path, "{}_shape.{}".format(data_name, str(i + 1)) + ) + with open(job_file) as job_f: + lines = job_f.readlines() + f.writelines(lines) + logging.info("Generating shape files done.") + + +def generate_data_list(args, data_dir, dataset, nj=64): + data_names = args.dataset_conf.get("data_names", "speech,text").split(",") + file_names = args.data_file_names.split(",") + concat_data_name = "_".join(data_names) + list_file = os.path.join(data_dir, dataset, "{}_data.list".format(concat_data_name)) + if os.path.exists(list_file): + logging.info("Data list for large dataset already exists.") + return + split_path = os.path.join(data_dir, dataset, "split") + if os.path.exists(split_path): + shutil.rmtree(split_path) + os.mkdir(split_path) + + data_lines_list = [] + for file_name in file_names: + with open(os.path.join(data_dir, dataset, file_name)) as f: + lines = f.readlines() + data_lines_list.append(lines) + num_lines = len(data_lines_list[0]) + num_job_lines = num_lines // nj + start = 0 + for i in range(nj): + end = start + num_job_lines + split_path_nj = os.path.join(split_path, str(i + 1)) + os.mkdir(split_path_nj) + for file_id, file_name in enumerate(file_names): + file = os.path.join(split_path_nj, file_name) + with open(file, "w") as f: + if i == nj - 1: + f.writelines(data_lines_list[file_id][start:]) + else: + f.writelines(data_lines_list[file_id][start:end]) + start = end + + with open(list_file, "w") as f_data: + for i in range(nj): + path = "" + for file_name in file_names: + path = path + " " + os.path.join(split_path, str(i + 1), file_name) + f_data.write(path + "\n") + + +def prepare_data(args, distributed_option): + data_names = args.dataset_conf.get("data_names", "speech,text").split(",") + data_types = args.dataset_conf.get("data_types", "sound,text").split(",") + file_names = args.data_file_names.split(",") + batch_type = args.dataset_conf["batch_conf"]["batch_type"] + print( + "data_names: {}, data_types: {}, file_names: {}".format( + data_names, data_types, file_names + ) + ) + assert len(data_names) == len(data_types) == len(file_names) + if args.dataset_type == "small": + args.train_shape_file = [ + os.path.join( + args.data_dir, args.train_set, "{}_shape".format(data_names[0]) + ) + ] + args.valid_shape_file = [ + os.path.join( + args.data_dir, args.valid_set, "{}_shape".format(data_names[0]) + ) + ] + ( + args.train_data_path_and_name_and_type, + args.valid_data_path_and_name_and_type, + ) = ([], []) + for file_name, data_name, data_type in zip(file_names, data_names, data_types): + args.train_data_path_and_name_and_type.append( + [ + "{}/{}/{}".format(args.data_dir, args.train_set, file_name), + data_name, + data_type, + ] + ) + args.valid_data_path_and_name_and_type.append( + [ + "{}/{}/{}".format(args.data_dir, args.valid_set, file_name), + data_name, + data_type, + ] + ) + if os.path.exists(args.train_shape_file[0]): + assert os.path.exists(args.valid_shape_file[0]) + print("shape file for small dataset already exists.") + return + else: + concat_data_name = "_".join(data_names) + args.train_data_file = os.path.join( + args.data_dir, args.train_set, "{}_data.list".format(concat_data_name) + ) + args.valid_data_file = os.path.join( + args.data_dir, args.valid_set, "{}_data.list".format(concat_data_name) + ) + if os.path.exists(args.train_data_file): + assert os.path.exists(args.valid_data_file) + print("data list for large dataset already exists.") + return + + distributed = distributed_option.distributed + if not distributed or distributed_option.dist_rank == 0: + if hasattr(args, "filter_input") and args.filter_input: + filter_wav_text(args.data_dir, args.train_set) + filter_wav_text(args.data_dir, args.valid_set) + + if args.dataset_type == "small" and batch_type != "unsorted": + calc_shape(args, args.train_set) + calc_shape(args, args.valid_set) + + if args.dataset_type == "large": + generate_data_list(args, args.data_dir, args.train_set) + generate_data_list(args, args.data_dir, args.valid_set) + + if distributed: + dist.barrier() diff --git a/almeval/models/stepaudio/funasr_detach/utils/speaker_utils.py b/almeval/models/stepaudio/funasr_detach/utils/speaker_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7a7b27f9fcb735796c0df6533c17ae3613ddeb95 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/utils/speaker_utils.py @@ -0,0 +1,200 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +"""Some implementations are adapted from https://github.com/yuyq96/D-TDNN""" + +import io +from typing import Union + +import librosa as sf +import numpy as np +import torch +import torch.nn.functional as F +import torchaudio.compliance.kaldi as Kaldi +from torch import nn + +from funasr_detach.utils.modelscope_file import File + + +def check_audio_list(audio: list): + audio_dur = 0 + for i in range(len(audio)): + seg = audio[i] + assert seg[1] >= seg[0], "modelscope error: Wrong time stamps." + assert isinstance(seg[2], np.ndarray), "modelscope error: Wrong data type." + assert ( + int(seg[1] * 16000) - int(seg[0] * 16000) == seg[2].shape[0] + ), "modelscope error: audio data in list is inconsistent with time length." + if i > 0: + assert seg[0] >= audio[i - 1][1], "modelscope error: Wrong time stamps." + audio_dur += seg[1] - seg[0] + return audio_dur + # assert audio_dur > 5, 'modelscope error: The effective audio duration is too short.' + + +def sv_preprocess(inputs: Union[np.ndarray, list]): + output = [] + for i in range(len(inputs)): + if isinstance(inputs[i], str): + file_bytes = File.read(inputs[i]) + data, fs = sf.load(io.BytesIO(file_bytes), dtype="float32") + if len(data.shape) == 2: + data = data[:, 0] + data = torch.from_numpy(data).unsqueeze(0) + data = data.squeeze(0) + elif isinstance(inputs[i], np.ndarray): + assert ( + len(inputs[i].shape) == 1 + ), "modelscope error: Input array should be [N, T]" + data = inputs[i] + if data.dtype in ["int16", "int32", "int64"]: + data = (data / (1 << 15)).astype("float32") + else: + data = data.astype("float32") + data = torch.from_numpy(data) + else: + raise ValueError( + "modelscope error: The input type is restricted to audio address and nump array." + ) + output.append(data) + return output + + +def sv_chunk(vad_segments: list, fs=16000) -> list: + config = { + "seg_dur": 1.5, + "seg_shift": 0.75, + } + + def seg_chunk(seg_data): + seg_st = seg_data[0] + data = seg_data[2] + chunk_len = int(config["seg_dur"] * fs) + chunk_shift = int(config["seg_shift"] * fs) + last_chunk_ed = 0 + seg_res = [] + for chunk_st in range(0, data.shape[0], chunk_shift): + chunk_ed = min(chunk_st + chunk_len, data.shape[0]) + if chunk_ed <= last_chunk_ed: + break + last_chunk_ed = chunk_ed + chunk_st = max(0, chunk_ed - chunk_len) + chunk_data = data[chunk_st:chunk_ed] + if chunk_data.shape[0] < chunk_len: + chunk_data = np.pad( + chunk_data, (0, chunk_len - chunk_data.shape[0]), "constant" + ) + seg_res.append([chunk_st / fs + seg_st, chunk_ed / fs + seg_st, chunk_data]) + return seg_res + + segs = [] + for i, s in enumerate(vad_segments): + segs.extend(seg_chunk(s)) + + return segs + + +def extract_feature(audio): + features = [] + for au in audio: + feature = Kaldi.fbank(au.unsqueeze(0), num_mel_bins=80) + feature = feature - feature.mean(dim=0, keepdim=True) + features.append(feature.unsqueeze(0)) + features = torch.cat(features) + return features + + +def postprocess( + segments: list, vad_segments: list, labels: np.ndarray, embeddings: np.ndarray +) -> list: + assert len(segments) == len(labels) + labels = correct_labels(labels) + distribute_res = [] + for i in range(len(segments)): + distribute_res.append([segments[i][0], segments[i][1], labels[i]]) + # merge the same speakers chronologically + distribute_res = merge_seque(distribute_res) + + # accquire speaker center + spk_embs = [] + for i in range(labels.max() + 1): + spk_emb = embeddings[labels == i].mean(0) + spk_embs.append(spk_emb) + spk_embs = np.stack(spk_embs) + + def is_overlapped(t1, t2): + if t1 > t2 + 1e-4: + return True + return False + + # distribute the overlap region + for i in range(1, len(distribute_res)): + if is_overlapped(distribute_res[i - 1][1], distribute_res[i][0]): + p = (distribute_res[i][0] + distribute_res[i - 1][1]) / 2 + distribute_res[i][0] = p + distribute_res[i - 1][1] = p + + # smooth the result + distribute_res = smooth(distribute_res) + + return distribute_res + + +def correct_labels(labels): + labels_id = 0 + id2id = {} + new_labels = [] + for i in labels: + if i not in id2id: + id2id[i] = labels_id + labels_id += 1 + new_labels.append(id2id[i]) + return np.array(new_labels) + + +def merge_seque(distribute_res): + res = [distribute_res[0]] + for i in range(1, len(distribute_res)): + if distribute_res[i][2] != res[-1][2] or distribute_res[i][0] > res[-1][1]: + res.append(distribute_res[i]) + else: + res[-1][1] = distribute_res[i][1] + return res + + +def smooth(res, mindur=1): + # short segments are assigned to nearest speakers. + for i in range(len(res)): + res[i][0] = round(res[i][0], 2) + res[i][1] = round(res[i][1], 2) + if res[i][1] - res[i][0] < mindur: + if i == 0: + res[i][2] = res[i + 1][2] + elif i == len(res) - 1: + res[i][2] = res[i - 1][2] + elif res[i][0] - res[i - 1][1] <= res[i + 1][0] - res[i][1]: + res[i][2] = res[i - 1][2] + else: + res[i][2] = res[i + 1][2] + # merge the speakers + res = merge_seque(res) + + return res + + +def distribute_spk(sentence_list, sd_time_list): + sd_sentence_list = [] + for d in sentence_list: + sentence_start = d["ts_list"][0][0] + sentence_end = d["ts_list"][-1][1] + sentence_spk = 0 + max_overlap = 0 + for sd_time in sd_time_list: + spk_st, spk_ed, spk = sd_time + spk_st = spk_st * 1000 + spk_ed = spk_ed * 1000 + overlap = max(min(sentence_end, spk_ed) - max(sentence_start, spk_st), 0) + if overlap > max_overlap: + max_overlap = overlap + sentence_spk = spk + d["spk"] = sentence_spk + sd_sentence_list.append(d) + return sd_sentence_list diff --git a/almeval/models/stepaudio/funasr_detach/utils/timestamp_tools.py b/almeval/models/stepaudio/funasr_detach/utils/timestamp_tools.py new file mode 100644 index 0000000000000000000000000000000000000000..6bc8cb1f8c4d01c5a4f8059d06303d053d8738d9 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/utils/timestamp_tools.py @@ -0,0 +1,201 @@ +import torch +import codecs +import logging +import argparse +import numpy as np + +# import edit_distance +from itertools import zip_longest + + +def cif_wo_hidden(alphas, threshold): + batch_size, len_time = alphas.size() + # loop varss + integrate = torch.zeros([batch_size], device=alphas.device) + # intermediate vars along time + list_fires = [] + for t in range(len_time): + alpha = alphas[:, t] + integrate += alpha + list_fires.append(integrate) + fire_place = integrate >= threshold + integrate = torch.where( + fire_place, + integrate - torch.ones([batch_size], device=alphas.device) * threshold, + integrate, + ) + fires = torch.stack(list_fires, 1) + return fires + + +def ts_prediction_lfr6_standard( + us_alphas, + us_peaks, + char_list, + vad_offset=0.0, + force_time_shift=-1.5, + sil_in_str=True, +): + if not len(char_list): + return "", [] + START_END_THRESHOLD = 5 + MAX_TOKEN_DURATION = 12 + TIME_RATE = 10.0 * 6 / 1000 / 3 # 3 times upsampled + if len(us_alphas.shape) == 2: + alphas, peaks = us_alphas[0], us_peaks[0] # support inference batch_size=1 only + else: + alphas, peaks = us_alphas, us_peaks + if char_list[-1] == "": + char_list = char_list[:-1] + fire_place = ( + torch.where(peaks > 1.0 - 1e-4)[0].cpu().numpy() + force_time_shift + ) # total offset + if len(fire_place) != len(char_list) + 1: + alphas /= alphas.sum() / (len(char_list) + 1) + alphas = alphas.unsqueeze(0) + peaks = cif_wo_hidden(alphas, threshold=1.0 - 1e-4)[0] + fire_place = ( + torch.where(peaks > 1.0 - 1e-4)[0].cpu().numpy() + force_time_shift + ) # total offset + num_frames = peaks.shape[0] + timestamp_list = [] + new_char_list = [] + # for bicif model trained with large data, cif2 actually fires when a character starts + # so treat the frames between two peaks as the duration of the former token + fire_place = ( + torch.where(peaks > 1.0 - 1e-4)[0].cpu().numpy() + force_time_shift + ) # total offset + # assert num_peak == len(char_list) + 1 # number of peaks is supposed to be number of tokens + 1 + # begin silence + if fire_place[0] > START_END_THRESHOLD: + # char_list.insert(0, '') + timestamp_list.append([0.0, fire_place[0] * TIME_RATE]) + new_char_list.append("") + # tokens timestamp + for i in range(len(fire_place) - 1): + new_char_list.append(char_list[i]) + if ( + MAX_TOKEN_DURATION < 0 + or fire_place[i + 1] - fire_place[i] <= MAX_TOKEN_DURATION + ): + timestamp_list.append( + [fire_place[i] * TIME_RATE, fire_place[i + 1] * TIME_RATE] + ) + else: + # cut the duration to token and sil of the 0-weight frames last long + _split = fire_place[i] + MAX_TOKEN_DURATION + timestamp_list.append([fire_place[i] * TIME_RATE, _split * TIME_RATE]) + timestamp_list.append([_split * TIME_RATE, fire_place[i + 1] * TIME_RATE]) + new_char_list.append("") + # tail token and end silence + # new_char_list.append(char_list[-1]) + if num_frames - fire_place[-1] > START_END_THRESHOLD: + _end = (num_frames + fire_place[-1]) * 0.5 + # _end = fire_place[-1] + timestamp_list[-1][1] = _end * TIME_RATE + timestamp_list.append([_end * TIME_RATE, num_frames * TIME_RATE]) + new_char_list.append("") + else: + timestamp_list[-1][1] = num_frames * TIME_RATE + if vad_offset: # add offset time in model with vad + for i in range(len(timestamp_list)): + timestamp_list[i][0] = timestamp_list[i][0] + vad_offset / 1000.0 + timestamp_list[i][1] = timestamp_list[i][1] + vad_offset / 1000.0 + res_txt = "" + for char, timestamp in zip(new_char_list, timestamp_list): + # if char != '': + if not sil_in_str and char == "": + continue + res_txt += "{} {} {};".format( + char, str(timestamp[0] + 0.0005)[:5], str(timestamp[1] + 0.0005)[:5] + ) + res = [] + for char, timestamp in zip(new_char_list, timestamp_list): + if char != "": + res.append([int(timestamp[0] * 1000), int(timestamp[1] * 1000)]) + return res_txt, res + + +def timestamp_sentence( + punc_id_list, timestamp_postprocessed, text_postprocessed, return_raw_text=False +): + punc_list = ["，", "。", "？", "、"] + res = [] + if text_postprocessed is None: + return res + if timestamp_postprocessed is None: + return res + if len(timestamp_postprocessed) == 0: + return res + if len(text_postprocessed) == 0: + return res + + if punc_id_list is None or len(punc_id_list) == 0: + res.append( + { + "text": text_postprocessed.split(), + "start": timestamp_postprocessed[0][0], + "end": timestamp_postprocessed[-1][1], + "timestamp": timestamp_postprocessed, + } + ) + return res + if len(punc_id_list) != len(timestamp_postprocessed): + logging.warning("length mismatch between punc and timestamp") + sentence_text = "" + sentence_text_seg = "" + ts_list = [] + sentence_start = timestamp_postprocessed[0][0] + sentence_end = timestamp_postprocessed[0][1] + texts = text_postprocessed.split() + punc_stamp_text_list = list( + zip_longest(punc_id_list, timestamp_postprocessed, texts, fillvalue=None) + ) + for punc_stamp_text in punc_stamp_text_list: + punc_id, timestamp, text = punc_stamp_text + # sentence_text += text if text is not None else '' + if text is not None: + if "a" <= text[0] <= "z" or "A" <= text[0] <= "Z": + sentence_text += " " + text + elif len(sentence_text) and ( + "a" <= sentence_text[-1] <= "z" or "A" <= sentence_text[-1] <= "Z" + ): + sentence_text += " " + text + else: + sentence_text += text + sentence_text_seg += text + " " + ts_list.append(timestamp) + + punc_id = int(punc_id) if punc_id is not None else 1 + sentence_end = timestamp[1] if timestamp is not None else sentence_end + sentence_text_seg = ( + sentence_text_seg[:-1] + if sentence_text_seg[-1] == " " + else sentence_text_seg + ) + if punc_id > 1: + sentence_text += punc_list[punc_id - 2] + if return_raw_text: + res.append( + { + "text": sentence_text, + "start": sentence_start, + "end": sentence_end, + "timestamp": ts_list, + "raw_text": sentence_text_seg, + } + ) + else: + res.append( + { + "text": sentence_text, + "start": sentence_start, + "end": sentence_end, + "timestamp": ts_list, + } + ) + sentence_text = "" + sentence_text_seg = "" + ts_list = [] + sentence_start = sentence_end + return res diff --git a/almeval/models/stepaudio/funasr_detach/utils/types.py b/almeval/models/stepaudio/funasr_detach/utils/types.py new file mode 100644 index 0000000000000000000000000000000000000000..6b36f9c4b87ed9258a5d1e254ba298ed5dbc01d2 --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/utils/types.py @@ -0,0 +1,149 @@ +from distutils.util import strtobool +from typing import Optional +from typing import Tuple +from typing import Union + +import humanfriendly + + +def str2bool(value: str) -> bool: + return bool(strtobool(value)) + + +def remove_parenthesis(value: str): + value = value.strip() + if value.startswith("(") and value.endswith(")"): + value = value[1:-1] + elif value.startswith("[") and value.endswith("]"): + value = value[1:-1] + return value + + +def remove_quotes(value: str): + value = value.strip() + if value.startswith('"') and value.endswith('"'): + value = value[1:-1] + elif value.startswith("'") and value.endswith("'"): + value = value[1:-1] + return value + + +def int_or_none(value: str) -> Optional[int]: + """int_or_none. + + Examples: + >>> import argparse + >>> parser = argparse.ArgumentParser() + >>> _ = parser.add_argument('--foo', type=int_or_none) + >>> parser.parse_args(['--foo', '456']) + Namespace(foo=456) + >>> parser.parse_args(['--foo', 'none']) + Namespace(foo=None) + >>> parser.parse_args(['--foo', 'null']) + Namespace(foo=None) + >>> parser.parse_args(['--foo', 'nil']) + Namespace(foo=None) + + """ + if value.strip().lower() in ("none", "null", "nil"): + return None + return int(value) + + +def float_or_none(value: str) -> Optional[float]: + """float_or_none. + + Examples: + >>> import argparse + >>> parser = argparse.ArgumentParser() + >>> _ = parser.add_argument('--foo', type=float_or_none) + >>> parser.parse_args(['--foo', '4.5']) + Namespace(foo=4.5) + >>> parser.parse_args(['--foo', 'none']) + Namespace(foo=None) + >>> parser.parse_args(['--foo', 'null']) + Namespace(foo=None) + >>> parser.parse_args(['--foo', 'nil']) + Namespace(foo=None) + + """ + if value.strip().lower() in ("none", "null", "nil"): + return None + return float(value) + + +def humanfriendly_parse_size_or_none(value) -> Optional[float]: + if value.strip().lower() in ("none", "null", "nil"): + return None + return humanfriendly.parse_size(value) + + +def str_or_int(value: str) -> Union[str, int]: + try: + return int(value) + except ValueError: + return value + + +def str_or_none(value: str) -> Optional[str]: + """str_or_none. + + Examples: + >>> import argparse + >>> parser = argparse.ArgumentParser() + >>> _ = parser.add_argument('--foo', type=str_or_none) + >>> parser.parse_args(['--foo', 'aaa']) + Namespace(foo='aaa') + >>> parser.parse_args(['--foo', 'none']) + Namespace(foo=None) + >>> parser.parse_args(['--foo', 'null']) + Namespace(foo=None) + >>> parser.parse_args(['--foo', 'nil']) + Namespace(foo=None) + + """ + if value.strip().lower() in ("none", "null", "nil"): + return None + return value + + +def str2pair_str(value: str) -> Tuple[str, str]: + """str2pair_str. + + Examples: + >>> import argparse + >>> str2pair_str('abc,def ') + ('abc', 'def') + >>> parser = argparse.ArgumentParser() + >>> _ = parser.add_argument('--foo', type=str2pair_str) + >>> parser.parse_args(['--foo', 'abc,def']) + Namespace(foo=('abc', 'def')) + + """ + value = remove_parenthesis(value) + a, b = value.split(",") + + # Workaround for configargparse issues: + # If the list values are given from yaml file, + # the value givent to type() is shaped as python-list, + # e.g. ['a', 'b', 'c'], + # so we need to remove double quotes from it. + return remove_quotes(a), remove_quotes(b) + + +def str2triple_str(value: str) -> Tuple[str, str, str]: + """str2triple_str. + + Examples: + >>> str2triple_str('abc,def ,ghi') + ('abc', 'def', 'ghi') + """ + value = remove_parenthesis(value) + a, b, c = value.split(",") + + # Workaround for configargparse issues: + # If the list values are given from yaml file, + # the value givent to type() is shaped as python-list, + # e.g. ['a', 'b', 'c'], + # so we need to remove quotes from it. + return remove_quotes(a), remove_quotes(b), remove_quotes(c) diff --git a/almeval/models/stepaudio/funasr_detach/utils/vad_utils.py b/almeval/models/stepaudio/funasr_detach/utils/vad_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..eb9cb6f9a9cb7585cd466f764ac340e73c27fb1b --- /dev/null +++ b/almeval/models/stepaudio/funasr_detach/utils/vad_utils.py @@ -0,0 +1,32 @@ +import torch +from torch.nn.utils.rnn import pad_sequence + + +def slice_padding_fbank(speech, speech_lengths, vad_segments): + speech_list = [] + speech_lengths_list = [] + for i, segment in enumerate(vad_segments): + + bed_idx = int(segment[0][0] * 16) + end_idx = min(int(segment[0][1] * 16), speech_lengths[0]) + speech_i = speech[0, bed_idx:end_idx] + speech_lengths_i = end_idx - bed_idx + speech_list.append(speech_i) + speech_lengths_list.append(speech_lengths_i) + feats_pad = pad_sequence(speech_list, batch_first=True, padding_value=0.0) + speech_lengths_pad = torch.Tensor(speech_lengths_list).int() + return feats_pad, speech_lengths_pad + + +def slice_padding_audio_samples(speech, speech_lengths, vad_segments): + speech_list = [] + speech_lengths_list = [] + for i, segment in enumerate(vad_segments): + bed_idx = int(segment[0][0] * 16) + end_idx = min(int(segment[0][1] * 16), speech_lengths) + speech_i = speech[bed_idx:end_idx] + speech_lengths_i = end_idx - bed_idx + speech_list.append(speech_i) + speech_lengths_list.append(speech_lengths_i) + + return speech_list, speech_lengths_list diff --git a/almeval/utils/__init__.py b/almeval/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6caf39a5dfcbcd2a4c08bf795bacd60337ec69ba --- /dev/null +++ b/almeval/utils/__init__.py @@ -0,0 +1 @@ +from .misc import * # noqa diff --git a/almeval/utils/config_manager.py b/almeval/utils/config_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..2827f0fbbfc4f5612c10bddb25dcf4909e25fc04 --- /dev/null +++ b/almeval/utils/config_manager.py @@ -0,0 +1,53 @@ +import os +from pathlib import Path +from typing import Any + +import yaml + + +class ConfigManager: + _instance = None + _config = None + + def __new__(cls): + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def __init__(self): + if self._config is None: + self.load_config() + + def _find_config_file(self) -> Path: + # 1. 首先检查环境变量 + config_path = os.getenv('PROJECT_CONFIG_PATH') + if config_path and Path(config_path).exists(): + return Path(config_path) + + # 2. 从当前目录向上查找，直到找到配置文件 + current_dir = Path(__file__).resolve().parent + while current_dir != current_dir.parent: + config_file = current_dir.parent / 'config.yaml' + if config_file.exists(): + return config_file + current_dir = current_dir.parent + + raise FileNotFoundError( + 'cannot find config.yaml. Please set PROJECT_CONFIG_PATH environment variable or put config.yaml in the project root directory.') + + def load_config(self): + config_path = self._find_config_file() + with open(config_path) as f: + self._config = yaml.safe_load(f) + + @property + def config(self) -> dict[str, Any]: + return self._config + + def get_dataset_root(self) -> str: + """get dataset root path, your dataset file should be in this path""" + return self._config['DATASETS']['dataset_root'] + + def get_dataset_path(self, dataset_name: str) -> str | None: + """get dataset file path""" + return self._config['DATASETS']['datasets'].get(dataset_name, None) \ No newline at end of file diff --git a/almeval/utils/misc.py b/almeval/utils/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..06bea5dc9aa7aad73315844bb5e98e7a6f783b08 --- /dev/null +++ b/almeval/utils/misc.py @@ -0,0 +1,173 @@ +# flake8: noqa: F401, F403 +import csv +import hashlib +import json +import mimetypes +import os +import os.path as osp +import pickle +import subprocess + +import numpy as np +import pandas as pd +import validators +from loguru import logger + + +def download_file(url, filename=None): + import urllib.request + + from tqdm import tqdm + + class DownloadProgressBar(tqdm): + def update_to(self, b=1, bsize=1, tsize=None): + if tsize is not None: + self.total = tsize + self.update(b * bsize - self.n) + + if filename is None: + filename = url.split('/')[-1] + + try: + with DownloadProgressBar(unit='B', unit_scale=True, miniters=1, desc=url.split('/')[-1]) as t: + urllib.request.urlretrieve( + url, filename=filename, reporthook=t.update_to) + except Exception: + # Handle Failed Downloads from huggingface.co + if 'huggingface.co' in url: + url_new = url.replace('huggingface.co', 'hf-mirror.com') + try: + os.system(f'wget {url_new} -O {filename}') + except Exception: + raise Exception(f'Failed to download {url}') + else: + raise Exception(f'Failed to download {url}') + + return filename + + +AUDIO_TYPES = {'mp3', 'ogg', 'wav', 'flac', 'm4a', 'wma', 'aac'} + + +def md5(s): + hash = hashlib.new('md5') + if osp.exists(s): + with open(s, 'rb') as f: + for chunk in iter(lambda: f.read(2**20), b''): + hash.update(chunk) + else: + hash.update(s.encode('utf-8')) + return str(hash.hexdigest()) + + +def parse_file(s): + if isinstance(s, str) and osp.exists(s) and s != '.': + assert osp.isfile(s) + suffix = osp.splitext(s)[1].lower() + if suffix in AUDIO_TYPES: + mime = 'audio' + else: + mime = mimetypes.types_map.get(suffix, 'unknown') + return (mime, s) + elif validators.url(s): + suffix = osp.splitext(s)[1].lower() + if suffix in AUDIO_TYPES: + mime = 'audio' + elif suffix in mimetypes.types_map: + mime = mimetypes.types_map[suffix] + + return (mime, s) + else: + return (None, s) + + +class NumpyEncoder(json.JSONEncoder): + def default(self, obj): + if isinstance(obj, (np.int_, np.intc, np.intp, np.int8, + np.int16, np.int32, np.int64, np.uint8, + np.uint16, np.uint32, np.uint64)): + return int(obj) + elif isinstance(obj, (np.float_, np.float16, np.float32, np.float64)): + return float(obj) + elif isinstance(obj, (np.complex_, np.complex64, np.complex128)): + return {'real': obj.real, 'imag': obj.imag} + elif isinstance(obj, (np.ndarray,)): + return obj.tolist() + elif isinstance(obj, (np.bool_)): + return bool(obj) + elif isinstance(obj, (np.void)): + return None + return json.JSONEncoder.default(self, obj) + + +# LOAD & DUMP +def dump(data, f, **kwargs): + def dump_pkl(data, pth, **kwargs): + pickle.dump(data, open(pth, 'wb')) + + def dump_json(data, pth, **kwargs): + json.dump(data, open(pth, 'w'), indent=4, + ensure_ascii=False, cls=NumpyEncoder) + + def dump_jsonl(data, f, **kwargs): + lines = [json.dumps(x, ensure_ascii=False, cls=NumpyEncoder) + for x in data] + with open(f, 'w', encoding='utf8') as fout: + fout.write('\n'.join(lines)) + + def dump_xlsx(data, f, **kwargs): + data.to_excel(f, index=False, engine='xlsxwriter', engine_kwargs={'options': {'strings_to_urls': False, + 'strings_to_formulas': False}}) + + def dump_csv(data, f, quoting=csv.QUOTE_ALL): + data.to_csv(f, index=False, encoding='utf-8', quoting=quoting) + + def dump_tsv(data, f, quoting=csv.QUOTE_ALL): + data.to_csv(f, sep='\t', index=False, + encoding='utf-8', quoting=quoting) + + handlers = dict(pkl=dump_pkl, json=dump_json, jsonl=dump_jsonl, + xlsx=dump_xlsx, csv=dump_csv, tsv=dump_tsv) + suffix = f.split('.')[-1] + return handlers[suffix](data, f, **kwargs) + + +def load(f): + def load_pkl(pth): + return pickle.load(open(pth, 'rb')) + + def load_json(pth): + return json.load(open(pth, encoding='utf-8')) + + def load_jsonl(f): + lines = open(f, encoding='utf-8').readlines() + lines = [x.strip() for x in lines] + if lines[-1] == '': + lines = lines[:-1] + data = [json.loads(x) for x in lines] + return data + + def load_xlsx(f): + return pd.read_excel(f) + + def load_csv(f): + return pd.read_csv(f) + + handlers = dict(pkl=load_pkl, json=load_json, + jsonl=load_jsonl, xlsx=load_xlsx, csv=load_csv) + suffix = f.split('.')[-1] + return handlers[suffix](f) + + +def run_command(cmd): + if isinstance(cmd, str): + cmd = cmd.split() + return subprocess.check_output(cmd).decode() + + +def print_once(msg): + if not hasattr(print_once, 'printed'): + print_once.printed = set() + if msg not in print_once.printed: + print_once.printed.add(msg) + logger.info(msg)