Upload 7 files

7208556 verified 19 days ago

1.99 kB

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 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,
            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:" 或 "ot:"，用正则或 strip 清理掉
    import re
    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}")

	import torch
	from transformers import AutoModelForCausalLM, AutoTokenizer

	# 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,
	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:" 或 "ot:"，用正则或 strip 清理掉
	import re
	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}")