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 withFlag:
                w = w.word
        wc = w.word if allowPOS and withFlag else w
        # 判断词的长度是否小于2，或者词是否为停用词
        if len(wc.strip()) < 2 or wc.lower() in self.stop_words:
            continue
        # 将其添加到词频词典中，次数加1
        freq[w] = freq.get(w, 0.0) + 1.0
    # 统计词频词典中的总次数
    total = sum(freq.values())
    for k in freq:
        kw = k.word if allowPOS and withFlag else k
        # 计算每个词的tf-idf值
        freq[k] *= self.idf_freq.get(kw, self.median_idf) / total
    
    # 根据tf-idf值进行排序
    if withWeight:
        tags = sorted(freq.items(), key=itemgetter(1), reverse=True)
    else:
        tags = sorted(freq, key=freq.__getitem__, reverse=True)
    # 输出topK个词作为关键词
    if topK:
        return tags[:topK]
    else:
        return tags

2.TextRank算法

基于TextRank算法抽取关键词的主调函数是TextRank.textrank函数，主要是在jieba/analyse/textrank.py中实现。

其中，TextRank是为TextRank算法抽取关键词所定义的类。类在初始化时，默认加载了分词函数和词性标注函数tokenizer = postokenizer = jieba.posseg.dt、停用词表stop_words = self.STOP_WORDS.copy()、词性过滤集合pos_filt = frozenset((‘ns’, ‘n’, ‘vn’, ‘v’))，窗口span = 5，((“ns”, “n”, “vn”, “v”))表示词性为地名、名词、动名词、动词。

首先定义一个无向有权图，然后对句子进行分词；依次遍历分词结果，如果某个词i满足过滤条件（词性在词性过滤集合中，并且词的长度大于等于2，并且词不是停用词），然后将这个词之后窗口范围内的词j（这些词也需要满足过滤条件），将它们两两（词i和词j）作为key，出现的次数作为value，添加到共现词典中；

然后，依次遍历共现词典，将词典中的每个元素，key = （词i，词j），value = 词i和词j出现的次数，其中词i，词j作为一条边起始点和终止点，共现的次数作为边的权重，添加到之前定义的无向有权图中。

然后对这个无向有权图进行迭代运算textrank算法，最终经过若干次迭代后，算法收敛，每个词都对应一个指标值；

如果设置了权重标志位，则根据指标值值对无向有权图中的词进行降序排序，最后输出topK个词作为关键词；

def textrank(self, sentence, topK=20, withWeight=False, allowPOS=('ns', 'n', 'vn', 'v'), withFlag=False):

    self.pos_filt = frozenset(allowPOS)
    # 定义无向有权图
    g = UndirectWeightedGraph()
    # 定义共现词典
    cm = defaultdict(int)
    # 分词
    words = tuple(self.tokenizer.cut(sentence))
    # 依次遍历每个词
    for i, wp in enumerate(words):
        # 词i 满足过滤条件
        if self.pairfilter(wp):
            # 依次遍历词i 之后窗口范围内的词
            for j in xrange(i + 1, i + self.span):
                # 词j 不能超出整个句子
                if j >= len(words):
                    break
                # 词j不满足过滤条件，则跳过
                if not self.pairfilter(words[j]):
                    continue
                # 将词i和词j作为key，出现的次数作为value，添加到共现词典中
                if allowPOS and withFlag:
                    cm[(wp, words[j])] += 1
                else:
                    cm[(wp.word, words[j].word)] += 1
    # 依次遍历共现词典的每个元素，将词i，词j作为一条边起始点和终止点，共现的次数作为边的权重
    for terms, w in cm.items():
        g.addEdge(terms[0], terms[1], w)
    
    # 运行textrank算法
    nodes_rank = g.rank()
    
    # 根据指标值进行排序
    if withWeight:
        tags = sorted(nodes_rank.items(), key=itemgetter(1), reverse=True)
    else:
        tags = sorted(nodes_rank, key=nodes_rank.__getitem__, reverse=True)

    # 输出topK个词作为关键词
    if topK:
        return tags[:topK]
    else:
        return tags

其中，无向有权图的的定义及实现是在UndirectWeightedGraph类中实现的。根据UndirectWeightedGraph类的初始化函数__ init__，我们可以发现，所谓的无向有权图就是一个词典，词典的key是后续要添加的词，词典的value，则是一个由（起始点，终止点，边的权重）构成的三元组所组成的列表，表示以这个词作为起始点的所有的边。

无向有权图添加边的操作是在addEdge函数中完成的，因为是无向图，所以我们需要依次将start作为起始点，end作为终止点，然后再将start作为终止点，end作为起始点，这两条边的权重是相同的。

def addEdge(self, start, end, weight):
    # use a tuple (start, end, weight) instead of a Edge object
    self.graph[start].append((start, end, weight))
    self.graph[end].append((end, start, weight))

执行textrank算法迭代是在rank函数中完成的。

首先对每个结点赋予相同的权重，以及计算出该结点的所有出度的次数之和；

然后迭代若干次，以确保得到稳定的结果；

在每一次迭代中，依次遍历每个结点；对于结点n，首先根据无向有权图得到结点n的所有
入度结点（对于无向有权图，入度结点与出度结点是相同的，都是与结点n相连的结点），在前面我们已经计算出这个入度结点的所有出度的次数，而它对于结点n的权值的贡献等于它本身的权值 乘以 它与结点n的共现次数 / 这个结点的所有出度的次数 ，将各个入度结点得到的权值相加，再乘以一定的阻尼系数，即可得到结点n的权值；

迭代完成后，对权值进行归一化，并返回各个结点及其对应的权值。

def rank(self):
    ws = defaultdict(float)
    outSum = defaultdict(float)

    wsdef = 1.0 / (len(self.graph) or 1.0)
    # 初始化各个结点的权值
    # 统计各个结点的出度的次数之和
    for n, out in self.graph.items():
        ws[n] = wsdef
        outSum[n] = sum((e[2] for e in out), 0.0)

    # this line for build stable iteration
    sorted_keys = sorted(self.graph.keys())
    # 遍历若干次
    for x in xrange(10):  # 10 iters
        # 遍历各个结点
        for n in sorted_keys:
            s = 0
            # 遍历结点的入度结点
            for e in self.graph[n]:
                # 将这些入度结点贡献后的权值相加
                # 贡献率 = 入度结点与结点n的共现次数 / 入度结点的所有出度的次数
                s += e[2] / outSum[e[1]] * ws[e[1]]
            # 更新结点n的权值
            ws[n] = (1 - self.d) + self.d * s

    (min_rank, max_rank) = (sys.float_info[0], sys.float_info[3])

    # 获取权值的最大值和最小值
    for w in itervalues(ws):
        if w < min_rank:
            min_rank = w
        if w > max_rank:
            max_rank = w

    # 对权值进行归一化
    for n, w in ws.items():
        # to unify the weights, don't *100.
        ws[n] = (w - min_rank / 10.0) / (max_rank - min_rank / 10.0)

    return ws