apps/FastDeploy
1
0
Fork 0
mirror of https://github.com/PaddlePaddle/FastDeploy.git synced 2025-10-05 16:48:03 +08:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

Add batch size argument for uie

Browse Source
This commit is contained in:
zhoushunjie
2022-12-27 15:22:09 +00:00
parent df940b750f
commit 60f8f0e11b
6 changed files with 114 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/fastdeploy/text/uie/__init__.py
View File

@@ -50,6 +50,7 @@ class UIEModel(FastDeployModel):
position_prob=0.5,
max_length=128,
schema=[],
batch_size=64,
runtime_option=RuntimeOption(),
model_format=ModelFormat.PADDLE,
schema_language=SchemaLanguage.ZH):
@@ -63,9 +64,10 @@ class UIEModel(FastDeployModel):
else:
assert "The type of schema should be list or dict."
schema_language = C.text.SchemaLanguage(schema_language)
self._model = C.text.UIEModel(
model_file, params_file, vocab_file, position_prob, max_length,
schema, runtime_option._option, model_format, schema_language)
self._model = C.text.UIEModel(model_file, params_file, vocab_file,
position_prob, max_length, schema,
batch_size, runtime_option._option,
model_format, schema_language)
assert self.initialized, "UIEModel initialize failed."
def set_schema(self, schema):