超详细中文注释的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
python3 http_server.py
或
python3 http_server.py --http_id '0.0.0.0' --port 5555后期工作
可能会将清华新闻数据、搜狗新闻数据等新闻数据集进行整理清洗，构建一个较完善的新闻摘要数据集。