hopetes/pdfChat

Use with transformers

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()