超详细中文注释的GPT2新闻标题生成项目

笔者开源了一个带有超详细中文注释的GPT2新闻标题生成项目。

该项目参考了GPT2-Chinese、GPT2-chitchat、CDial-GPT、GPT2等多个GPT2开源项目（感谢大佬们的开源），并根据自己的理解，将代码进行重构，添加详细注释，希望可以帮助到有需要的同学。

项目是基于HuggingFace的transformers实现GPT2模型代码进行修改、训练及测试。并且通过Flask框架搭建了一个Web服务，将新闻标题生成模型进行工程化，可以通过页面，可视化地体验新闻标题生成效果。

该项目的目的是带领大家走一遍GPT2生成模型的训练、测试及部署全部流程。

项目地址：https://github.com/liucongg/GPT2-NewsTitle

本文主要是对项目中的代码进行讲解，主要从数据预处理、数据类实现、模型代码实现、模型训练、模型测试和模型上线，六个部分进行介绍，如下。

数据预处理

数据来源于新浪微博，由He Zhengfang大佬整理，详细链接如下：https://www.jianshu.com/p/8f52352f0748?tdsourcetag=s_pcqq_aiomsg。

由于数据来自微博，在标题中常常带有“话题”、“表情”标记，在正文中常常带有“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()

# 定义损失函数CrossEntropyLoss，并且设置忽略计算loss的索引，以及返回loss的形式
# 忽略shift_labels中为0的loss，也就是仅计算title部分的损失值
# 对loss的计算方式设为sum，由于我们仅计算了itle部分的损失值，如果使用mean，会使loss变小（实际除的是sequence_length-1，不是title部分的真实长度）
loss_fct = CrossEntropyLoss(ignore_index=0, reduction="sum")
loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
# 获取title部分的真实长度，并计算真实loss
num = shift_labels.ne(0).long().sum().item()
loss = loss / num
outputs = (loss,) + outputs
return outputs # (loss), lm_logits, presents, (all hidden_states), (attentions)

详细代码见Github项目的model.py文件。

模型训练

模型训练参数如下图所示：

模型训练执行代码如下：

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服务，见：https://zhuanlan.zhihu.com/p/143678340

https://zhuanlan.zhihu.com/p/148224626

后期工作

可能会将清华新闻数据、搜狗新闻数据等新闻数据集进行整理清洗，构建一个较完善的新闻摘要数据集。

总结