| | |
|---|
| | | import torch.nn as nn |
|---|
| | | from typeguard import check_argument_types |
|---|
| | | |
|---|
| | | from funasr.modules.eend_ola.encoder import TransformerEncoder |
|---|
| | | from funasr.models.frontend.wav_frontend import WavFrontendMel23 |
|---|
| | | from funasr.modules.eend_ola.encoder import EENDOLATransformerEncoder |
|---|
| | | from funasr.modules.eend_ola.encoder_decoder_attractor import EncoderDecoderAttractor |
|---|
| | | from funasr.modules.eend_ola.utils.power import generate_mapping_dict |
|---|
| | | from funasr.models.base_model import FunASRModel |
|---|
| | | from funasr.torch_utils.device_funcs import force_gatherable |
|---|
| | | from funasr.train.abs_espnet_model import AbsESPnetModel |
|---|
| | | |
|---|
| | | if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): |
|---|
| | | pass |
|---|
| | |
|---|
| | | return att |
|---|
| | | |
|---|
| | | |
|---|
| | | class DiarEENDOLAModel(AbsESPnetModel): |
|---|
| | | """CTC-attention hybrid Encoder-Decoder model""" |
|---|
| | | class DiarEENDOLAModel(FunASRModel): |
|---|
| | | """EEND-OLA diarization model""" |
|---|
| | | |
|---|
| | | def __init__( |
|---|
| | | self, |
|---|
| | | encoder: TransformerEncoder, |
|---|
| | | eda: EncoderDecoderAttractor, |
|---|
| | | frontend: WavFrontendMel23, |
|---|
| | | encoder: EENDOLATransformerEncoder, |
|---|
| | | encoder_decoder_attractor: EncoderDecoderAttractor, |
|---|
| | | n_units: int = 256, |
|---|
| | | max_n_speaker: int = 8, |
|---|
| | | attractor_loss_weight: float = 1.0, |
|---|
| | |
|---|
| | | assert check_argument_types() |
|---|
| | | |
|---|
| | | super().__init__() |
|---|
| | | self.encoder = encoder |
|---|
| | | self.eda = eda |
|---|
| | | self.frontend = frontend |
|---|
| | | self.enc = encoder |
|---|
| | | self.eda = 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.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): |
|---|
| | |
|---|
| | | 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.encoder(xs, xs_mask) |
|---|
| | | 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, batch_first=True, enforce_sorted=False) |
|---|
| | | outputs, (_, _) = self.PostNet(logits) |
|---|
| | | 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] |
|---|
| | |
|---|
| | | text = text[:, : text_lengths.max()] |
|---|
| | | |
|---|
| | | # 1. Encoder |
|---|
| | | encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) |
|---|
| | | encoder_out, encoder_out_lens = self.enc(speech, speech_lengths) |
|---|
| | | intermediate_outs = None |
|---|
| | | if isinstance(encoder_out, tuple): |
|---|
| | | intermediate_outs = encoder_out[1] |
|---|
| | |
|---|
| | | pred[i] = pred[i - 1] |
|---|
| | | else: |
|---|
| | | pred[i] = 0 |
|---|
| | | pred = [self.reporter.inv_mapping_func(i, self.mapping_dict) for i in pred] |
|---|
| | | 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 |
|---|
