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.
if is_fast_model:
text += self._punc_list[l - 1] 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,6 +163,7 @@ class TextEngine(BaseEngine):
return False return False
self.executor = TextServerExecutor() self.executor = TextServerExecutor()
if 'fast' in config.model_type:
self.executor._init_from_path_new( self.executor._init_from_path_new(
task=config.task, task=config.task,
model_type=config.model_type, model_type=config.model_type,
@ -166,7 +171,15 @@ class TextEngine(BaseEngine):
cfg_path=config.cfg_path, cfg_path=config.cfg_path,
ckpt_path=config.ckpt_path, ckpt_path=config.ckpt_path,
vocab_file=config.vocab_file) 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