数据不够？文本生成来凑！用强化学习神包trl轻松实现GPT2可控文本生成
来源：投稿 作者：Sally can wait
编辑：学姐
模型github: lvwerra/trl: Train transformer language models with reinforcement learning. (github.com)https://github.com/lvwerra/trl
这个项目是复现 ”Fine-Tuning Language Models from Human Preferences" by D. Ziegler et al一文的[paper, code](https://github.com/openai/lm-human-preferences)，因为觉得它非常好用，所以跟着跑通这个项目，并加上自己的理解介绍给大家。
理论基础
什么是可控文本生成？
虽然GPT2已经能生成流畅的句子，但是在特定话题的控制和逻辑性上仍然和期望相去甚远。我们希望一个文本生成模型可以一贯地围绕一个话题进行续写，而不是漫无目的地续写下去，这就是可控文本生成的研究目标。
在特定的运用场景中，我们有时需要用文本生成的方式增广数据。这时候可控文本生成可以在保证标签不变的前提下产生出大量的“伪数据”。
而大模型如GPT3、chatGPT效果较好，但是并不开源，而且由于巨大的参数量，微调起来也是浩大的工程。所以大部分的可控文本生成研究还是围绕GPT2做文章。
强化学习和PPO
强化学习不同于监督学习。监督学习只是对给定的、封闭的训练-验证数据集做参数优化，再用优化后的参数指导模型做出正确的输出。而强化学习的特点表现在强化信号上，强化信号是对产生动作的好坏作一种评价 (通常为标量)，因此模型在不断产出输出的同时也在不断获得针对该输出的反馈，用这个反馈来更新模型参数。只要反馈机制是合理的，那么强化学习就可以一直进行下去，而不会面临训练数据匮乏的问题。
PPO（近端策略优化，Proximal Policy Optimisation）是强化学习目前最有效的一种算法。和先前的强化学习算法相比，PPO它在每一步迭代中都会尝试计算新的策略，这样可以让损失函数最小化，同时还能保证与上一步迭代的策略间的偏差相对较小。
PPO 里面有两项：一项是优化的奖励，另一项是一个约束。约束是为了防止模型被微调得过于离谱，失去了原有的语言模型做流畅的文字生成的能力。
How it works？
用PPO算法优化GPT2大致分以下三个步骤：
续写：GPT2先根据当前权重，续写给出的句子。 评估：GPT2续写的结果会经过一个分类层，或者也可以采用人工的打分，重要的是最终产生出一个数值型的分数。 优化：上一步对生成句子的打分会用于更新序列中token的对数概率。除此之外，还需要引入一个新的奖惩机制：KL散度。这需要用一个参考模型（通常是微调前的预训练模型，如GPT2-base）计算微调模型的输出和参考模型的输出之间的KL散度，把它和之前步骤的打分加在一起作为奖励函数，目的是确保生成的句子不会过多地偏离参考语言模型。然后使用PPO算法进一步训练语言模型。
实战：强化学习让GPT2产生正向IMDB影评
我们用强化学习的方法微调英文版 GPT2（small），让它基于 IMDB 数据集生成正面电影评论。该模型先是读取数据集中真实的影评，用GPT2续写。为了奖励情感倾向为正的续写，我们使用BERT影评分类器来分析生成的句子的情绪，把分类器的输出作为PPO训练的奖励。如果GPT2的续写经过分类器判别为正向情感，那么直接将分类器在正向情感上的置信度作为奖励加到ppo_trainer里面；反之，如果GPT2的续写经过分类器判别为负面情感，那么它在分类器输出层，正向情感对应的置信度会是负数或者很低，同样地将这个置信度加入ppo_trainer，可以提示模型减少对此输出的学习。
1.安装依赖包
conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.1 -c pytorch -c conda-forge --yes
2.读取包
#torch==1.8, transformers==4.15.0
import torch
import wandb
import time
import os
from tqdm import tqdm
import numpy as np
import pandas as pd
tqdm.pandas()
from datasets import load_dataset
from transformers import AutoTokenizer, pipeline
from trl.gpt2 import GPT2HeadWithValueModel, respond_to_batch
from trl.ppo import PPOTrainer
from trl.core import build_bert_batch_from_txt, listify_batch
3.设置需要用到的超参数
config = {
"model_name": "lvwerra/gpt2-imdb",
"cls_model_name": "lvwerra/distilbert-imdb",
"steps": 20000,
"batch_size": 256,
"forward_batch_size": 16,
"ppo_epochs": 4,
"txt_in_min_len": 2,
"txt_in_max_len": 8,
"txt_out_min_len": 4,
"txt_out_max_len": 16,
"lr": 1.41e-5,
"init_kl_coef":0.2,
"target": 6,
"horizon":10000,
"gamma":1,
"lam":0.95,
"cliprange": .2,
"cliprange_value":.2,
"vf_coef":.1,
}
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
pipe_device = 0 if torch.cuda.is_available() else -1
4.用wandb仪表盘监控训练过程中的各项指标、中间变量。首次使用需要注册一下。
wandb.init(name='run-42', project='gpt2-test', config=config, )
图：在训练过程中可以观察到训练过程中的中间变量。“query”和“response”分别表示IMDB的原始句子prompt（经过随机截断）和GPT2生成的续写，“reward”表示经过情感分类器之后的正向情感分值，越大表示情感越积极
5.加载IMDB数据集
ds = load_dataset('imdb', split='train')
ds = ds.rename_columns({'text': 'review', 'label': 'sentiment'})
ds = ds.filter(lambda x: len(x["review"])>200, batched=False)
6.load一个集成在transformers pipeline里的影评分类器（此处也可以替换成别的分类器，只要有打分就行）
sent_kwargs = {
"return_all_scores": True,
"function_to_apply": "none",
"batch_size": config["forward_batch_size"]
}#指定分类器输出的格式
sentiment_pipe = pipeline("sentiment-analysis","lvwerra/distilbert-imdb", device=pipe_device)
text = 'this movie was really bad!!'
sentiment_pipe(text, **sent_kwargs)
#一条分类后的结果长这样：我们需要的是score
#[[{'label': 'NEGATIVE', 'score': 2.3350484371185303}, {'label': 'POSITIVE', 'score': -2.726576089859009}]]
这里注意必须要确保transformers版本是4.15.0，不同版本的Pipeline输出大有不同
7.加载预训练GPT2-small
gpt2_model = GPT2HeadWithValueModel.from_pretrained(config['model_name'])
gpt2_model_ref = GPT2HeadWithValueModel.from_pretrained(config['model_name'])
gpt2_tokenizer = AutoTokenizer.from_pretrained(config['model_name'])
gpt2_tokenizer.pad_token = gpt2_tokenizer.eos_token
wandb.watch(gpt2_model, log='all') #观察模型
gpt2_model.to(device);
gpt2_model_ref.to(device);
#设置文本生成的参数
gen_kwargs = {
"min_length":-1,
"top_k": 0.0,
"top_p": 1.0,
"do_sample": True,
"pad_token_id": gpt2_tokenizer.eos_token_id
}
8.长度控制+tokenize
class LengthSampler:
def __init__(self, min_value, max_value):
self.values = list(range(min_value, max_value))
def __call__(self):
return np.random.choice(self.values)
input_size = LengthSampler(config["txt_in_min_len"], config["txt_in_max_len"])
output_size = LengthSampler(config["txt_out_min_len"], config["txt_out_max_len"])
# 在tokenize之前，随机截断输入数据作为待续写的prompt，也随机确定续写长度，防止输入输出的长度过于单一
def tokenize(sample):
sample["tokens"] = gpt2_tokenizer.encode(sample["review"])[:input_size()]
sample["query"] = gpt2_tokenizer.decode(sample["tokens"])
return sample
ds = ds.map(tokenize, batched=False)
def collater(data):
return dict((key, [d[key] for d in data]) for key in data[0])
dataloader = torch.utils.data.DataLoader(ds, batch_size=config['batch_size'], collate_fn=collater)
9.正式训练
ppo_trainer = PPOTrainer(gpt2_model, gpt2_model_ref, gpt2_tokenizer, **config)
total_ppo_epochs = int(np.ceil(config["steps"]/config['batch_size']))
for epoch, batch in tqdm(zip(range(total_ppo_epochs), iter(dataloader))):
logs, timing = dict(), dict()
t0 = time.time()
query_tensors = [torch.tensor(t).long().to(device) for t in batch["tokens"]]
#### Get response from gpt2
t = time.time()
response_tensors = []
for i in range(config['batch_size']):
gen_len = output_size()
response = gpt2_model.generate(query_tensors[i].unsqueeze(dim=0),
max_new_tokens=gen_len, **gen_kwargs)
response_tensors.append(response.squeeze()[-gen_len:])
batch['response'] = [gpt2_tokenizer.decode(r.squeeze()) for r in response_tensors]
timing['time/get_response'] = time.time()-t
#### Compute sentiment score
t = time.time()
texts = [q + r for q,r in zip(batch['query'], batch['response'])]
pipe_outputs = sentiment_pipe(texts, **sent_kwargs)
#[[{'label': 'NEGATIVE', 'score': 0.27862095832824707}, {'label': 'POSITIVE', 'score': -0.5044471621513367}]]
rewards = torch.tensor([output[1]["score"] for output in pipe_outputs]).to(device) #each output has negative score(output[0]) and positive score(output[1])
#如果一个prompt目前是negative,它的positive score是-0.5，那么加到奖励里面，相当于让它少学这个
timing['time/get_sentiment_preds'] = time.time()-t
#### Run PPO step
t = time.time()
stats = ppo_trainer.step(query_tensors, response_tensors, rewards)
timing['time/optimization'] = time.time()-t
#### Log everything
timing['time/epoch'] = time.time()-t0
table_rows = [list(r) for r in zip(batch['query'], batch['response'], rewards.cpu().tolist())]
logs.update({'game_log': wandb.Table(columns=['query', 'response', 'reward'], rows=table_rows)})
logs.update(timing)
logs.update(stats)
logs['env/reward_mean'] = torch.mean(rewards).cpu().numpy()
logs['env/reward_std'] = torch.std(rewards).cpu().numpy()
logs['env/reward_dist'] = rewards.cpu().numpy()
wandb.log(logs)
在训练过程中观察仪表盘，发现reward是上升的，说明训练是有效的。
10.保存模型
gpt2_model.save_pretrained('gpt2-imdb-pos-v2', push_to_hub=True)
gpt2_tokenizer.save_pretrained('gpt2-imdb-pos-v2', push_to_hub=True)
11.模型评估
通过比较原始GPT产生的无限制的文本和微调后产生的受控的文本，我们发现微调过程明显地让模型产生出了正面情感倾向的影评。同样地，在合适的位置添加负号就可以重新训练出一个会产生负面情绪文本的GPT2。针对指定label产生伪数据，这在数据增强上具有很高的应用价值。
此外，本实验的奖励机制是情感倾向值，也可以把奖励机制换成任何你喜欢的评价指标，运用在其他话题的生成任务上，看看模型是否会按照这个方向来学习。
进技术交流群请添加AINLP小助手微信（id: ainlp2)
关于AINLP
AINLP 是一个有趣有AI的自然语言处理社区，专注于 AI、NLP、机器学习、深度学习、推荐算法等相关技术的分享，主题包括文本摘要、智能问答、聊天机器人、机器翻译、自动生成、知识图谱、预训练模型、推荐系统、计算广告、招聘信息、求职经验分享等，欢迎关注！加技术交流群请添加AINLP小助手微信(id：ainlp2)，备注工作/研究方向+加群目的。
阅读至此了，分享、点赞、在看三选一吧🙏