Spaces:
Sleeping
Sleeping
File size: 4,647 Bytes
df0647f 8f7ca17 df0647f 8f7ca17 df0647f 8f7ca17 df0647f 8f7ca17 df0647f 8f7ca17 df0647f 8f7ca17 df0647f 8f7ca17 df0647f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | #!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
https://github.com/jingyaogong/minimind/blob/master/eval_llm.py
"""
import argparse
import os
from pathlib import Path
import platform
import time
if platform.system() in ("Windows", "Darwin"):
from project_settings import project_path, temp_directory
else:
project_path = os.path.abspath("../../")
project_path = Path(project_path)
temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
import torch
from modelscope import AutoTokenizer, AutoModelForCausalLM
from transformers import TextStreamer
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"--pretrained_model_name_or_path",
default="qgyd2021/Qwen2.5-0.5B-ultrachat-sft-deepspeed",
type=str
)
parser.add_argument(
"--model_cache_dir",
default=(temp_directory / "hub_models").as_posix(),
type=str
)
parser.add_argument(
"--max_new_tokens",
default=8192, # 8192, 128
type=int, help="最大生成长度（注意：并非模型实际长文本能力）"
)
parser.add_argument("--top_p", default=0.85, type=float, help="nucleus采样阈值（0-1）")
parser.add_argument("--temperature", default=0.85, type=float, help="生成温度，控制随机性（0-1，越大越随机）")
parser.add_argument(
"--show_speed",
default=1, # 1, 0
type=int, help="显示decode速度（tokens/s）"
)
args = parser.parse_args()
return args
def main():
args = get_args()
os.environ["MODELSCOPE_CACHE"] = args.model_cache_dir
if torch.cuda.is_available():
device = "cuda"
elif torch.backends.mps.is_available():
# device = "mps"
device = "cpu"
else:
device = "cpu"
print(f"device: {device}")
model = AutoModelForCausalLM.from_pretrained(
args.pretrained_model_name_or_path,
cache_dir=args.model_cache_dir,
trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained(
args.pretrained_model_name_or_path,
cache_dir=args.model_cache_dir,
trust_remote_code=True,
)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
tokenizer.pad_token_id = tokenizer.eos_token_id
model = model.eval().to(device)
# print(tokenizer)
# print(model)
prompts = [
"你有什么特长？",
"为什么天空是蓝色的",
"请用Python写一个计算斐波那契数列的函数",
'解释一下"光合作用"的基本过程',
"如果明天下雨，我应该如何出门",
"比较一下猫和狗作为宠物的优缺点",
"解释什么是机器学习",
"推荐一些中国的美食"
]
input_mode = int(input("[0] 自动测试\n[1] 手动输入\n"))
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
# conversation = list()
conversation = [
{"role": "system", "content": "You are a helpful assistant"}
]
while True:
if input_mode == 0:
if len(prompts) == 0:
break
user_input = prompts.pop(0)
print(f"💬: {user_input}")
else:
user_input = input("💬: ")
user_input = str(user_input).strip()
conversation.append({"role": "user", "content": user_input})
inputs = tokenizer.apply_chat_template(
conversation=conversation,
tokenize=False,
add_generation_prompt=True
)
inputs = tokenizer.__call__(
inputs,
return_tensors="pt",
truncation=True
)
inputs = inputs.to(device)
# print(inputs)
print("🤖: ", end="")
st = time.time()
generated_ids = model.generate(
inputs=inputs["input_ids"], attention_mask=inputs["attention_mask"],
max_new_tokens=args.max_new_tokens, do_sample=True, streamer=streamer,
pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id,
top_p=args.top_p, temperature=args.temperature, repetition_penalty=3.0,
)
response = tokenizer.decode(generated_ids[0][len(inputs["input_ids"][0]):], skip_special_tokens=True)
conversation.append({"role": "assistant", "content": response})
gen_tokens = len(generated_ids[0]) - len(inputs["input_ids"][0])
print(f"\n[Speed]: {gen_tokens / (time.time() - st):.2f} tokens/s\n\n") if args.show_speed else print("\n\n")
return
if __name__ == "__main__":
main()
|