Update text_engine.py

pull/2464/head
Zhao Yuting 2 years ago committed by GitHub
parent 80837fd658
commit 304dc2603c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -107,11 +107,14 @@ class PaddleTextConnectionHandler:
assert len(tokens) == len(labels) assert len(tokens) == len(labels)
text = '' text = ''
is_fast_model = 'fast' in self.text_engine.config.model_type
for t, l in zip(tokens, labels): for t, l in zip(tokens, labels):
text += t text += t
if l != 0: # Non punc. if l != 0: # Non punc.
text += self._punc_list[l - 1] if is_fast_model:
text += self._punc_list[l - 1]
else:
text += self._punc_list[l]
return text return text
else: else:
raise NotImplementedError raise NotImplementedError
@ -131,6 +134,7 @@ class TextEngine(BaseEngine):
""" """
super(TextEngine, self).__init__() super(TextEngine, self).__init__()
logger.debug("Create the TextEngine Instance") logger.debug("Create the TextEngine Instance")
def init(self, config: dict): def init(self, config: dict):
"""Init the Text Engine """Init the Text Engine
@ -159,14 +163,23 @@ class TextEngine(BaseEngine):
return False return False
self.executor = TextServerExecutor() self.executor = TextServerExecutor()
self.executor._init_from_path_new( if 'fast' in config.model_type:
task=config.task, self.executor._init_from_path_new(
model_type=config.model_type, task=config.task,
lang=config.lang, model_type=config.model_type,
cfg_path=config.cfg_path, lang=config.lang,
ckpt_path=config.ckpt_path, cfg_path=config.cfg_path,
vocab_file=config.vocab_file) ckpt_path=config.ckpt_path,
vocab_file=config.vocab_file)
else:
self.executor._init_from_path(
task=config.task,
model_type=config.model_type,
lang=config.lang,
cfg_path=config.cfg_path,
ckpt_path=config.ckpt_path,
vocab_file=config.vocab_file)
logger.info("Using model: %s." % (config.model_type))
logger.info("Initialize Text server engine successfully on device: %s." logger.info("Initialize Text server engine successfully on device: %s."
% (self.device)) % (self.device))
return True return True

Loading…
Cancel
Save