Initial commit: FunASR Speech Recognition Toolkit

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

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