File size: 2,015 Bytes
c4ce149 | 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 | import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import re
# 1. 加载训练好的模型和分词器
model_path = "./model" # 对应你训练保存的路径
print("正在加载模型...")
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model.eval() # 切换到评估模式
print("模型加载完成！")
# 2. 定义推理对话函数
def generate_response(user_message, max_new_tokens=128):
prompt = f"Human:{user_message}\nBot:"
# 编码输入
inputs = tokenizer(prompt, return_tensors="pt").to(device)
prompt_len = inputs["input_ids"].shape[1] # 获取输入的 token 数量
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=max_new_tokens,
do_sample=True,
temperature=0.7,
top_p=0.9,
repetition_penalty=1.2, # 抑制重复的核心参数
no_repeat_ngram_size=3, # 防止短语死循环
pad_token_id=tokenizer.pad_token_id,
eos_token_id=tokenizer.eos_token_id
)
# 直接切片去掉输入 Prompt 部分的 token
generated_tokens = outputs[0][prompt_len:]
response = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
# 清理可能残留的 "Bot:" 前缀
response = re.sub(r'^(Bot:|bot:|ot:|\s)+', '', response).strip()
return response
# 3. 简单的交互式对话循环
if __name__ == "__main__":
print("\n--- GPT-2 中文聊天机器人已启动 (输入 'exit' 或 'quit' 退出) ---")
while True:
user_input = input("\nHuman: ")
if user_input.strip().lower() in ["exit", "quit"]:
print("再见！")
break
if not user_input.strip():
continue
bot_reply = generate_response(user_input).replace(" ", "")
print(f"Bot: {bot_reply}") |