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