改进代码
时间: 2023-07-02 19:06:32 浏览: 42
非常抱歉,我之前的回答有误。CSDN ChitGPT 是一个基于GPT-2的中文聊天机器人,与OpenAI的ChatGPT确实有关系。以下是改进代码:
```python
import torch
import torch.nn.functional as F
from transformers import GPT2Tokenizer, GPT2LMHeadModel
# 加载预训练模型和分词器
tokenizer = GPT2Tokenizer.from_pretrained('uer/gpt2-chinese-cluecorpussmall')
model = GPT2LMHeadModel.from_pretrained('uer/gpt2-chinese-cluecorpussmall')
# 设置GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
# 最大生成长度
max_length = 50
# 对话循环
while True:
# 输入问题
question = input("你:")
# 将问题转换为模型输入格式
input_ids = tokenizer.encode(question, return_tensors='pt').to(device)
# 生成答案
output = model.generate(
input_ids=input_ids,
max_length=max_length,
pad_token_id=tokenizer.eos_token_id
)
# 将生成的答案转换为人类可读的格式
answer = tokenizer.decode(output[0], skip_special_tokens=True)
# 输出答案
print("机器人:", answer)
```
这里使用了中文预训练模型 `uer/gpt2-chinese-cluecorpussmall`,并且将生成的答案转换为了人类可读的格式。同时,我们也将模型移动到GPU上进行加速。