Initial commit: FunASR Speech Recognition Toolkit

Add complete FunASR codebase including models, runtime, and documentation.

funasr/models/eend/e2e_diar_eend_ola.py

@@ -0,0 +1,343 @@
+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.frontends.wav_frontend import WavFrontendMel23
+from funasr.models.eend.encoder import EENDOLATransformerEncoder
+from funasr.models.eend.encoder_decoder_attractor import EncoderDecoderAttractor
+from funasr.models.eend.utils.losses import (
+    standard_loss,
+    cal_power_loss,
+    fast_batch_pit_n_speaker_loss,
+)
+from funasr.models.eend.utils.power import create_powerlabel
+from funasr.models.eend.utils.power import generate_mapping_dict
+from funasr.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):
+        """Autocast.
+        
+            Args:
+                enabled: TODO.
+            """
+        yield
+
+
+def pad_attractor(att, max_n_speakers):
+    """Pad attractor.
+    
+        Args:
+            att: TODO.
+            max_n_speakers: TODO.
+        """
+    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):
+    """Pad labels.
+    
+        Args:
+            ts: TODO.
+            out_size: Size/dimension parameter.
+        """
+    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):
+    """Pad results.
+    
+        Args:
+            ys: TODO.
+            out_size: Size/dimension parameter.
+        """
+    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,
+    ):
+        """Initialize DiarEENDOLAModel.
+        
+            Args:
+                frontend: Audio frontend for feature extraction.
+                encoder: TODO.
+                encoder_decoder_attractor: TODO.
+                n_units: TODO.
+                max_n_speaker: TODO.
+                attractor_loss_weight: TODO.
+                mapping_dict: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        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):
+        """Forward encoder.
+        
+            Args:
+                xs: TODO.
+                ilens: TODO.
+            """
+        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):
+        """Forward post net.
+        
+            Args:
+                logits: TODO.
+                ilens: TODO.
+            """
+        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
+        """Forward pass for training.
+        
+            Args:
+                speech: Speech audio tensor, shape (batch, time).
+                speaker_labels: TODO.
+                orders: TODO.
+            """
+        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,
+    ):
+        """Estimate sequential.
+        
+            Args:
+                speech: Speech audio tensor, shape (batch, time).
+                n_speakers: TODO.
+                shuffle: TODO.
+                threshold: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        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):
+        """Recover y from powerlabel.
+        
+            Args:
+                logit: TODO.
+                n_speaker: TODO.
+            """
+        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):
+
+        """Inv mapping func.
+        
+            Args:
+                label: TODO.
+            """
+        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]:
+        """Collect feats.
+        
+            Args:
+                **batch: Additional keyword arguments.
+            """
+        pass

funasr/models/eend/eend_ola_dataloader.py

@@ -0,0 +1,86 @@
+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):
+    """Custom collate.
+    
+        Args:
+            batch: TODO.
+        """
+    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,
+    ):
+        """Initialize EENDOLADataset.
+        
+            Args:
+                data_file: TODO.
+            """
+        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):
+        """Internal: len  ."""
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        """Internal: getitem  .
+        
+            Args:
+                idx: TODO.
+            """
+        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):
+        """Initialize EENDOLADataLoader.
+        
+            Args:
+                data_file: TODO.
+                batch_size: Number of samples per batch.
+                shuffle: TODO.
+                num_workers: TODO.
+            """
+        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):
+        """Build iter.
+        
+            Args:
+                epoch: TODO.
+            """
+        return self.data_loader

funasr/models/eend/encoder.py

@@ -0,0 +1,179 @@
+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):
+        """Initialize MultiHeadSelfAttention.
+        
+            Args:
+                n_units: TODO.
+                h: TODO.
+                dropout_rate: TODO.
+            """
+        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):
+        """Internal: call  .
+        
+            Args:
+                x: TODO.
+                batch_size: Number of samples per batch.
+                x_mask: TODO.
+            """
+        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):
+        """Initialize PositionwiseFeedForward.
+        
+            Args:
+                n_units: TODO.
+                d_units: TODO.
+                dropout_rate: TODO.
+            """
+        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):
+        """Internal: call  .
+        
+            Args:
+                x: TODO.
+            """
+        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):
+        """Initialize PositionalEncoding.
+        
+            Args:
+                d_model: D Model instance.
+                dropout_rate: TODO.
+                max_len: TODO.
+                reverse: TODO.
+            """
+        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):
+        """Extend pe.
+        
+            Args:
+                x: TODO.
+            """
+        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):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        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,
+    ):
+        """Initialize EENDOLATransformerEncoder.
+        
+            Args:
+                idim: TODO.
+                n_layers: TODO.
+                n_units: TODO.
+                e_units: TODO.
+                h: TODO.
+                dropout_rate: TODO.
+                use_pos_emb: TODO.
+            """
+        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):
+        """Internal: call  .
+        
+            Args:
+                x: TODO.
+                x_mask: TODO.
+            """
+        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)

funasr/models/eend/encoder_decoder_attractor.py

