你可以给一段示例代码吗
时间: 2024-01-26 19:04:01 浏览: 20
以下是一个使用ChitGPT进行对话的示例代码:
```python
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
# 加载tokenizer和模型
tokenizer = AutoTokenizer.from_pretrained("chitgpt-small")
model = AutoModelForCausalLM.from_pretrained("chitgpt-small")
# 设置对话起始语句
user_input = "你好,我想问一个问题"
# 对话循环
while True:
# 将用户输入进行编码
input_ids = tokenizer.encode(user_input, return_tensors='pt')
# 生成回复
reply_ids = model.generate(
input_ids,
max_length=50, # 最大生成长度
pad_token_id=tokenizer.eos_token_id, # 结束标志
do_sample=True, # 开启采样
top_k=50, # 选择概率最大的前50个词
top_p=0.95, # 选择概率累计大于0.95的词
temperature=0.8, # 生成温度,控制多样性
)
# 将回复进行解码并打印
reply_text = tokenizer.decode(reply_ids[0], skip_special_tokens=True)
print("ChitGPT:", reply_text)
# 获取新的用户输入
user_input = input("你:")
```
以上代码中,我们首先加载了ChitGPT的tokenizer和模型,然后设置了对话起始语句。在对话循环中,我们首先将用户输入进行编码,然后使用ChitGPT生成回复。在生成回复时,我们开启了采样,并且设置了一些参数以控制生成的多样性和质量。最后将回复进行解码并打印,并获取新的用户输入,循环进行对话。