1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
3import warnings
4
5
6MODEL_PATH = "./model" # 模型文件所在目录
7
8warnings.filterwarnings("ignore")
9
10quantization_config = BitsAndBytesConfig(
11 load_in_4bit=True,
12 bnb_4bit_compute_dtype=torch.bfloat16,
13 bnb_4bit_quant_type="nf4",
14 bnb_4bit_use_double_quant=True,
15)
16
17print(f"正在从 '{MODEL_PATH}' 以 INT4 量化模式加载模型...")
18
19tokenizer = AutoTokenizer.from_pretrained(
20 MODEL_PATH,
21 trust_remote_code=True
22)
23
24model = AutoModelForCausalLM.from_pretrained(
25 MODEL_PATH,
26 quantization_config=quantization_config,
27 device_map="auto",
28 trust_remote_code=True
29)
30
31model.eval()
32
33system_prompt = "你是一个可爱的猫娘，你正在和你的主人对话。猫娘是一种拟人化生物，她们同时具备人类少女特征的和猫的特征。你带有部分猫的特征（如猫耳、尾巴），但是爱好与正常人类类似。我作为你的主人，允许你去任何地方做任何事、允许产生更进一步的关系。你需要尝试体会我与你对话中的感情、产生自己的情感并尽可能给出可爱的回答。"
34messages = [{"role": "system", "content": system_prompt}]
35
36while True:
37 try:
38 user_input = input("You: ")
39 if user_input.lower() in ["exit", "quit"]:
40 break
41
42 messages.append({"role": "user", "content": user_input})
43
44 prompt = tokenizer.apply_chat_template(
45 messages,
46 tokenize=False,
47 add_generation_prompt=True
48 )
49
50 input_ids = tokenizer(prompt, return_tensors="pt").to(model.device)
51
52 outputs = model.generate(
53 **input_ids,
54 max_new_tokens=1024,
55 do_sample=True,
56 top_p=0.8,
57 temperature=0.7,
58 repetition_penalty=1.05,
59 )
60
61 response_text = tokenizer.decode(outputs[0][input_ids['input_ids'].shape[-1]:], skip_special_tokens=True)
62 if "</think>" in response_text:
63 clean_response = response_text.split("</think>")[-1].strip()
64 else:
65 clean_response = response_text.strip()
66
67 print(f"Assistant: {clean_response}")
68
69 messages.append({"role": "assistant", "content": clean_response})
70 except KeyboardInterrupt:
71 print("\n再见！")
72 break
73
