jmwang66
2023-06-29 98abc0e5ac1a1da0fe1802d9ffb623802fbf0b2f
funasr/tasks/punctuation.py
@@ -9,8 +9,6 @@
import numpy as np
import torch
from typeguard import check_argument_types
from typeguard import check_return_type
from funasr.datasets.collate_fn import CommonCollateFn
from funasr.datasets.preprocessor import PuncTrainTokenizerCommonPreprocessor
@@ -47,7 +45,6 @@
    @classmethod
    def add_task_arguments(cls, parser: argparse.ArgumentParser):
        # NOTE(kamo): Use '_' instead of '-' to avoid confusion
        assert check_argument_types()
        group = parser.add_argument_group(description="Task related")
        # NOTE(kamo): add_arguments(..., required=True) can't be used
@@ -126,7 +123,6 @@
            # e.g. --encoder and --encoder_conf
            class_choices.add_arguments(group)
        assert check_return_type(parser)
        return parser
    @classmethod
@@ -136,14 +132,12 @@
        [Collection[Tuple[str, Dict[str, np.ndarray]]]],
        Tuple[List[str], Dict[str, torch.Tensor]],
    ]:
        assert check_argument_types()
        return CommonCollateFn(int_pad_value=0)
    @classmethod
    def build_preprocess_fn(
            cls, args: argparse.Namespace, train: bool
    ) -> Optional[Callable[[str, Dict[str, np.array]], Dict[str, np.ndarray]]]:
        assert check_argument_types()
        token_types = [args.token_type, args.token_type]
        token_lists = [args.token_list, args.punc_list]
        bpemodels = [args.bpemodel, args.bpemodel]
@@ -161,7 +155,6 @@
            )
        else:
            retval = None
        assert check_return_type(retval)
        return retval
    @classmethod
@@ -182,7 +175,6 @@
    @classmethod
    def build_model(cls, args: argparse.Namespace) -> PunctuationModel:
        assert check_argument_types()
        if isinstance(args.token_list, str):
            with open(args.token_list, encoding="utf-8") as f:
                token_list = [line.rstrip() for line in f]
@@ -223,5 +215,4 @@
        if args.init is not None:
            initialize(model, args.init)
        assert check_return_type(model)
        return model