paddlenlp

1
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
2
#
3
# Licensed under the Apache License, Version 2.0 (the "License");
4
# you may not use this file except in compliance with the License.
5
# You may obtain a copy of the License at
6
#
7
#     http://www.apache.org/licenses/LICENSE-2.0
8
#
9
# Unless required by applicable law or agreed to in writing, software
10
# distributed under the License is distributed on an "AS IS" BASIS,
11
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
# See the License for the specific language governing permissions and
13
# limitations under the License.
14

15
import argparse
16
import os
17

18
import paddle
19
from model import PointwiseMatching
20

21
from paddlenlp.transformers import AutoModel
22

23
if __name__ == "__main__":
24
    parser = argparse.ArgumentParser()
25
    parser.add_argument(
26
        "--params_path",
27
        type=str,
28
        required=True,
29
        default="./checkpoint/model_900/model_state.pdparams",
30
        help="The path to model parameters to be loaded.",
31
    )
32
    parser.add_argument(
33
        "--output_path", type=str, default="./output", help="The path of model parameter in static graph to be saved."
34
    )
35
    args = parser.parse_args()
36

37
    pretrained_model = AutoModel.from_pretrained("ernie-3.0-medium-zh")
38
    model = PointwiseMatching(pretrained_model)
39

40
    if args.params_path:
41
        if os.path.isfile(args.params_path):
42
            state_dict = paddle.load(args.params_path)
43
            model.set_dict(state_dict)
44
            print("Loaded parameters from %s" % args.params_path)
45
        elif os.path.isdir(args.params_path):
46
            path = os.path.join(args.params_path, "model_state.pdparams")
47
            state_dict = paddle.load(path)
48
            model.set_dict(state_dict)
49
            print("Loaded parameters from %s" % path)
50
    model.eval()
51

52
    # Convert to static graph with specific input description
53
    model = paddle.jit.to_static(
54
        model,
55
        input_spec=[
56
            paddle.static.InputSpec(shape=[None, None], dtype="int64"),  # input_ids
57
            paddle.static.InputSpec(shape=[None, None], dtype="int64"),  # segment_ids
58
        ],
59
    )
60
    # Save in static graph model.
61
    save_path = os.path.join(args.output_path, "inference")
62
    paddle.jit.save(model, save_path)
63