Keyword任务_Extraction-bert模型

一. 参考博客or文献

NLP实践教程

二. 知识点

三. 思路与代码

当我们想从特定文档中了解关键信息时，通常会转向关键词提取（keyword extraction）。关键词提取(keyword extraction)是提取与输入文本最相关的词汇、短语的自动化过程。
目前已经有很多的集成工具库类似KeyBERT，已经能够达到非常好的效果，但是在这里我们选择使用BERT创建自己的关键词提取模型。

3.1 思路

• 对于Bert预训练模型, 使用distiluse-base-multilingual-cased, 因为它在相似性任务中表现出了很好的性能, 这也是 Keywords Extraction 任务的关键.
• 由于Transformer模型有token长度限制, 所以在输入大型文档时, 我们需要考虑将文档分割成几个小的段落, 并对其产生的向量进行平均池化(mean pooling, 需要去平均值).
• 通过Bert模型, 可以得到文本的embedding, 此时我们使用文本的embedding与文本标题的embedding进行比较, 因为文章标题与文章的关键字/词一般会有着比较高的相似度, 我们这里目前暂时使用向量之间的余弦相似度来作为相似度的度量方式.
• 最终返回相似度最高的前k个关键字/词.

3.2 数据集描述

需要想讯飞的比赛进行申请, 数据集来源: 基于论文摘要的文本分类与关键词抽取挑战赛

3.2.1 训练集

uuid	title	author	abstract	Keywords	label
…	…	…	…	…	…

训练集的shape为: (6000, 6)

3.2.2 测试集

uuid	title	author	abstract	Keywords
…	…	…	…	…

测试集的shape为: (2358, 5)

3.3 代码与细节

3.3.1 导入前置依赖

# 导入pandas用于读取表格数据
import pandas as pd

# 导入BOW（词袋模型），可以选择将CountVectorizer替换为TfidfVectorizer（TF-IDF（词频-逆文档频率）），注意上下文要同时修改，亲测后者效果更佳
from sklearn.feature_extraction.text import TfidfVectorizer
# 导入Bert模型
from sentence_transformers import SentenceTransformer

# 导入计算相似度前置库，为了计算候选者和文档之间的相似度，我们将使用向量之间的余弦相似度，因为它在高维度下表现得相当好。
from sklearn.metrics.pairwise import cosine_similarity

# 过滤警告消息
from warnings import simplefilter
from sklearn.exceptions import ConvergenceWarning
simplefilter("ignore", category=ConvergenceWarning)

3.3.2 读取数据集并处理

# 读取数据集
test = pd.read_csv('./data/test2.csv')
test['title'] = test['title'].fillna('')
test['abstract'] = test['abstract'].fillna('')
test['text'] = test['title'].fillna('') + ' ' +test['abstract'].fillna('')

# 定义停用词，去掉出现较多，但对文章不关键的词语
stops =[i.strip() for i in open(r'stop.txt',encoding='utf-8').readlines()] 
model = SentenceTransformer(r'xlm-r-distilroberta-base-paraphrase-v1')

3.3.3 提取关键词

test_words = []
for row in test.iterrows():
    # 读取第每一行数据的标题与摘要并提取关键词
    # 修改n_gram_range来改变结果候选词的词长大小。例如，如果我们将它设置为(3，3)，那么产生的候选词将是包含3个关键词的短语。
    n_gram_range = (2,2)
    # 这里我们使用TF-IDF算法来获取候选关键词 
    count = TfidfVectorizer(ngram_range=n_gram_range, stop_words=stops).fit([row[1].text])
    candidates = count.get_feature_names_out()
    # 将文本标题以及候选关键词/关键短语转换为数值型数据（numerical data）。我们使用BERT来实现这一目的
    title_embedding = model.encode([row[1].title])
    
    candidate_embeddings = model.encode(candidates)
    
    # 通过修改这个参数来更改关键词数量
    top_n = 15
    # 利用文章标题进一步提取关键词
    distances = cosine_similarity(title_embedding, candidate_embeddings)
    keywords = [candidates[index] for index in distances.argsort()[0][-top_n:]]
    
    if len( keywords) == 0:
         keywords = ['A', 'B']
    test_words.append('; '.join( keywords))

3.3.4 输出

test['Keywords'] = test_words
test[['uuid', 'Keywords']].to_csv('submit_task2.csv', index=None)