Views
No views yet
import transformers
import torch
# 加载训练好的模型
model_path = "hopetes/pdfChat"
model = GPT2LMHeadModel.from_pretrained(model_path)
tokenizer = GPT2Tokenizer.from_pretrained(model_path)
# 连续对话存储历史
chat_history = []
def chatbot():
global chat_history
print("开始对话(输入 'exit' 退出)")
while True:
user_input = input("User: ")
if user_input.lower() == "exit":
break
# 添加用户输入到对话历史
chat_history.append(f"User: {user_input}")
inputs = tokenizer(" ".join(chat_history), return_tensors="pt", truncation=False)
# 模型生成回复
outputs = model.generate(inputs["input_ids"], max_length=1024, pad_token_id=tokenizer.eos_token_id)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
# 处理并显示回复
bot_reply = response.split("User:")[-1].strip() # 获取最后的回复部分
print(f"Bot: {bot_reply}")
# 添加模型的回复到历史记录
chat_history.append(f"Bot: {bot_reply}")
# 开始交互
chatbot()