File size: 5,755 Bytes
45cf443 2c87b68 92c7321 45cf443 4020c99 c6bc767 45cf443 4020c99 45cf443 4020c99 45cf443 2c87b68 45cf443 4020c99 45cf443 | 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 | import time
import argparse
import random
import warnings
import torch
from transformers import AutoTokenizer, TextStreamer
from models import LMConfig, LMForCausalLM
from models.lm.lora import * # noqa: F401,F403
from utils.training import setup_seed, get_model_params
warnings.filterwarnings('ignore')
def init_model(args):
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
if args.native:
moe_suffix = '_moe' if args.use_moe else ''
ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
state = torch.load(ckp, map_location=args.device)
n_layers = max(int(k.split('.')[2]) for k in state if k.startswith('model.layers.')) + 1
model = LMForCausalLM(LMConfig(
hidden_size=args.hidden_size,
num_hidden_layers=n_layers,
use_moe=bool(args.use_moe),
inference_rope_scaling=args.inference_rope_scaling
))
model.load_state_dict(state, strict=True)
if args.lora_weight != 'None':
apply_lora(model)
load_lora(model, f'./{args.save_dir}/{args.lora_weight}_{args.hidden_size}.pth')
else:
model = LMForCausalLM.from_pretrained(args.load_from)
get_model_params(model, model.config)
return model.half().eval().to(args.device), tokenizer
def main():
parser = argparse.ArgumentParser(description="MiniMind模型推理与对话")
parser.add_argument('--load_from', default='', type=str, help="模型加载路径（transformers格式，native模式不感知此参数）")
parser.add_argument('--tokenizer_path', default='checkpoint/tokenizer', type=str, help="tokenizer 路径")
parser.add_argument('--native', action='store_true', help="加载原生 torch checkpoint（由 save_dir/weight/hidden_size 定位）")
parser.add_argument('--save_dir', default='out', type=str, help="模型权重目录")
parser.add_argument('--weight', default='full_sft', type=str, help="权重名称前缀（pretrain, full_sft, rlhf, reason, ppo_actor, grpo, spo）")
parser.add_argument('--lora_weight', default='None', type=str, help="LoRA权重名称（None表示不使用，可选：lora_identity, lora_medical）")
parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
parser.add_argument('--inference_rope_scaling', default=False, action='store_true', help="启用RoPE位置编码外推（4倍，仅解决位置编码问题）")
parser.add_argument('--max_new_tokens', default=8192, type=int, help="最大生成长度（注意：并非模型实际长文本能力）")
parser.add_argument('--temperature', default=0.85, type=float, help="生成温度，控制随机性（0-1，越大越随机）")
parser.add_argument('--top_p', default=0.95, type=float, help="nucleus采样阈值（0-1）")
parser.add_argument('--open_thinking', default=0, type=int, help="是否开启自适应思考（0=否，1=是）")
parser.add_argument('--historys', default=0, type=int, help="携带历史对话轮数（需为偶数，0表示不携带历史）")
parser.add_argument('--show_speed', default=1, type=int, help="显示decode速度（tokens/s）")
parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备")
args = parser.parse_args()
prompts = [
'你有什么特长？',
'为什么天空是蓝色的',
'请用Python写一个计算斐波那契数列的函数',
'解释一下"光合作用"的基本过程',
'如果明天下雨，我应该如何出门',
'比较一下猫和狗作为宠物的优缺点',
'解释什么是机器学习',
'推荐一些中国的美食'
]
conversation = []
model, tokenizer = init_model(args)
input_mode = int(input('[0] 自动测试\n[1] 手动输入\n'))
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
prompt_iter = prompts if input_mode == 0 else iter(lambda: input('💬: '), '')
for prompt in prompt_iter:
setup_seed(random.randint(0, 31415926))
if input_mode == 0: print(f'💬: {prompt}')
conversation = conversation[-args.historys:] if args.historys else []
conversation.append({"role": "user", "content": prompt})
if 'pretrain' in args.weight:
inputs = tokenizer.bos_token + prompt
else:
inputs = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True, open_thinking=bool(args.open_thinking))
inputs = tokenizer(inputs, return_tensors="pt", truncation=True).to(args.device)
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=1
)
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')
if __name__ == "__main__":
main()
|