超详细中文注释的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:///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

本地测试直接使用'127.0.0.1:5555/news-title-generate'，如果给他人访问，只需将'127.0.0.1'替换成的电脑的IP地址即可。

初始页面如下图所示：

输入新闻正文后，点击“一键生成”，可以获取到生成的新闻标题，如下图所示：

后期工作

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

可能会使用新闻数据训练一个小的GPT2预训练模型。

可能会对已上传的新闻标题生成模型进行更新，训练一个效果较好的模型。

总结

GPT2模型已经非常成熟，也有很多很好的开源项目。笔者本着开源之心，将代码进行整理，增加详细注释，希望可以帮助大家更好地理解代码。也欢迎大家留言讨论。

参考

GPT2-Chinese：https://github.com/Morizeyao/GPT2-Chinese

GPT2-chitchat：https://github.com/yangjianxin1/GPT2-chitchat

CDial-GPT：https://github.com/thu-coai/CDial-GPT

GPT2：https://github.com/ConnorJL/GPT2

transformers：https://github.com/huggingface/transformers