Upload 6 files

c4ce149 verified 12 days ago

2.02 kB

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}")

	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}")