Views
No views yet
1# -- coding: utf-8 --
2# @time : 2024/12/1 16:26
3# @author : shajiu
4# @email : 18810979033@163.com
5# @file : .py
6# @software: pycharm
7
8
9from transformers import AutoTokenizer
10from transformers import AutoModelForCausalLM, BitsAndBytesConfig
11import torch
12from peft import PeftModel
13
14class ModelUtils(object):
15
16 @classmethod
17 def load_model(cls, model_name_or_path, load_in_4bit=False, adapter_name_or_path=None):
18 # 是否使用4bit量化进行推理
19 if load_in_4bit:
20 quantization_config = BitsAndBytesConfig(
21 load_in_4bit=True,
22 bnb_4bit_compute_dtype=torch.float16,
23 bnb_4bit_use_double_quant=True,
24 bnb_4bit_quant_type="nf4",
25 llm_int8_threshold=6.0,
26 llm_int8_has_fp16_weight=False,
27 )
28 else:
29 quantization_config = None
30
31 # 加载base model
32 model = AutoModelForCausalLM.from_pretrained(
33 model_name_or_path,
34 load_in_4bit=load_in_4bit,
35 trust_remote_code=True,
36 low_cpu_mem_usage=True,
37 torch_dtype=torch.float16,
38 device_map='auto',
39 quantization_config=quantization_config
40 )
41
42 # 加载adapter
43 if adapter_name_or_path is not None:
44 model = PeftModel.from_pretrained(model, adapter_name_or_path)
45
46 return model
47
48
49def main(model_name_or_path):
50 # 使用合并后的模型进行推理
51 adapter_name_or_path = None
52
53 # 使用base model和adapter进行推理
54 # model_name_or_path = 'shajiu/Tibetan_Llama2_7B_Mental_Health'
55 # adapter_name_or_path = 'shajiu/Tibetan_Llama2_7B_Mental_Health'
56
57 # 是否使用4bit进行推理，能够节省很多显存，但效果可能会有一定的下降
58 load_in_4bit = False
59 device = 'cuda'
60
61 # 生成超参配置
62 max_new_tokens = 500 # 每轮对话最多生成多少个token
63 history_max_len = 1000 # 模型记忆的最大token长度
64 top_p = 0.9
65 temperature = 0.35
66 repetition_penalty = 1.0
67
68 # 加载模型
69 model = ModelUtils.load_model(
70 model_name_or_path,
71 load_in_4bit=load_in_4bit,
72 adapter_name_or_path=adapter_name_or_path
73 ).eval()
74 # 加载tokenizer
75 tokenizer = AutoTokenizer.from_pretrained(
76 model_name_or_path,
77 trust_remote_code=True,
78 # llama不支持fast
79 use_fast=False if model.config.model_type == 'llama' else True
80 )
81 # QWenTokenizer比较特殊，pad_token_id、bos_token_id、eos_token_id均为None。eod_id对应的token为<|endoftext|>
82 if tokenizer.__class__.__name__ == 'QWenTokenizer':
83 tokenizer.pad_token_id = tokenizer.eod_id
84 tokenizer.bos_token_id = tokenizer.eod_id
85 tokenizer.eos_token_id = tokenizer.eod_id
86
87 # 记录所有历史记录
88 if model.config.model_type != 'chatglm':
89 history_token_ids = torch.tensor([[tokenizer.bos_token_id]], dtype=torch.long)
90 else:
91 history_token_ids = torch.tensor([[]], dtype=torch.long)
92
93 # 开始对话
94 utterance_id = 0 # 记录当前是第几轮对话，为了契合chatglm的数据组织格式
95 user_input = input('User：')
96 while True:
97 utterance_id += 1
98 # chatglm使用官方的数据组织格式
99 if model.config.model_type == 'chatglm':
100 user_input = '[Round {}]\n\n问：{}\n\n答：'.format(utterance_id, user_input)
101 user_input_ids = tokenizer(user_input, return_tensors="pt", add_special_tokens=False).input_ids
102 # firefly的数据组织格式
103 # 为了兼容qwen-7b，因为其对eos_token进行tokenize，无法得到对应的eos_token_id
104 else:
105 input_ids = tokenizer(user_input, return_tensors="pt", add_special_tokens=False).input_ids
106 eos_token_id = torch.tensor([[tokenizer.eos_token_id]], dtype=torch.long)
107 user_input_ids = torch.concat([input_ids, eos_token_id], dim=1)
108 history_token_ids = torch.concat((history_token_ids, user_input_ids), dim=1)
109 model_input_ids = history_token_ids[:, -history_max_len:].to(device)
110 with torch.no_grad():
111 outputs = model.generate(
112 input_ids=model_input_ids, max_new_tokens=max_new_tokens, do_sample=True, top_p=top_p,
113 temperature=temperature, repetition_penalty=repetition_penalty, eos_token_id=tokenizer.eos_token_id
114 )
115 model_input_ids_len = model_input_ids.size(1)
116 response_ids = outputs[:, model_input_ids_len:]
117 history_token_ids = torch.concat((history_token_ids, response_ids.cpu()), dim=1)
118 response = tokenizer.batch_decode(response_ids)
119 print("Firefly：" + response[0].strip().replace(tokenizer.eos_token, ""))
120 user_input = input('User：')
121
122
123if __name__ == '__main__':
124 model_name_or_path = 'E:\models\shajiuTibetan_Llama2_7B_Mental_Health'
125 main(model_name_or_path)