@@ -0,0 +1,89 @@
+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):
+        """Initialize EncoderDecoderAttractor.
+        
+            Args:
+                n_units: TODO.
+                encoder_dropout: TODO.
+                decoder_dropout: TODO.
+            """
+        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):
+        """Forward core.
+        
+            Args:
+                xs: TODO.
+                zeros: TODO.
+            """
+        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):
+        """Forward pass for training.
+        
+            Args:
+                xs: TODO.
+                n_speakers: TODO.
+            """
+        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):
+        """Estimate.
+        
+            Args:
+                xs: TODO.
+                max_n_speakers: TODO.
+            """
+        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

funasr/models/eend/utils/feature.py

@@ -0,0 +1,276 @@
+# 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,
+):
+    """Get input dim.
+    
+        Args:
+            frame_size: Size/dimension parameter.
+            context_size: Size/dimension parameter.
+            transform_type: TODO.
+        """
+    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)
+    """Internal: count frames.
+    
+        Args:
+            data_len: TODO.
+            size: TODO.
+            shift: TODO.
+        """
+    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

funasr/models/eend/utils/kaldi_data.py

@@ -0,0 +1,174 @@
+# 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):
+    """Load segments hash.
+    
+        Args:
+            segments_file: TODO.
+        """
+    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):
+    """Load segments rechash.
+    
+        Args:
+            segments_file: TODO.
+        """
+    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):
+        """Initialize KaldiData.
+        
+            Args:
+                data_dir: TODO.
+            """
+        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):
+        """Load wav.
+        
+            Args:
+                recid: TODO.
+                start: TODO.
+                end: TODO.
+            """
+        data, rate = load_wav(self.wavs[recid], start, end)
+        return data, rate

funasr/models/eend/utils/losses.py

@@ -0,0 +1,67 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from scipy.optimize import linear_sum_assignment
+
+
+def standard_loss(ys, ts):
+    """Standard loss.
+    
+        Args:
+            ys: TODO.
+            ts: TODO.
+        """
+    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):
+    """Fast batch pit n speaker loss.
+    
+        Args:
+            ys: TODO.
+            ts: TODO.
+        """
+    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):
+    """Cal power loss.
+    
+        Args:
+            logits: TODO.
+            power_ts: TODO.
+        """
+    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

funasr/models/eend/utils/power.py

@@ -0,0 +1,156 @@
+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):
+    """Generate mapping dict.
+    
+        Args:
+            max_speaker_num: TODO.
+            max_olp_speaker_num: TODO.
+        """
+    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):
+    """Raw dec trans.
+    
+        Args:
+            x: TODO.
+            max_speaker_num: TODO.
+        """
+    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):
+    """Mapping func.
+    
+        Args:
+            num: TODO.
+            mapping_dict: TODO.
+        """
+    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):
+    """Dec trans.
+    
+        Args:
+            x: TODO.
+            max_speaker_num: TODO.
+            mapping_dict: TODO.
+        """
+    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):
+    """Create powerlabel.
+    
+        Args:
+            label: TODO.
+            mapping_dict: TODO.
+            max_speaker_num: TODO.
+            max_olp_speaker_num: TODO.
+        """
+    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):
+    """Generate perm pse.
+    
+        Args:
+            label: TODO.
+            n_speaker: TODO.
+            mapping_dict: TODO.
+            max_speaker_num: TODO.
+            max_olp_speaker_num: TODO.
+        """
+    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
+):
+    """Generate min pse.
+    
+        Args:
+            label: TODO.
+            n_speaker: TODO.
+            mapping_dict: TODO.
+            max_speaker_num: TODO.
+            pse_logit: TODO.
+            max_olp_speaker_num: TODO.
+        """
+    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

funasr/models/eend/utils/report.py

@@ -0,0 +1,229 @@
+import copy
+import numpy as np
+import time
+import torch
+from funasr.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):
+    """Recover prediction.
+    
+        Args:
+            y: TODO.
+            n_speaker: TODO.
+        """
+    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):
+        """Initialize PowerReporter.
+        
+            Args:
+                valid_data_loader: TODO.
+                mapping_dict: TODO.
+                max_n_speaker: TODO.
+            """
+        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):
+        """Report.
+        
+            Args:
+                model: Model instance or model name.
+                eidx: TODO.
+                device: Target device ("cuda:0", "cpu", etc.).
+            """
+        self.report_val(model, eidx, device)
+
+    def report_val(self, model, eidx, device):
+        """Report val.
+        
+            Args:
+                model: Model instance or model name.
+                eidx: TODO.
+                device: Target device ("cuda:0", "cpu", etc.).
+            """
+        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):
+        """Inv mapping func.
+        
+            Args:
+                label: TODO.
+                mapping_dict: TODO.
+            """
+        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):
+        """Report core.
+        
+            Args:
+                model: Model instance or model name.
+                data_loader: TODO.
+                device: Target device ("cuda:0", "cpu", etc.).
+            """
+        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):
+        """Calc diarization error.
+        
+            Args:
+                decisions: TODO.
+                label: TODO.
+                label_delay: TODO.
+            """
+        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