import logging
|
|
from funasr.models.target_delay_transformer import TargetDelayTransformer
|
from funasr.models.vad_realtime_transformer import VadRealtimeTransformer
|
from funasr.torch_utils.initialize import initialize
|
from funasr.train.abs_model import PunctuationModel
|
from funasr.train.class_choices import ClassChoices
|
|
punc_choices = ClassChoices(
|
"punctuation",
|
classes=dict(
|
target_delay=TargetDelayTransformer,
|
vad_realtime=VadRealtimeTransformer
|
),
|
default="target_delay",
|
)
|
model_choices = ClassChoices(
|
"model",
|
classes=dict(
|
punc=PunctuationModel,
|
),
|
default="punc",
|
)
|
class_choices_list = [
|
# --punc and --punc_conf
|
punc_choices,
|
# --model and --model_conf
|
model_choices
|
]
|
|
|
def build_punc_model(args):
|
# token_list and punc list
|
if isinstance(args.token_list, str):
|
with open(args.token_list, encoding="utf-8") as f:
|
token_list = [line.rstrip() for line in f]
|
args.token_list = token_list.copy()
|
if isinstance(args.punc_list, str):
|
with open(args.punc_list, encoding="utf-8") as f2:
|
pairs = [line.rstrip().split(":") for line in f2]
|
punc_list = [pair[0] for pair in pairs]
|
punc_weight_list = [float(pair[1]) for pair in pairs]
|
args.punc_list = punc_list.copy()
|
elif isinstance(args.punc_list, list):
|
punc_list = args.punc_list.copy()
|
punc_weight_list = [1] * len(punc_list)
|
if isinstance(args.token_list, (tuple, list)):
|
token_list = args.token_list.copy()
|
else:
|
raise RuntimeError("token_list must be str or dict")
|
|
vocab_size = len(token_list)
|
punc_size = len(punc_list)
|
logging.info(f"Vocabulary size: {vocab_size}")
|
|
# punc
|
punc_class = punc_choices.get_class(args.punctuation)
|
punc = punc_class(vocab_size=vocab_size, punc_size=punc_size, **args.punctuation_conf)
|
|
if "punc_weight" in args.model_conf:
|
args.model_conf.pop("punc_weight")
|
model = PunctuationModel(punc_model=punc, vocab_size=vocab_size, punc_weight=punc_weight_list, **args.model_conf)
|
|
# initialize
|
if args.init is not None:
|
initialize(model, args.init)
|
|
return model
|
