筆者開源了一個帶有超詳細中文注釋的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”標記，如下：

因此需要對數據進行清洗，具體如下：

（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”字符

（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）判斷是否存在緩存文件，如果存在，則直接加載；否則重新將文本數據轉換為索引數據，并存為緩存。

（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？
為了更好地區(qū)分Content和Title，并且根據token type id可以僅計算title部分的損失值。

詳細代碼見Github項目的model.py文件。

模型訓練

模型訓練參數如下圖所示：

模型訓練執(zhí)行代碼如下：

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，后續(xù)在選取標記的過程中會取不到值設為無窮小。

（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包含全部的序列序號，則停止預測；否則繼續(xù)預測
            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上，繼續(xù)下一次預測
            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文件。

測試結果如下：

模型上線

通過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

后期工作

可能會將清華新聞數據、搜狗新聞數據等新聞數據集進行整理清洗，構建一個較完善的新聞摘要數據集。

總結