Views
No views yet
1import os
2import platform
3import torch
4from transformers import AutoTokenizer, AutoModel
5
6#current_dir = os.path.dirname(os.path.abspath(__file__))
7#model_path = os.path.join(current_dir, 'cntd','CNTDAI-6B')
8model_path = "cntd/CNTDAI-6B"
9print("是否可用:", torch.cuda.is_available()) # 查看GPU是否可用
10print("GPU数量:", torch.cuda.device_count()) # 查看GPU数量
11print("torch方法查看CUDA版本:", torch.version.cuda) # torch方法查看CUDA版本
12print("GPU索引号:", torch.cuda.current_device()) # 查看GPU索引号
13tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
14model = AutoModel.from_pretrained(model_path, trust_remote_code=True).half().cuda()
15# 多显卡支持,使用下面两行代替上面一行,将num_gpus改为你实际的显卡数量
16# from utils import load_model_on_gpus
17# model = load_model_on_gpus(model_path, num_gpus=2)
18model = model.eval()
19os_name = platform.system()
20clear_command = 'cls' if os_name == 'Windows' else 'clear'
21stop_stream = False
22
23
24def build_prompt(history):
25 prompt = "欢迎使用 CNTDAI-6B 模型,输入内容即可进行对话,clear 清空对话历史,stop 终止程序"
26 for query, response in history:
27 prompt += f"\n\n用户:{query}"
28 prompt += f"\n\nCNTDAI-6B:{response}"
29 return prompt
30
31
32
33
34def main():
35 past_key_values, history = None, []
36 global stop_stream
37 print("欢迎使用 CNTDAI-6B 模型,输入内容即可进行对话,clear 清空对话历史,stop 终止程序")
38 while True:
39 query = input("\n用户:")
40 if query.strip() == "stop":
41 break
42 if query.strip() == "clear":
43 past_key_values, history = None, []
44 os.system(clear_command)
45 print("欢迎使用 CNTDAI-6B 模型,输入内容即可进行对话,clear 清空对话历史,stop 终止程序")
46 continue
47 print("\nCNTDAI:", end="")
48 current_length = 0
49 for response, history, past_key_values in model.stream_chat(tokenizer, query, history=history,
50 past_key_values=past_key_values,
51 return_past_key_values=True):
52 if stop_stream:
53 stop_stream = False
54 break
55 else:
56 print(response[current_length:], end="", flush=True)
57 current_length = len(response)
58 print("")
59
60
61if __name__ == "__main__":
62 main()
63
64
65