TF-IDF和TextRank算法抽取关键词源码分析

TF-IDF和TextRank算法抽取关键词源码分析

---

jieba分词的关键词抽取功能，是在jieba/analyse目录下实现的。

其中，

• __ init__.py主要用于封装jieba分词的关键词抽取接口；
• tfidf.py实现了基于TF-IDF算法抽取关键词；
• textrank.py实现了基于TextRank算法抽取关键词。

1.TF-IDF算法

基于TF-IDF算法抽取关键词的主调函数是TFIDF.extract_tags函数，主要是在jieba/analyse/tfidf.py中实现。

其中TFIDF是为TF-IDF算法抽取关键词所定义的类。类在初始化时，默认加载了分词函数tokenizer = jieba.dt、词性标注函数postokenizer = jieba.posseg.dt、停用词stop_words = self.STOP_WORDS.copy()、idf词典idf_loader = IDFLoader(idf_path or DEFAULT_IDF)等，并获取idf词典及idf中值（如果某个词没有出现在idf词典中，则将idf中值作为这个词的idf值）。

def __init__(self, idf_path=None):
    # 加载
    self.tokenizer = jieba.dt
    self.postokenizer = jieba.posseg.dt
    self.stop_words = self.STOP_WORDS.copy()
    self.idf_loader = IDFLoader(idf_path or DEFAULT_IDF)
    self.idf_freq, self.median_idf = self.idf_loader.get_idf()

然后开始通过TF-IDF算法进行关键词抽取。

首先根据是否传入了词性限制集合，来决定是调用词性标注接口还是调用分词接口。例如，词性限制集合为[“ns”, “n”, “vn”, “v”, “nr”]，表示只能从词性为地名、名词、动名词、动词、人名这些词性的词中抽取关键词。

• 如果传入了词性限制集合，首先调用词性标注接口，对输入句子进行词性标注，得到分词及对应的词性；依次遍历分词结果，如果该词的词性不在词性限制集合中，则跳过；如果词的长度小于2，或者词为停用词，则跳过；最后将满足条件的词添加到词频词典中，出现的次数加1；然后遍历词频词典，根据idf词典得到每个词的idf值，并除以词频词典中的次数总和，得到每个词的tf * idf值；如果设置了权重标志位，则根据tf-idf值对词频词典中的词进行降序排序，然后输出topK个词作为关键词；
• 如果没有传入词性限制集合，首先调用分词接口，对输入句子进行分词，得到分词；依次遍历分词结果，如果词的长度小于2，或者词为停用词，则跳过；最后将满足条件的词添加到词频词典中，出现的次数加1；然后遍历词频词典，根据idf词典得到每个词的idf值，并除以词频词典中的次数总和，得到每个词的tf * idf值；如果设置了权重标志位，则根据tf-idf值对词频词典中的词进行降序排序，然后输出topK个词作为关键词。

源码分析：

def extract_tags(self, sentence, topK=20, withWeight=False, allowPOS=(), withFlag=False):
    # 传入了词性限制集合
    if allowPOS:
        allowPOS = frozenset(allowPOS)
        # 调用词性标注接口
        words = self.postokenizer.cut(sentence)
    # 没有传入词性限制集合
    else:
        # 调用分词接口
        words = self.tokenizer.cut(sentence)
    freq = {
   }
    for w in words:
        if allowPOS:
            if w.flag not in allowPOS:
                continue
            elif not wi