Merge pull request #165 from alibaba-damo-academy/dev_cmz

punctuation:add training code, support largedataset

Author: LauraGPT

funasr/bin/punc_train.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+import os
+from funasr.tasks.punctuation import PunctuationTask
+
+
+def parse_args():
+    parser = PunctuationTask.get_parser()
+    parser.add_argument(
+        "--gpu_id",
+        type=int,
+        default=0,
+        help="local gpu id.",
+    )
+    parser.add_argument(
+        "--punc_list",
+        type=str,
+        default=None,
+        help="Punctuation list",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main(args=None, cmd=None):
+    """
+    punc training.
+    """
+    PunctuationTask.main(args=args, cmd=cmd)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    # setup local gpu_id
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_id)
+
+    # DDP settings
+    if args.ngpu > 1:
+        args.distributed = True
+    else:
+        args.distributed = False
+
+    main(args=args)

funasr/bin/punc_train_vadrealtime.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+import os
+from funasr.tasks.punctuation import PunctuationTask
+
+
+def parse_args():
+    parser = PunctuationTask.get_parser()
+    parser.add_argument(
+        "--gpu_id",
+        type=int,
+        default=0,
+        help="local gpu id.",
+    )
+    parser.add_argument(
+        "--punc_list",
+        type=str,
+        default=None,
+        help="Punctuation list",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main(args=None, cmd=None):
+    """
+    punc training.
+    """
+    PunctuationTask.main(args=args, cmd=cmd)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    # setup local gpu_id
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_id)
+
+    # DDP settings
+    if args.ngpu > 1:
+        args.distributed = True
+    else:
+        args.distributed = False
+    assert args.num_worker_count == 1
+
+    main(args=args)

funasr/datasets/large_datasets/build_dataloader.py

@@ -34,16 +34,20 @@ def load_seg_dict(seg_dict_file):
     return seg_dict
 
 class ArkDataLoader(AbsIterFactory):
-    def __init__(self, data_list, dict_file, dataset_conf, seg_dict_file=None, mode="train"):
+    def __init__(self, data_list, dict_file, dataset_conf, seg_dict_file=None, punc_dict_file=None, mode="train"):
         symbol_table = read_symbol_table(dict_file) if dict_file is not None else None
         if seg_dict_file is not None:
             seg_dict = load_seg_dict(seg_dict_file)
         else:
             seg_dict = None
+        if punc_dict_file is not None:
+            punc_dict = read_symbol_table(punc_dict_file)
+        else:
+            punc_dict = None
         self.dataset_conf = dataset_conf
         logging.info("dataloader config: {}".format(self.dataset_conf))
         batch_mode = self.dataset_conf.get("batch_mode", "padding")
