From 5fa895c9befdd09d65423188f821d6e93ae8aaec Mon Sep 17 00:00:00 2001 From: hanlintang Date: Fri, 11 Apr 2025 08:05:56 +0000 Subject: [PATCH] [PIR] update document_intelligence.py --- .../pipelines/nodes/document/document_intelligence.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/slm/pipelines/pipelines/nodes/document/document_intelligence.py b/slm/pipelines/pipelines/nodes/document/document_intelligence.py index eaa7ab6d2f3a..ec55cbacf187 100644 --- a/slm/pipelines/pipelines/nodes/document/document_intelligence.py +++ b/slm/pipelines/pipelines/nodes/document/document_intelligence.py @@ -29,7 +29,11 @@ sort_res, ) from paddlenlp.transformers import AutoTokenizer -from paddlenlp.utils.env import PPNLP_HOME +from paddlenlp.utils.env import ( + PADDLE_INFERENCE_MODEL_SUFFIX, + PADDLE_INFERENCE_WEIGHTS_SUFFIX, + PPNLP_HOME, +) logger = logging.getLogger(__name__) @@ -94,8 +98,8 @@ def __init__( def _get_inference_model(self): inference_model_path = os.path.join(self._task_path, "static", "inference") - self._static_model_file = inference_model_path + ".pdmodel" - self._static_params_file = inference_model_path + ".pdiparams" + self._static_model_file = inference_model_path + f"{PADDLE_INFERENCE_MODEL_SUFFIX}" + self._static_params_file = inference_model_path + f"{PADDLE_INFERENCE_WEIGHTS_SUFFIX}" self._config = paddle.inference.Config(self._static_model_file, self._static_params_file) self._prepare_static_mode()