代码如下
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
# 加载模型和分词器
def load_model_and_tokenizer(model_name):
"""
加载预训练模型和分词器
:param model_name: 模型名称或路径
:return: 模型和分词器
"""
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
return model, tokenizer
# 生成回复
def generate_response(prompt, model, tokenizer, max_length=50, temperature=0.7):
"""
生成模型的回复
:param prompt: 用户输入的文本
:param model: 加载的模型
:param tokenizer: 分词器
:param max_length: 生成文本的最大长度
:param temperature: 控制生成文本的随机性
:return: 生成的回复
"""
# 对输入进行编码
inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
# 生成回复
outputs = model.generate(
inputs,
max_length=max_length,
num_return_sequences=1,
temperature=temperature,
pad_token_id=tokenizer.eos_token_id,
do_sample=True
)
# 解码生成的回复
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
return response
# 主函数
def main():
# 模型名称或路径
model_name = "./Llama-3.2-1B-Instruct" # 你可以替换为你自己的模型路径
# 加载模型和分词器
print("正在加载模型，请稍等...")
model, tokenizer = load_model_and_tokenizer(model_name)
# 设置设备为 GPU（如果可用）
global device
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
print(f"模型已加载到设备: {device}")
# 开始对话
print("开始对话吧！输入 '退出' 结束对话。")
while True:
user_input = input("你: ")
if user_input.lower() in ["退出", "exit", "quit"]:
print("对话结束。")
break
# 生成回复
response = generate_response(user_input, model, tokenizer)
print(f"模型: {response}")
if __name__ == "__main__":
main()
代码说明
- 加载模型和分词器：
-
使用 AutoModelForCausalLM 和 AutoTokenizer 加载预训练模型（如 GPT-2）。
-
支持本地模型路径或 Hugging Face 模型名称。
- 生成回复：
-
使用 model.generate 方法生成回复。
-
支持参数调整，如 max_length（生成文本的最大长度）和 temperature（控制生成文本的随机性）。
- 连续对话：
-
使用 while 循环实现连续对话。
-
用户输入“退出”或类似指令时，对话结束。
- 设备设置：
- 自动检测 GPU 是否可用，如果可用则将模型加载到 GPU，否则使用 CPU。
运行代码
-
将上述代码保存为
chatbot.py。 -
在命令行中运行：
python chatbot.py
- 开始对话：
-
输入你的问题或对话内容。
-
输入“退出”或类似指令结束对话。
示例对话
正在加载模型，请稍等...
模型已加载到设备: cuda
开始对话吧！输入 '退出' 结束对话。
你: 你好，今天天气怎么样？
模型: 你好！今天天气很好，阳光明媚，适合出去走走。
你: 你觉得人工智能会取代人类吗？
模型: 人工智能可能会在某些领域取代人类，但它不会完全取代人类。人类有独特的创造力和情感。
你: 退出
对话结束。