以我在做的对话回复生成方向一个比较著名的benchmark,DialogGPT,由微软开发,在huggingface中的链接如下:microsoft/DialoGPT-small · Hugging Face
它也有对应的medium,large版本。
首先假设我们有一个数据集:
【“Hey my name is Julien! How are you?”,“I'm a girl.”】
前者为对话历史,后者为对应的回复,我们首先要做的是加载模型,将两句话进行编码
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
# 这里如果你将模型下载到本地了,就改成本地的地址xxx文件夹
# 也就是将“microsoft/DialoGPT-small”改成"xxx/"
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")
context="Hey my name is Julien! How are you?"
response="I'm a girl."
context=tokenizer.encode(context,return_tensors='pt').squeeze(0) #[batch_size,seq_len1]
response=tokenizer.encode(response,return_tensors='pt').squeeze(0) #[batch_size,seq_len2]
然后获得我们所需的输入,要注意的是如果采用的是Encoder-Decoder架构,例如T5,input_ids对应的就是对话历史的编码,labels对应的就是回复的编码。对于自回归语言模型(例如GPT-2),input_ids是对话历史与回复编码的concate,labels是将input_ids中的context部分变为-100后的tensor(也就是只保留response部分,其余为-100).
input_ids=torch.cat([torch.tensor([50256]),context,response])
labels = torch.cat([torch.tensor([-100]+[-100 for _ in contextt]), response])
# 对于自回归模型,我们不是使用response作为label,而是将response补齐至input_ids相同形状
# Train
# 如果使用了padding,则需要加入attention_mask参数,padding的部分填为0,未pad的部分变为1
# 可以使用如下代码
# attention_mask=(input_ids!=50256).float() #True的部分就是pad的部分,经过float后会变成1
outputs = model(input_ids=input_ids, labels=labels)
optimizer.zero_grad()
outputs.loss.backward() #在默认状态下,采用的是outputs.XXX获得对应的属性
optimizer.step()
# 生成
pred = model.generate(
input_ids=input_ids,
max_new_tokens=100,
temperature=1.0,
top_k=10,
top_p=0.9,
do_sample=True,
num_return_sequences=1,
pad_token_id=model.config.pad_token_id)[:,length:]
#这返回的是对应的模型预测出的token_ids,使用tokenizer.decode解码
pred=tokenizer.batch_decode(preds, skip_special_tokens=True)