paddlepaddle--paddlenlp
wehub-resource-sync
81 行
2.9 KiB
Python
81 行
2.9 KiB
Python
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import argparse
|
|
import time
|
|
|
|
import torch
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
|
|
def parse_args(prog=None):
|
|
"""
|
|
parse_args
|
|
"""
|
|
parser = argparse.ArgumentParser(prog=prog)
|
|
parser.add_argument("--model_name_or_path", type=str, help="model name or local path", required=True)
|
|
parser.add_argument("--do_forward", action="store_true", help="fowrward test")
|
|
parser.add_argument("--do_generate", action="store_true", help="generate test")
|
|
return parser.parse_args()
|
|
|
|
|
|
@torch.no_grad()
|
|
def predict_generate(model, inputs):
|
|
for i in range(10):
|
|
start = time.perf_counter()
|
|
generate_ids = model.generate(
|
|
inputs.input_ids,
|
|
max_length=100,
|
|
do_sample=False,
|
|
)
|
|
hf_cost = (time.perf_counter() - start) * 1000
|
|
print("Speed test:", hf_cost)
|
|
result = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
|
|
print(result)
|
|
|
|
|
|
@torch.no_grad()
|
|
def predict_forward(model, inputs):
|
|
for i in range(10):
|
|
start = time.perf_counter()
|
|
_ = model(**inputs)
|
|
hf_cost = (time.perf_counter() - start) * 1000
|
|
print("Speed test:", hf_cost)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_args()
|
|
all_texts = [
|
|
"你好",
|
|
"去年9月，拼多多海外版“Temu”正式在美国上线。数据显示，截至2023年2月23日，Temu App新增下载量4000多万，新增用户增速第一，AppStore购物榜霸榜69天、Google Play购物榜霸榜114天。",
|
|
]
|
|
tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
|
|
if "llama" in args.model_name_or_path:
|
|
tokenizer.pad_token_id = 0
|
|
model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path).cuda()
|
|
model = model.eval()
|
|
if args.do_forward:
|
|
for input_text in all_texts:
|
|
print(f"text: {input_text}")
|
|
inputs = tokenizer([input_text], return_tensors="pt", max_length=50, padding=True)
|
|
inputs = inputs.to("cuda")
|
|
predict_forward(model, inputs)
|
|
|
|
if args.do_generate:
|
|
for input_text in all_texts:
|
|
print(f"text: {input_text}")
|
|
inputs = tokenizer([input_text], return_tensors="pt", max_length=50, padding=True)
|
|
inputs = inputs.to("cuda")
|
|
predict_generate(model, inputs)
|