手写AI算法之TF-IDF关键词提取

TF-IDF（term frequency–inverse document frequency）是一种常用于NLP的统计方法，用以评估一字词，对于一个语料库中的某一份文档的重要程度。该词重要性，跟它在文件中出现的次数成正比，但同时跟它在语料库中出现的频率成反比。本文介绍使用TF-IDF提取文档关键词的方法。

一、公式表达

tfidf(w) = tf(d, w) * idf(w)
tf(d, w) = Nw/N  # 词频 = 词在文档中出现的次数 / 文档总词数
idf(w) = lg(D/Dw)  # 逆文档频率 = lg(语料库的文档总数 / 包含改词的文档数 + 1)

二、优缺点

1、TF-IDF算法的优点是简单快速，结果比较符合实际情况。

2、单纯以"词频"衡量一个词的重要性，不够全面，有时重要的词可能出现次数并不多。

3、该算法无法体现词的位置信息，出现位置靠前的词与出现位置靠后的词，都被视为重要性相同。

三、代码示例

1、定义语料并分词

docA = 'the cat sat on my bed'
docB = 'the dog sat on my knees'
# 分词
bowA = docA.split()
bowB = docB.split()

2、定义TF-IDF处理类

class TfidfCounter():
   def __init__(self) -> None:
       self.termset = set()  #('a','b','c')
       self.tfs = {}  #{'A':{'a':0.1, 'b':0.2}, 'B':{'b':0.3, 'c':0.2}}
       self.idfs = {}  #{'a':0.1, 'b':0.2, 'c':0.5}
       self.tfidfs = {}  #{'A':{'a':0.1, 'b':0.2, 'c':0.1}, 'B':{'a':0.3, 'b':0.3, 'c':0.2}}

3、实例化并添加语料

counter = TfidfCounter()
counter.add('a', bowA)
counter.add('b', bowB)

def add(self, docid, term_list):
   self.termset = self.termset.union(term_list)

4、计算TF值

def add(self, docid, term_list):
   ......
   self.compute_tf(docid, term_list)

def compute_tf(self, docid, term_list):
   count = len(term_list)
   tf = {}
   for term in term_list:
       tf[term] = term_list.count(term) / count
   self.tfs[docid] = tf

5、计算IDF值

def compute(self):
   self.compute_idf()
   self.compute_tfidf()

def compute_idf(self):
   total = len(self.tfs)
   for term in self.termset:
       count = 0
       for tfs in self.tfs.values():
           if term in tfs.keys():
               count += 1
       self.idfs[term] = np.log10((total + 1) / (count))

6、计算TF-IDF并提取关键词

def compute_tfidf(self):
   for docid, tfs in self.tfs.items():
       tfidf = {}
       for term, tf in tfs.items():
           tfidf[term] = tf * self.idfs[term]
       self.tfidfs[docid] = dict(sorted(tfidf.items(), key=lambda x:x[1], reverse=True))

tfidfs_dict = counter.get_tfidf()
for docid, tfidf in tfidfs_dict.items():
   #获取topN个关键词
   print(list(tfidf.keys())[:5])

四、更多运用

除了本文介绍的《TF-IDF关键词提取》运用外，还有两个延伸运用《TF-IDF余弦相似找相似文章》和《TF-IDF通过词频对文章自动摘要》。