超详细中文注释的GPT2新闻标题生成项目：https://zhuanlan.zhihu.com/p/338171330
笔者开源了一个带有超详细中文注释的GPT2新闻标题生成项目。
该项目参考了GPT2-Chinese、GPT2-chitchat、CDial-GPT、GPT2等多个GPT2开源项目（感谢大佬们的开源），并根据自己的理解，将代码进行重构，添加详细注释，希望可以帮助到有需要的同学。
项目是基于HuggingFace的transformers实现GPT2模型代码进行修改、训练及测试。并且通过Flask框架搭建了一个Web服务，将新闻标题生成模型进行工程化，可以通过页面，可视化地体验新闻标题生成效果。
该项目的目的是带领大家走一遍GPT2生成模型的训练、测试及部署全部流程。
项目地址：
https://github.com/liucongg/GPT2-NewsTitle github.com本文主要是对项目中的代码进行讲解，主要从数据预处理、数据类实现、模型代码实现、模型训练、模型测试和模型上线，六个部分进行介绍，如下。
数据预处理
数据来源于新浪微博，由He Zhengfang大佬整理，详细链接如下：中文短文本摘要数据集。
由于数据来自微博，在标题中常常带有“话题”、“表情”标记，在正文中常常带有“HTML”标记，如下：
Title：
2014#福布斯中国名人榜#：她再夺冠[威武]
Content：
为什么我们要工作？听演讲者Barry Schwartz告诉你工作的另一个重要意义。非常有深度的一个演讲，值得一看！http://t.cn/RqzKvtn 转发学习，给自己的工作加油打气吧！[good] 因此需要对数据进行清洗，具体如下：
（1）对标题清洗时，会去除“##”符号（一般为微博数据的话题标记）、去除“[]”中间的文字（一般为微博数据中的表情）、合并标题中过多的空格
def clean_weibo_title(title: str):
"""
对微博数据中的标题内容（待生成）进行清洗
Args:
title: 标题
Returns:
"""
# 去除##符号（一般为微博数据的话题标记）
title = re.sub(r"#", "", title)
# 去除[]中间的文字（一般为微博数据中的表情）
title = re.sub(r"(\[{1,2})(.*?)(\]{1,2})", "", title)
# 合并标题中过多的空格
title = re.sub(r"\s+", " ", title)
return title（2）对正文清洗时，会去除网址、合并正文中过多的空格、去除“\u200b”字符
def clean_weibo_content(content: str):
"""
对微博数据中的文本内容进行清洗
Args:
content: 文本
Returns:
"""
# 去除网址
content = re.sub(r"(https|http)?:\/\/(\w|\.|\/|\?|\=|\&|\%)*\b", "", content)
# 合并正文中过多的空格
content = re.sub(r"\s+", " ", content)
# 去除\u200b字符
content = content.replace("\u200b", "")
return content（3）对清洗后的数据进行整合，去除重复数据、正文内容字数小于100的数据和标题内容字数小于2的数据；并且拆分训练集和测试集。
def build_news_data(content_path, title_path, train_save_path, test_save_path):
"""
对微博数据进行清洗，构建训练集和测试集
Args:
content_path: 正文内容文件路径
title_path: 标题内容文件路径
train_save_path: 训练集文件路径
test_save_path: 测试集文件路径
Returns:
"""
# 打开文件，并将其zip成一个文件
content_data = open(content_path, "r", encoding="utf-8")
title_data = open(title_path, "r", encoding="utf-8")
data = zip(content_data.readlines(), title_data.readlines())
# 使用多进程处理数据
threads = min(8, cpu_count())
with Pool(threads) as p:
annoate_ = partial(clean_data)
data = list(tqdm(p.imap(annoate_, data, chunksize=8),
desc="build data"
)
)
# 对数据进行过滤，去除重复数据、正文内容字长小于100的数据和标题内容字长小于100的数据
data_set = set()
data_new = []
for d in data:
if d["content"] in data_set or len(d["content"]) < 100 or len(d["title"]) < 2:
continue
else:
data_set.add(d["content"])
data_new.append(d)
# 拆分数据，构建训练集和测试集
random.shuffle(data_new)
train_data = data_new[:-3000]
test_data = data_new[-3000:]
fin = open(train_save_path, "w", encoding="utf-8")
fin.write(json.dumps(train_data, indent=4, ensure_ascii=False))
fin.close()
fin = open(test_save_path, "w", encoding="utf-8")
fin.write(json.dumps(test_data, indent=4, ensure_ascii=False))
fin.close()详细代码见Github项目的data_helper.py文件。
数据类实现
数据类的作用是将文本数据转换成模型可以使用的索引数据，并预先存储下来。避免模型每训练一步，都进行无效的数据转换操作。
（1）判断是否存在缓存文件，如果存在，则直接加载；否则重新将文本数据转换为索引数据，并存为缓存。
if os.path.exists(cached_feature_file) and not is_overwrite:
logger.info("已经存在缓存文件{}，直接加载".format(cached_feature_file))
self.data_set = torch.load(cached_feature_file)["data_set"]
# 如果缓存数据不存在，则对原始数据进行数据处理操作，并将处理后的数据存成缓存文件
else:
logger.info("不存在缓存文件{}，进行数据预处理操作".format(cached_feature_file))
self.data_set = self.load_data(path_file)
logger.info("数据预处理操作完成，将处理后的数据存到{}中，作为缓存文件".format(cached_feature_file))
torch.save({"data_set": self.data_set}, cached_feature_file)（2）将文本数据转换为索引数据的函数
def convert_feature(self, sample):
"""
数据处理函数
Args:
sample: 一个字典，包含新闻的正文和新闻的标题，格式为{"content": content, "title": title}
Returns:
"""
input_ids = []
token_type_ids = []
# 对新闻正文进行tokenizer.tokenize分词
content_tokens = self.tokenizer.tokenize(sample["content"])
# 对新闻标题进行tokenizer.tokenize分词，注意tokenizer中已经将[Space]作为一个分隔符，不会切割成多个字符
title_tokens = self.tokenizer.tokenize(sample["title"].replace(" ", "[Space]"))
# 判断如果正文过长，进行截断
if len(content_tokens) > self.max_len - len(title_tokens) - 3:
content_tokens = content_tokens[:self.max_len - len(title_tokens) - 3]
# 生成模型所需的input_ids和token_type_ids
input_ids.append(self.tokenizer.cls_token_id)
token_type_ids.append(self.content_id)
input_ids.extend(self.tokenizer.convert_tokens_to_ids(content_tokens))
token_type_ids.extend([self.content_id] * len(content_tokens))
input_ids.append(self.tokenizer.sep_token_id)
token_type_ids.append(self.content_id)
input_ids.extend(self.tokenizer.convert_tokens_to_ids(title_tokens))
token_type_ids.extend([self.title_id] * len(title_tokens))
input_ids.append(self.tokenizer.sep_token_id)
token_type_ids.append(self.title_id)
# 判断input_ids与token_type_ids长度是否一致
assert len(input_ids) == len(token_type_ids)
# 判断input_ids长度是否小于等于最大长度
assert len(input_ids) <= self.max_len
return input_ids, token_type_ids详细代码见Github项目的data_set.py文件。
模型代码实现
模型部分，主要对transformers包中GPT2LMHeadModel类进行重写，修改计算loss部分，只计算预测title部分的loss。
模型的输入由word embedding、segment embedding和position embedding三部分组成，具体如下图所示：
为什么需要加segment embedding？
为了更好地区分Content和Title，并且根据token type id可以仅计算title部分的损失值。
def forward(self, input_ids=None, past=None, token_type_ids=None, labels=None, title_id=None): """ 前向函数，计算GPT2预测结果值 Args: input_ids: 输入序列在词表中的索引序列，size:[batch_size, sequence_length] past: 包含由模型预先计算好的隐藏状态，一般使用在预测阶段，用于加速顺序解码，防止重复计算前面计算过的token token_type_ids: 用于区分输入序列中content和title的分隔符序列，size:[batch_size, sequence_length] labels: 标签序列，size:[batch_size, sequence_length]，一般情况下，与input_ids相同 title_id: title部分分隔符的id Returns: """ # 获取GPT2模型的输出结果 transformer_outputs = self.transformer(input_ids, past=past, token_type_ids=token_type_ids) # 获取GPT2模型的最后一层的隐层节点状态，size:[batch_size, sequence_length, config.n_embd] hidden_states = transformer_outputs[0] # 预测隐层节点状态中的每一个token的下一个token，size:[batch_size, sequence_length, config.vocab_size] lm_logits = self.lm_head(hidden_states) # 拼接输出结果 outputs = (lm_logits,) + transformer_outputs[1:] # 如果labels不为None时，计算损失值loss，并拼接到输出结果中 if labels is not None: # 计算loss时，title_id不可以为None，因为需要title_id找到title的部分 if title_id is None or token_type_ids is None: raise Exception("当labels不为None时， title_id和token_type_ids均不可以为None。") # 获取mask值，如果token_type_ids中等于title_id的部分需要计算loss，标记为1；否则为0。 # size:[batch_size, sequence_length] mask = (token_type_ids == title_id).long() # 获取新的标签，size:[batch_size, sequence_length] labels = labels * mask # 对预测结果和标签进行偏移操作 # GPT2的生成机制为通过前面的token，预测下一个token；并且labels与input_ids相同， # 因此input_ids中的第一个token的预测结果，实际上是标签中的第二个token，以此类推，最终仅计算sequence_length-1个token的loss shift_logits = lm_logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous()<span class="c1"># 定义损失函数CrossEntropyLoss，并且设置忽略计算loss的索引，以及返回loss的形式</span> <span class="c1"># 忽略shift_labels中为0的loss，也就是仅计算title部分的损失值</span> <span class="c1"># 对loss的计算方式设为sum，由于我们仅计算了itle部分的损失值，如果使用mean，会使loss变小（实际除的是sequence_length-1，不是title部分的真实长度）</span> <span class="n">loss_fct</span> <span class="o">=</span> <span class="n">CrossEntropyLoss</span><span class="p">(</span><span class="n">ignore_index</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">reduction</span><span class="o">=</span><span class="s2">"sum"</span><span class="p">)</span> <span class="n">loss</span> <span class="o">=</span> <span class="n">loss_fct</span><span class="p">(</span><span class="n">shift_logits</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="n">shift_logits</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)),</span> <span class="n">shift_labels</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">))</span> <span class="c1"># 获取title部分的真实长度，并计算真实loss</span> <span class="n">num</span> <span class="o">=</span> <span class="n">shift_labels</span><span class="o">.</span><span class="n">ne</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">long</span><span class="p">()</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span><span class="o">.</span><span class="n">item</span><span class="p">()</span> <span class="n">loss</span> <span class="o">=</span> <span class="n">loss</span> <span class="o">/</span> <span class="n">num</span> <span class="n">outputs</span> <span class="o">=</span> <span class="p">(</span><span class="n">loss</span><span class="p">,)</span> <span class="o">+</span> <span class="n">outputs</span> <span class="k">return</span> <span class="n">outputs</span> <span class="c1"># (loss), lm_logits, presents, (all hidden_states), (attentions)</span></code></pre></div><p>详细代码见Github项目的model.py文件。</p><h2>模型训练</h2><p>模型训练参数如下图所示：</p><figure data-size="normal"><noscript><img src="https://pic4.zhimg.com/v2-c0a650d43fefb2ac864fca4facfa3f87_b.jpg" data-caption="" data-size="normal" data-rawwidth="948" data-rawheight="898" class="origin_image zh-lightbox-thumb" width="948" data-original="https://pic4.zhimg.com/v2-c0a650d43fefb2ac864fca4facfa3f87_r.jpg"/></noscript><img src="data:image/svg+xml;utf8,<svg xmlns='http://www.w3.org/2000/svg' width='948' height='898'></svg>" data-caption="" data-size="normal" data-rawwidth="948" data-rawheight="898" class="origin_image zh-lightbox-thumb lazy" width="948" data-original="https://pic4.zhimg.com/v2-c0a650d43fefb2ac864fca4facfa3f87_r.jpg" data-actualsrc="https://pic4.zhimg.com/v2-c0a650d43fefb2ac864fca4facfa3f87_b.jpg"></figure><p>模型训练执行代码如下：</p><div class="highlight"><pre><code class="language-text">python3 train.py
或
python3 train.py --output_dir output_dir/(自定义保存模型路径)
模型训练文件主要由以下几个函数组成：（1）设置训练模型所需参数函数set_args；（2）训练模型函数train；（3）对测试数据集进行模型测试evaluate；（4）主函数main。
详细代码见Github项目的train.py文件。
值得注意的是，在实例化tokenizer时，一定要使用tokenizer.add_tokens("[Space]", special_tokens=True)，目的是为了将[Space]作为一个切分整体，例如：“我爱[Space]北京*。”，使用原始tokenizer分词结果为"[‘我’, ‘爱’, ‘[’, ‘Space’, ‘]’, ‘北’, ‘京’, ‘天’, ‘安’,‘门’,’。’]"；增加切分符号后的结果为"[‘我’, ‘爱’, ‘[Space]’, ‘北’, ‘京’, ‘天’, ‘安’,‘门’,’。’]"。
模型测试
模型测试部分，主要是通过不同的解码策略，对已经训练好的模型进行单个样本的预测。
（1）top_k或top_p解码策略，仅保留top_k个或累积概率到达top_p的标记，其他标记设为filter_value，后续在选取标记的过程中会取不到值设为无穷小。
def top_k_top_p_filtering(logits, top_k, top_p, filter_value=-float(“Inf”)):
“”"
top_k或top_p解码策略，仅保留top_k个或累积概率到达top_p的标记，其他标记设为filter_value，后续在选取标记的过程中会取不到值设为无穷小。
Args:
logits: 预测结果，即预测成为词典中每个词的分数
top_k: 只保留概率最高的top_k个标记
top_p: 只保留概率累积达到top_p的标记
filter_value: 过滤标记值
Returns:
“”"
# logits的维度必须为2，即size:[batch_size, vocab_size]
assert logits.dim() 2
# 获取top_k和字典大小中较小的一个，也就是说，如果top_k大于字典大小，则取字典大小个标记
top_k = min(top_k, logits[0].size(-1))
# 如果top_k不为0，则将在logits中保留top_k个标记
if top_k > 0:
# 由于有batch_size个预测结果，因此对其遍历，选取每个预测结果的top_k标记
for logit in logits:
indices_to_remove = logit < torch.topk(logit, top_k)[0][…, -1, None]
logit[indices_to_remove] = filter_value
# 如果top_p不为0，则将在logits中保留概率值累积达到top_p的标记
if top_p > 0.0:
# 对logits进行递减排序
sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
# 对排序后的结果使用softmax归一化，再获取累积概率序列
# 例如：原始序列[0.1, 0.2, 0.3, 0.4]，则变为：[0.1, 0.3, 0.6, 1.0]
cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
# 删除累积概率高于top_p的标记
sorted_indices_to_remove = cumulative_probs > top_p
# 将索引向右移动，使第一个标记也保持在top_p之上
sorted_indices_to_remove[…, 1:] = sorted_indices_to_remove[…, :-1].clone()
sorted_indices_to_remove[…, 0] = 0
for index, logit in enumerate(logits):
# 由于有batch_size个预测结果，因此对其遍历，选取每个预测结果的累积概率达到top_p的标记
indices_to_remove = sorted_indices[index][sorted_indices_to_remove[index]]
logit[indices_to_remove] = filter_value
return logits（2）对单个样本进行预测
def predict_one_sample(model, tokenizer, device, args, content):
“”"
对单个样本进行预测
Args:
model: 模型
tokenizer: 分词器
device: 设备信息
args: 配置项信息
content: 新闻正文
Returns:
“”"
# 对新闻正文进行预处理，并判断如果超长则进行截断
content_tokens = tokenizer.tokenize(content)
if len(content_tokens) > args.max_len - 3 - args.generate_max_len:
content_tokens = content_tokens[:args.max_len - 3 - args.generate_max_len]
# 获取content_id、title_id、unk_id、sep_id值
content_id = tokenizer.convert_tokens_to_ids("[Content]")
title_id = tokenizer.convert_tokens_to_ids("[Title]")
unk_id = tokenizer.convert_tokens_to_ids("[UNK]")
sep_id = tokenizer.convert_tokens_to_ids("[SEP]")
# 将tokens索引化，变成模型所需格式
content_tokens = ["[CLS]"] + content_tokens + ["[SEP]"]
input_ids = tokenizer.convert_tokens_to_ids(content_tokens)
# 将input_ids和token_type_ids进行扩充，扩充到需要预测标题的个数，即batch_size
input_ids = [copy.deepcopy(input_ids) for in range(args.batch_size)]
token_type_ids = [[content_id] * len(content_tokens) for in range(args.batch_size)]
# 将input_ids和token_type_ids变成tensor
input_tensors = torch.tensor(input_ids).long().to(device)
token_type_tensors = torch.tensor(token_type_ids).long().to(device)
next_token_type = torch.tensor([[title_id] for in range(args.batch_size)]).long().to(device)
# 用于存放每一步解码的结果
generated = []
# 用于存放，完成解码序列的序号
finish_set = set()
with torch.no_grad():
# 遍历生成标题最大长度
for in range(args.generate_max_len):
outputs = model(input_ids=input_tensors, token_type_ids=token_type_tensors)
# 获取预测结果序列的最后一个标记，next_token_logits size：[batch_size, vocab_size]
next_token_logits = outputs[0][:, -1, :]
# 对batch_size进行遍历，将词表中出现在序列中的词的概率进行惩罚
for index in range(args.batch_size):
for token_id in set([token_ids[index] for token_ids in generated]):
next_token_logits[index][token_id] /= args.repetition_penalty
# 对batch_size进行遍历，将词表中的UNK的值设为无穷小
for next_token_logit in next_token_logits:
next_token_logit[unk_id] = -float(“Inf”)
# 使用top_k_top_p_filtering函数，按照top_k和top_p的值，对预测结果进行筛选
filter_logits = top_k_top_p_filtering(next_token_logits, top_k=args.top_k, top_p=args.top_p)
# 对filter_logits的每一行做一次取值，输出结果是每一次取值时filter_logits对应行的下标，即词表位置（词的id）
# filter_logits中的越大的值，越容易被选中
next_tokens = torch.multinomial(F.softmax(filter_logits, dim=-1), num_samples=1)
# 判断如果哪个序列的预测标记为sep_id时，则加入到finish_set
for index, token_id in enumerate(next_tokens[:, 0]):
if token_id sep_id:
finish_set.add(index)
# 判断，如果finish_set包含全部的序列序号，则停止预测；否则继续预测
finish_flag = True
for index in range(args.batch_size):
if index not in finish_set:
finish_flag = False
break
if finish_flag:
break
# 将预测标记添加到generated中
generated.append([token.item() for token in next_tokens[:, 0]])
# 将预测结果拼接到input_tensors和token_type_tensors上，继续下一次预测
input_tensors = torch.cat((input_tensors, next_tokens), dim=-1)
token_type_tensors = torch.cat((token_type_tensors, next_token_type), dim=-1)
# 用于存储预测结果
candidate_responses = []
# 对batch_size进行遍历，并将token_id变成对应汉字
for index in range(args.batch_size):
responses = []
for token_index in range(len(generated)):
# 判断，当出现sep_id时，停止在该序列中添加token
if generated[token_index][index] != sep_id:
responses.append(generated[token_index][index])
else:
break
# 将token_id序列变成汉字序列，去除"##"，并将[Space]替换成空格
candidate_responses.append(
“”.join(tokenizer.convert_ids_to_tokens(responses)).replace("##", “”).replace("[Space]", " “))
return candidate_responses详细代码见Github项目的generate_title.py文件。
测试结果如下：
从测试集中抽一篇
content：
今日，中国三条重要高铁干线——兰新高铁、贵广铁路和南广铁路将开通运营。其中兰新高铁是中国首条高原高铁，全长1776公里，最高票价658元。贵广铁路最贵车票320元，南广铁路最贵车票206.5元，这两条线路大大缩短西南与各地的时空距离。出行更方便了！中国“高铁版图”再扩容 三条重要高铁今日开通
title：
生成的第1个标题为：中国“高铁版图”再扩容 三条重要高铁今日开通
生成的第2个标题为：贵广铁路最高铁版图
生成的第3个标题为：出行更方便了！中国“高铁版图”再扩容三条重要高铁今日开通模型上线
通过Flask框架搭建了一个Web服务，将新闻摘要生成模型进行工程化，可以通过页面可视化地体验新闻摘要生成效果。
详细代码见Github项目的http_server.py文件。
并且在我之前文章中，详细介绍过如何使用Flask框架搭建Web服务，见：
刘聪NLP：Web服务部署深度学习模型 zhuanlan.zhihu.com启动服务命令：
python3 http_server.py
或
python3 http_server.py --http_id “0.0.0.0” --port 5555本地测试直接使用"127.0.0.1:5555/news-title-generate”，如果给他人访问，只需将"127.0.0.1"替换成的电脑的IP地址即可。
初始页面如下图所示：
输入新闻正文后，点击“一键生成”，可以获取到生成的新闻标题，如下图所示：
后期工作
可能会将清华新闻数据、搜狗新闻数据等新闻数据集进行整理清洗，构建一个较完善的新闻摘要数据集。
可能会使用新闻数据训练一个小的GPT2预训练模型。
可能会对已上传的新闻标题生成模型进行更新，训练一个效果较好的模型。
总结
GPT2模型已经非常成熟，也有很多很好的开源项目。笔者本着开源之心，将代码进行整理，增加详细注释，希望可以帮助大家更好地理解代码。也欢迎大家留言讨论。
鸽了这么久，终于回来了~~~~~~
其他文章推荐：
刘聪NLP：MacBERT：MLM as correction BERT
刘聪NLP：BERT-QE: 基于上下文化查询扩展的文档ReRank
刘聪NLP：SIGIR 2020之MarkedBERT模型：加入传统检索线索的Rerank模型
刘聪NLP：SIGIR 2020之DC-BERT模型：解耦问题-文档编码，提速QA-Rerank模块
刘聪NLP：ACL2020论文整理之问题生成、自然语言推理、预训练语言模型及部分应用、QA问答系统及机器阅读理解