基于 TF-IDF、TextRank 和 Word2Vec 的关键词提取方法对比与实践

在自然语言处理（NLP）任务中，关键词提取是一个非常重要的环节。它能够帮助我们快速理解文本的核心内容，并为后续的文本分类、信息检索、摘要生成等任务提供支持。本文将介绍三种常用的关键词提取方法：TF-IDF、TextRank 和 Word2Vec，并通过实际代码演示它们的实现过程。

---

一、关键词提取的背景与意义

关键词提取是从文本中自动识别出最能代表文本主题的词语或短语。它在以下场景中非常有用：

1. 文本摘要：通过提取关键词，快速生成文本的摘要。
2. 信息检索：帮助搜索引擎更好地理解文档内容，提高检索效率。
3. 文本分类：作为特征输入，提升分类模型的性能。

本文将对比三种常用的关键词提取方法，并通过 Python 实现它们。

二、方法介绍与对比

准备：

使用conda环境,安装以下包：

conda create -n nat_lang python=3.9

conda activate nat_lang
conda install numpy scipy gensim

pip install jieba -i https://mirrors.aliyun.com/pypi/simple/
    
pip install pandas scikit-learn

数据预处理

# data_prepare.py

# 导入工具包
import os
import csv


# 文本文件合并
def text_combine(path):
    # 1. 获取文件列表
    files = []
    for file in os.listdir(path):
        if file.endswith(".txt"):
            files.append(path + "/" + file)
    # 2. 创建text.csv文件，保存结果
    with open('data/text.csv', 'w', newline='',
              encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(['id', 'title', 'abstract'])
        # 3. 遍历txt文件，获取文件编号
        for file_name in files:
            number = (file_name.split('/')[1]).split('_')[0]
            title, text = '', ''
            count = 0
            # 4. 读取标题和内容
            with open(file_name, encoding='utf-8-sig') as f:
                for line in f:
                    if count == 0:
                        title += line.strip()
                    else:
                        text += line.strip()
                    count += 1
            res = [number, title, text]
            writer.writerow(res)


# 主函数处理
def main():
    path = 'text_file'
    text_combine(path)

if __name__ == '__main__':
    main()

1. TF-IDF（词频-逆文档频率）

原理：

TF-IDF 通过计算词语在文档中的词频（TF）和逆文档频率（IDF）来衡量词语的重要性。词频越高且逆文档频率越低的词语，越可能是关键词。

优点：

• 简单高效，适合处理大规模文本数据。
• 能够较好地反映文档的核心内容。

缺点：

• 无法捕捉词语之间的语义关系。
• 对短文本效果较差。

代码实现：

# tfidf.py

import codecs
import pandas as pd
import numpy as np
# 导入jieba分词
import jieba.posseg
import jieba.analyse
# 导入文本向量化函数
from sklearn.feature_extraction.text import TfidfTransformer
# 导入词频统计函数
from sklearn.feature_extraction.text import CountVectorizer


# 读取text.csv文件：分词，去停用词，词性筛选
def data_read(text, stopkey):
    l = []
    pos = ['n', 'nz', 'v', 'vd', 'vn', 'l', 'a', 'd']  # 定义选取的词性
    seg = jieba.posseg.cut(text)  # 分词
    for i in seg:
        if i.word not in stopkey and i.flag in pos:  # 去停用词 + 词性筛选
            l.append(i.word)
    return l


# tf-idf获取文本top10关键词
def words_tfidf(data, stopkey, topK):
    idList, titleList, abstractList = \
        data['id'], data['title'], data['abstract']
    corpus = []  # 将所有文档输出到一个list中，一行就是一个文档
    for index in range(len(idList)):
        # 拼接标题和摘要
        text = '%s。%s' % (titleList[index], abstractList[index])
        text = data_read(text, stopkey)  # 文本预处理
        text = " ".join(text)  # 连接成字符串，空格分隔
        corpus.append(text)

    # 1、构建词频矩阵，将文本中的词语转换成词频矩阵
    vectorizer = CountVectorizer()
    # 词频矩阵,a[i][j]:表示j词在第i个文本中的词频
    X = vectorizer.fit_transform(corpus)
    # 2、统计每个词的tf-idf权值
    transformer = TfidfTransformer()
    tfidf = transformer.fit_transform(X)
    # 3、获取词袋模型中的关键词
    word = vectorizer.get_feature_names_out()
    # 4、获取tf-idf矩阵，a[i][j]表示j词在i篇文本中的tf-idf权重
    weight = tfidf.toarray()
    # 5、打印词语权重
    ids, titles, keys = [], [], []
    for i in range(len(weight)):
        print(u"-------这里输出第", i + 1, u"篇文本的词语tf-idf------")
        ids.append(idList[i])
        titles.append(titleList[i])
        df_word, df_weight = [], []  # 当前文章的所有词汇列表、词汇对应权重列表
        for j in range(len(word)):
            print(word[j], weight[i][j])
            df_word.append(word[j])
            df_weight.append(weight[i][j])
        df_word = pd.DataFrame(df_word, columns=['word'])
        df_weight = pd.DataFrame(df_weight, columns=['weight'])
        word_weight = pd.concat([df_word, df_weight], axis=1)  # 拼接词汇列表和权重列表
        word_weight = word_weight.sort_values(by="weight", ascending=False)  # 按照权重值降序排列
        keyword = np.array(word_weight['word'])  # 选择词汇列并转成数组格式
        word_split = [keyword[x] for x in range(0, topK)]  # 抽取前topK个词汇作为关键词
        word_split = " ".join(word_split)
        keys.append(word_split.encode("utf-8").decode("utf-8"))

    result = pd.DataFrame({"id": ids, "title": titles, "key": keys},
                          columns=['id', 'title', 'key'])
    return result


def main():
    # 读取数据集
    dataFile = 'data/text.csv'
    data = pd.read_csv(dataFile)
    # 停用词表
    stopkey = [w.strip() for w in codecs.open('data/stopWord.txt', 'r', encoding="utf-8").readlines()]
    # tf-idf关键词抽取
    result = words_tfidf(data, stopkey, 10)
    result.to_csv("result/tfidf.csv", index=False)


if __name__ == '__main__':
    main()

2. TextRank

原理：

TextRank 是一种基于图排序的算法，通过构建词语共现图并计算图中节点的权重来提取关键词。

优点：

• 无需训练，适合处理短文本。
• 能够捕捉词语之间的共现关系。

缺点：

• 对长文本效果较差。
• 计算复杂度较高。

代码实现：

# textrank.py

import pandas as pd
import jieba.analyse


# 处理标题和摘要，提取关键词
def words_textrank(data, topK):
    idList, titleList, abstractList = data['id'], data['title'], data['abstract']
    ids, titles, keys = [], [], []
    for index in range(len(idList)):
        # 拼接标题和摘要
        text = '%s。%s' % (titleList[index], abstractList[index])
        jieba.analyse.set_stop_words("data/stopWord.txt") # 加载自定义停用词表
        print("\"", titleList[index], "\"", " 10 Keywords - TextRank :")
        # TextRank关键词提取，词性筛选
        keywords = jieba.analyse.textrank(text, topK=topK,
            allowPOS=('n', 'nz', 'v','vd', 'vn','l', 'a', 'd'))
        word_split = " ".join(keywords)
        keys.append(word_split.encode("utf-8").decode("utf-8"))
        ids.append(idList[index])
        titles.append(titleList[index])

    result = pd.DataFrame({"id": ids, "title": titles,
                           "key": keys},
                          columns=['id', 'title', 'key'])
    return result


def main():
    dataFile = 'data/text.csv'
    data = pd.read_csv(dataFile)
    result = words_textrank(data, 10)
    result.to_csv("result/textrank.csv", index=False)


if __name__ == '__main__':
    main()

3. Word2Vec + K-means

原理：

Word2Vec 是一种词向量模型，能够将词语映射到高维空间，捕捉词语的语义信息。通过 K-means 聚类，可以从候选关键词中提取出最具代表性的词语。

优点：

• 能够捕捉词语的语义信息。
• 适合处理语义相关的任务。

缺点：

• 依赖预训练模型。
• 计算复杂度较高。

代码实现：

# word2vec_prepare.py

import warnings
warnings.filterwarnings(action='ignore',
                        category=UserWarning,
                        module='gensim')  # 忽略警告
import codecs
import pandas as pd
import numpy as np
import jieba  # 分词
import jieba.posseg
import gensim  # 加载词向量模型

# 返回特征词向量bai
def word_vecs(wordList, model):
    name = []
    vecs = []
    for word in wordList:
        word = word.replace('\n', '')
        try:
            if word in model:  # 模型中存在该词的向量表示
                name.append(word.encode('utf8').decode("utf-8"))
                vecs.append(model[word])
        except KeyError:
            continue
    a = pd.DataFrame(name, columns=['word'])
    b = pd.DataFrame(np.array(vecs, dtype='float'))
    return pd.concat([a, b], axis=1)


# 数据预处理操作：分词，去停用词，词性筛选
def data_prepare(text, stopkey):
    l = []
    # 定义选取的词性
    pos = ['n', 'nz', 'v', 'vd', 'vn', 'l', 'a', 'd']
    # 确保输入是字符串类型
    if not isinstance(text, str):
        text = str(text)
    seg = jieba.posseg.cut(text)  # 分词

    for i in seg:
        # 去重 + 去停用词 + 词性筛选
        if i.word not in l and i.word not in stopkey and i.flag in pos:
            # print i.word
            l.append(i.word)
    return l

# 根据数据获取候选关键词词向量
def build_words_vecs(data, stopkey, model):
    idList, titleList, abstractList = data['id'], data['title'], data['abstract']
    for index in range(len(idList)):
        id = idList[index]
        title = titleList[index]
        abstract = abstractList[index]
        l_ti = data_prepare(title, stopkey)  # 处理标题
        l_ab = data_prepare(abstract, stopkey)  # 处理摘要
        # 获取候选关键词的词向量
        words = np.append(l_ti, l_ab)  # 拼接数组元素
        words = list(set(words))  # 数组元素去重,得到候选关键词列表
        wordvecs = word_vecs(words, model)  # 获取候选关键词的词向量表示
        # 词向量写入csv文件，每个词400维
        data_vecs = pd.DataFrame(wordvecs)
        data_vecs.to_csv('result/vecs/wordvecs_' + str(id) + '.csv', index=False)
        print ("document ", id, " well done.")

def main():
    # 读取数据集
    dataFile = 'data/text.csv'
    data = pd.read_csv(dataFile)

    # # 检查是否有 NaN 值
    # print("Number of NaN values in 'title':", data['title'].isnull().sum())
    # print("Number of NaN values in 'abstract':", data['abstract'].isnull().sum())
    #
    # # 将 NaN 值替换为空字符串
    # data['title'] = data['title'].fillna('')
    # data['abstract'] = data['abstract'].fillna('')

    # 停用词表
    stopkey = [w.strip() for w in codecs.open('data/stopWord.txt', 'r', encoding='utf-8').readlines()]
    # 词向量模型
    inp = 'wiki.zh.text.vector'
    model = gensim.models.KeyedVectors.load_word2vec_format(inp, binary=False)
    build_words_vecs(data, stopkey, model)


if __name__ == '__main__':
    main()

word2vec_result将word2vec_prepare文件构建的词向量进行提取

 # word2vec_result.py

import os
# 导入kmeans聚类算法
from sklearn.cluster import KMeans
import pandas as pd
import numpy as np
import math


# 对词向量采用K-means聚类抽取TopK关键词
def words_kmeans(data, topK):
    words = data["word"]  # 词汇
    vecs = data.iloc[:, 1:]  # 向量表示

    kmeans = KMeans(n_clusters=1, random_state=10).fit(vecs)
    labels = kmeans.labels_  # 类别结果标签
    labels = pd.DataFrame(labels, columns=['label'])
    new_df = pd.concat([labels, vecs], axis=1)
    vec_center = kmeans.cluster_centers_  # 聚类中心

    # 计算距离（相似性） 采用欧几里得距离（欧式距离）
    distances = []
    vec_words = np.array(vecs)  # 候选关键词向量，dataFrame转array
    vec_center = vec_center[0]  # 第一个类别聚类中心,本例只有一个类别
    length = len(vec_center)  # 向量维度
    for index in range(len(vec_words)):  # 候选关键词个数
        cur_wordvec = vec_words[index]  # 当前词语的词向量
        dis = 0  # 向量距离
        for index2 in range(length):
            dis += (vec_center[index2] - cur_wordvec[index2]) * \
                   (vec_center[index2] - cur_wordvec[index2])
        dis = math.sqrt(dis)
        distances.append(dis)
    distances = pd.DataFrame(distances, columns=['dis'])
    # 拼接词语与其对应中心点的距离
    result = pd.concat([words, labels, distances], axis=1)
    # 按照距离大小进行升序排序
    result = result.sort_values(by="dis", ascending=True)

    # 抽取排名前topK个词语作为文本关键词
    wordlist = np.array(result['word'])

    # 检查 wordlist 的长度是否足够
    actual_topK = min(topK, len(wordlist))  # 取 topK 和 wordlist 长度的较小值
    word_split = [wordlist[x] for x in range(0, actual_topK)]  # 提取前 actual_topK 个词汇
    word_split = " ".join(word_split)
    return word_split


def main():
    # 读取数据集
    dataFile = 'data/text.csv'
    articleData = pd.read_csv(dataFile)
    ids, titles, keys = [], [], []

    rootdir = "result/vecs"  # 词向量文件根目录
    fileList = os.listdir(rootdir)  # 列出文件夹下所有的目录与文件
    # 遍历文件
    for i in range(len(fileList)):
        filename = fileList[i]
        path = os.path.join(rootdir, filename)
        if os.path.isfile(path):
            # 读取词向量文件数据
            data = pd.read_csv(path, encoding='utf-8')
            # 聚类算法得到当前文件的关键词
            artile_keys = words_kmeans(data, 5)
            # 根据文件名获得文章id以及标题
            (shortname, extension) = os.path.splitext(filename)
            t = shortname.split("_")
            article_id = int(t[len(t) - 1])  # 获得文章id
            # 获得文章标题
            artile_tit = articleData[articleData.id ==
                                     article_id]['title']
            print(artile_tit)
            print(list(artile_tit))
            artile_tit = list(artile_tit)[0]  # series转成字符串
            ids.append(article_id)
            titles.append(artile_tit)
            keys.append(artile_keys.encode("utf-8").decode("utf-8"))
    # 所有结果写入文件
    result = pd.DataFrame({"id": ids, "title": titles, "key": keys},
                          columns=['id', 'title', 'key'])
    result = result.sort_values(by="id", ascending=True)  # 排序
    result.to_csv("result/word2vec.csv", index=False,
                  encoding='utf_8_sig')


if __name__ == '__main__':
    main()

三、实验与结果分析

1. 数据准备

使用 data_prepare.py 脚本将多个 .txt 文件合并为一个 text.csv 文件。每个 .txt 文件包含文章的标题和摘要，合并后的 text.csv 文件包含三列：id、title 和 abstract。

2. 实验结果

• TF-IDF：提取的关键词基于词频和逆文档频率，能够较好地反映文章的核心内容。
• TextRank：提取的关键词基于图排序算法，能够捕捉词语之间的共现关系。
• Word2Vec：提取的关键词基于词向量和聚类算法，能够捕捉词语的语义信息。

3. 方法对比

方法	优点	缺点	适用场景
TF-IDF	简单高效，适合大规模文本	无法捕捉语义关系	大规模文本数据
TextRank	无需训练，适合短文本	对长文本效果较差	短文本
Word2Vec	能够捕捉语义信息	依赖预训练模型，计算复杂度高	语义相关任务

五、总结

本文介绍了三种常用的关键词提取方法：TF-IDF、TextRank 和 Word2Vec，并通过 Python 实现了它们。每种方法都有其优缺点，适用于不同的场景。在实际应用中，可以根据任务需求选择合适的方法，或结合多种方法提高效果。

资料

通过网盘分享的文件：3.3.2-关键词提取综合案例.zip
链接: https://pan.baidu.com/s/1Nbg1fCKYLAhxRVB9jgRM_Q 提取码: w89q

希望本文对你理解关键词提取有所帮助！如果你有任何问题或建议，欢迎在评论区留言。