-        self.dataset = Dataset(data_list, symbol_table, seg_dict,
+        self.dataset = Dataset(data_list, symbol_table, seg_dict, punc_dict,
                                self.dataset_conf, mode=mode, batch_mode=batch_mode)
 
     def build_iter(self, epoch, shuffle=True):

funasr/datasets/large_datasets/dataset.py

@@ -127,14 +127,17 @@ def __iter__(self):
                             sample_dict["key"] = key
                     else:
                         text = item
-                        sample_dict[data_name] = text.strip().split()[1:]
+                        segs = text.strip().split()
+                        sample_dict[data_name] = segs[1:]
+                        if "key" not in sample_dict:
+                            sample_dict["key"] = segs[0]
                 yield sample_dict
 
             self.close_reader(reader_list)
 
 
 def len_fn_example(data):
-    return len(data)
+    return 1
 
 
 def len_fn_token(data):
@@ -148,6 +151,7 @@ def len_fn_token(data):
 def Dataset(data_list_file,
             dict,
             seg_dict,
+            punc_dict,
             conf,
             mode="train",
             batch_mode="padding"):
@@ -162,7 +166,7 @@ def Dataset(data_list_file,
     dataset = FilterIterDataPipe(dataset, fn=filter_fn)
 
     if "text" in data_names:
-        vocab = {'vocab': dict, 'seg_dict': seg_dict}
+        vocab = {'vocab': dict, 'seg_dict': seg_dict, 'punc_dict': punc_dict}
         tokenize_fn = partial(tokenize, **vocab)
         dataset = MapperIterDataPipe(dataset, fn=tokenize_fn)
 
@@ -191,6 +195,10 @@ def Dataset(data_list_file,
                                              sort_size=sort_size,
                                              batch_mode=batch_mode)
 
-    dataset = MapperIterDataPipe(dataset, fn=padding if batch_mode == "padding" else clipping)
+    int_pad_value = conf.get("int_pad_value", -1)
+    float_pad_value = conf.get("float_pad_value", 0.0)
+    padding_conf = {"int_pad_value": int_pad_value, "float_pad_value": float_pad_value}
+    padding_fn = partial(padding, **padding_conf)
+    dataset = MapperIterDataPipe(dataset, fn=padding_fn if batch_mode == "padding" else clipping)
 
     return dataset

funasr/datasets/large_datasets/utils/padding.py

@@ -6,9 +6,8 @@
 def padding(data, float_pad_value=0.0, int_pad_value=-1):
     assert isinstance(data, list)
     assert "key" in data[0]
-    assert "speech" in data[0]
-    assert "text" in data[0]
-
+    assert "speech" in data[0] or "text" in data[0]
+    
     keys = [x["key"] for x in data]
 
     batch = {}

funasr/datasets/large_datasets/utils/tokenize.py

@@ -31,22 +31,43 @@ def seg_tokenize(txt, seg_dict):
 
 def tokenize(data,
              vocab=None,
-             seg_dict=None):
+             seg_dict=None,
+             punc_dict=None):
     assert "text" in data
     assert isinstance(vocab, dict)
     text = data["text"]
     token = []
+    vad = -2
 
     if seg_dict is not None:
         assert isinstance(seg_dict, dict)
         txt = forward_segment("".join(text).lower(), seg_dict)
         text = seg_tokenize(txt, seg_dict)
-    
-    for x in text:
-        if x in vocab:
+
+    length = len(text)
+    for i in range(length):
+        x = text[i]
+        if i == length-1 and "punc" in data and text[i].startswith("vad:"):
+            vad = x[-1][4:]
+            if len(vad) == 0:
+                vad = -1
+            else:
+                vad = int(vad)
+        elif x in vocab:
             token.append(vocab[x])
         else:
             token.append(vocab['<unk>'])
 
+    if "punc" in data and punc_dict is not None:
+        punc_token = []
+        for punc in data["punc"]:
+            if punc in punc_dict:
+                punc_token.append(punc_dict[punc])
+            else:
+                punc_token.append(punc_dict["_"])
+        data["punc"] =  np.array(punc_token)
+
     data["text"] = np.array(token)
+    if vad is not -2:
+        data["vad_indexes"]=np.array([vad], dtype=np.int64)
     return data

funasr/datasets/preprocessor.py

@@ -704,3 +704,103 @@ def pop_split_text_data(self, data: Dict[str, Union[str, np.ndarray]]):
         del data[self.split_text_name]
         return result
 
+class PuncTrainTokenizerCommonPreprocessor(CommonPreprocessor):
+    def __init__(
+            self,
+            train: bool,
+            token_type: List[str] = [None],
+            token_list: List[Union[Path, str, Iterable[str]]] = [None],
+            bpemodel: List[Union[Path, str, Iterable[str]]] = [None],
+            text_cleaner: Collection[str] = None,
+            g2p_type: str = None,
+            unk_symbol: str = "<unk>",
+            space_symbol: str = "<space>",
+            non_linguistic_symbols: Union[Path, str, Iterable[str]] = None,
+            delimiter: str = None,
+            rir_scp: str = None,
+            rir_apply_prob: float = 1.0,
+            noise_scp: str = None,
+            noise_apply_prob: float = 1.0,
+            noise_db_range: str = "3_10",
+            speech_volume_normalize: float = None,
+            speech_name: str = "speech",
+            text_name: List[str] = ["text"],
+            vad_name: str = "vad_indexes",
+    ):
+        # TODO(jiatong): sync with Kamo and Jing on interface for preprocessor
+        super().__init__(
+            train=train,
+            token_type=token_type[0],
+            token_list=token_list[0],
+            bpemodel=bpemodel[0],
+            text_cleaner=text_cleaner,
+            g2p_type=g2p_type,
+            unk_symbol=unk_symbol,
+            space_symbol=space_symbol,
+            non_linguistic_symbols=non_linguistic_symbols,
+            delimiter=delimiter,
+            speech_name=speech_name,
+            text_name=text_name[0],
+            rir_scp=rir_scp,
+            rir_apply_prob=rir_apply_prob,
+            noise_scp=noise_scp,
+            noise_apply_prob=noise_apply_prob,
+            noise_db_range=noise_db_range,
+            speech_volume_normalize=speech_volume_normalize,
+        )
+
+        assert (
+                len(token_type) == len(token_list) == len(bpemodel) == len(text_name)
+        ), "token_type, token_list, bpemodel, or processing text_name mismatched"
+        self.num_tokenizer = len(token_type)
+        self.tokenizer = []
+        self.token_id_converter = []
+
+        for i in range(self.num_tokenizer):
+            if token_type[i] is not None:
+                if token_list[i] is None:
+                    raise ValueError("token_list is required if token_type is not None")
+
+                self.tokenizer.append(
+                    build_tokenizer(
+                        token_type=token_type[i],
+                        bpemodel=bpemodel[i],
+                        delimiter=delimiter,
+                        space_symbol=space_symbol,
+                        non_linguistic_symbols=non_linguistic_symbols,
+                        g2p_type=g2p_type,
+                    )
+                )
+                self.token_id_converter.append(
+                    TokenIDConverter(
+                        token_list=token_list[i],
+                        unk_symbol=unk_symbol,
+                    )
+                )
+            else:
+                self.tokenizer.append(None)
+                self.token_id_converter.append(None)
+
+        self.text_cleaner = TextCleaner(text_cleaner)
+        self.text_name = text_name  # override the text_name from CommonPreprocessor
+        self.vad_name = vad_name
+
+    def _text_process(
+            self, data: Dict[str, Union[str, np.ndarray]]
+    ) -> Dict[str, np.ndarray]:
+        for i in range(self.num_tokenizer):
+            text_name = self.text_name[i]
+            if text_name in data and self.tokenizer[i] is not None:
+                text = data[text_name]
+                text = self.text_cleaner(text)
+                tokens = self.tokenizer[i].text2tokens(text)
+                if "vad:" in tokens[-1]:
+                    vad = tokens[-1][4:]
+                    tokens = tokens[:-1]
+                    if len(vad) == 0:
+                        vad = -1
+                    else:
+                        vad = int(vad)
+                    data[self.vad_name] = np.array([vad], dtype=np.int64)
+                text_ints = self.token_id_converter[i].tokens2ids(tokens)
+                data[text_name] = np.array(text_ints, dtype=np.int64)

funasr/modules/attention.py

@@ -439,6 +439,18 @@ def forward(self, x, mask, mask_shfit_chunk=None, mask_att_chunk_encoder=None):
         att_outs = self.forward_attention(v_h, scores, mask, mask_att_chunk_encoder)
         return att_outs + fsmn_memory
 
+class MultiHeadedAttentionSANMwithMask(MultiHeadedAttentionSANM):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x, mask, mask_shfit_chunk=None, mask_att_chunk_encoder=None):
+        q_h, k_h, v_h, v = self.forward_qkv(x)
+        fsmn_memory = self.forward_fsmn(v, mask[0], mask_shfit_chunk)
+        q_h = q_h * self.d_k ** (-0.5)
+        scores = torch.matmul(q_h, k_h.transpose(-2, -1))
+        att_outs = self.forward_attention(v_h, scores, mask[1], mask_att_chunk_encoder)
+        return att_outs + fsmn_memory
+
 class MultiHeadedAttentionSANMDecoder(nn.Module):
     """Multi-Head Attention layer.
 

funasr/modules/mask.py

@@ -33,3 +33,20 @@ def target_mask(ys_in_pad, ignore_id):
     ys_mask = ys_in_pad != ignore_id
     m = subsequent_mask(ys_mask.size(-1), device=ys_mask.device).unsqueeze(0)
     return ys_mask.unsqueeze(-2) & m
+
+def vad_mask(size, vad_pos, device="cpu", dtype=torch.bool):
+    """Create mask for decoder self-attention.
+
+    :param int size: size of mask
+    :param int vad_pos: index of vad index
+    :param str device: "cpu" or "cuda" or torch.Tensor.device
+    :param torch.dtype dtype: result dtype
+    :rtype: torch.Tensor (B, Lmax, Lmax)
+    """
+    ret = torch.ones(size, size, device=device, dtype=dtype)
+    if vad_pos <= 0 or vad_pos >= size:
+        return ret
+    sub_corner = torch.zeros(
+        vad_pos - 1, size - vad_pos, device=device, dtype=dtype)
+    ret[0:vad_pos - 1, vad_pos:] = sub_corner
+    return ret

funasr/punctuation/abs_model.py

@@ -25,3 +25,7 @@ class AbsPunctuation(torch.nn.Module, BatchScorerInterface, ABC):
     @abstractmethod
     def forward(self, input: torch.Tensor, hidden: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         raise NotImplementedError
+
+    @abstractmethod
+    def with_vad(self) -> bool:
+        raise NotImplementedError