实战小项目：使用 TF-IDF 算法提取文章关键词
实战小项目：使用 TF-IDF 算法提取文章关键词
1 背景描述
提取文本关键词是很常见的一个需求，比较常见简单的算法，像 TF-IDF 就可以用来关键词提取。
在 Python 中有很多库都实现了这个算法，如果仅仅是为了做一些实验研究使用python中的库来作为提取算法是比较便捷的方式，但是如果是应用到生产环境中 python 将会有很多限制，比如需要将提取关键词算法部署到服务器中，并提供一个 Rest API 。
本篇文章，提供另一种常用语言的实现思路。
Java 是目前 Web 应用中常用的语言，其性能、兼容性、稳定性是经得住长时间考验的。
本文为了符合生产环境的需求，选用 Java 作为开发语言，开发一个能够应用到服务器环境中的提取关键词应用，而不是仅仅停留在实验阶段。
关于 TF-IDF 算法原理很多博客写的都非常的棒，本文不会对原理有详细的阐述，而是具体的描述如何将公式算法使用 Java 语言实现出来。
2 计算TF-IDF步骤
TF-IDF 是衡量某个词的重要程度的一个指标，因此如果想要提取某个文档的关键词，只需要把这个文档分词，然后对所有词的 TF-IDF 排序，TF-IDF 越大，权值越高，说明越重要，通过这个思路就可以提取出这个文档的关键词。
2.1 首先，计算词频
对公式 1 的分子来说，只要把语料库中的文档分词，然后统计不同词的个数。对于分母来说，直接求 cout 就行
getDataSource 方法是获取语料库的方法，对于中文语料库来说，你可以使用结巴分词，对不同的文档进行分词，然后放到List中。
因此我们的语料库就成了一个二维数组，第一维度是文档，第二维度是这个文档中的单个词，当然这个二维数组很有可能是长度不一的。
List<List<String>> docs=getDataSource()
使用 Java8 的新特性能够抛弃臃肿的 for 循环，因此本文使用 Java8 来简化代码。
首先定义一个 wordList 变量，把所有单词不去除重复都存起来，用来作为统计词频的数组。
List<String> wordList = new ArrayList<>();
//填充数组
docs.stream().forEach(tokens -> {
wordList.addAll(tokens);
});
//groupby
Map<String, Long> wordCount = wordList.stream().collect(Collectors.groupingBy(String::toString, Collectors.counting()));
定义一个词频Map，key是单词，value 是公式 1 计算的词频结果。现在词频已经计算出来。
2.2 下一步，计算逆文档频率
Map<String, Double> wordTF = new HashMap<>();
//公式（1）
wordCount.keySet().stream().forEach(key -> {
wordTF.put(key, wordCount.get(key) / Integer.valueOf(wordList.size()).doubleValue());
});
公式 2 中逆文档频率，首先计算分子，统计一共有多少个文档。然后计算分母，分别统计不同的词在多少个文档中出现。
首先定义一个 uniqueWords 变量，用来缓存所有文档中去除重复后的词。
再定义一个 wordDocCount 变量，用来缓存词在文档中出现的个数。
最后定义 wordIDF 变量用来缓存，不同词的IDF值。
Set<String> uniqueWords = new HashSet<>();
Map<String, Long> wordDocCount = new HashMap<>();
//统计唯一词
docs.stream().forEach(tokens -> {
uniqueWords.addAll(tokens);
});
//统计不同词在多少个文档中出现
docs.stream().forEach(doc -> {
uniqueWords.stream().forEach(token -> {
if (doc.contains(token)) {
if (!wordDocCount.containsKey(token)) {
wordDocCount.put(token, 0l);
}
wordDocCount.put(token, wordDocCount.get(token) + 1);
}
});
});
//计算IDF
Map<String, Double> wordIDF = new HashMap<>();
//公式（2）
wordDocCount.keySet().stream().forEach(key -> {
wordIDF.put(key, Math.log(Float.valueOf(Integer.valueOf(docs.size()).floatValue() / (wordDocCount.get(key)) + 1).doubleValue()));
});
计算出 TF与IDF之后，利用公式（3）计算，不同词的 TF-IDF
uniqueWords.stream().forEach(token -> {
wordTFIDF.put(token, wordTF.get(token) * wordIDF.get(token));
});
现在已经计算出不同词的TF-IDF值。
如果需要提取某个文档的关键词，只需要将这个文档，分词、去重，然后根据 TF-IDF排序，TF-IDF比较大的就是关键词，具体要返回几个关键词，这个需要自己根据需求考虑。
2.3 封装获取关键词代码
public List<String> keyword(Set<String> tokens, int topN) {
List<List<String>> tokensArr = tokens.stream().filter(token -> wordTFIDF.containsKey(token))
.map(token -> Arrays.asList(token, String.valueOf(wordTFIDF.get(token))))
.sorted(Comparator.comparing(t -> Double.valueOf(t.get(1)))).collect(Collectors.toList());
Collections.reverse(tokensArr);
if (tokensArr.size() < topN) {
topN = tokensArr.size();
}
List<String> keywords = tokensArr.subList(0, topN - 1).stream().map(tokenValue -> tokenValue.get(0)).collect(Collectors.toList());
return keywords;
}
3 实验测试
将代码定义一个 KeyWordAnalyzer 类
public static void main(String[] args) {
//训练模型
List<List<String>> docs=new ArrayList<>();
docs.add(Arrays.asList("中国 领导人 习近平 在 新冠疫情 爆发 以来 首次 视察 武汉 当天 外界 开始 分析 新冠肺炎 疫情 在 中国 是否 接近 尾声".split(" ")));
docs.add(Arrays.asList("同心 抗疫 武汉 加油 ".split(" ")));
docs.add(Arrays.asList("武汉 一直 是 习近平 总书记 心中 的 牵挂 疫情 发生 以来 习近平 为 统筹 疫情 防控 工作 周密 部署 为 谋划 经济 社会 发展 日夜 操劳 为 英雄 的 武汉 人民 真诚 赞叹".split(" ")));
docs.add(Arrays.asList("疫情 发生 以来 以 习近平 同志 为 核心 的 党中央 高度 重视 始终 把 人民 群众 生命 安全 和 身体 健康 放在 第一 位".split(" ")));
docs.add(Arrays.asList("武汉 不愧 为 英雄 的 城市 武汉 人民 不愧 为 英雄 的 人民 必将 通过 打赢 这次 抗击 新冠 肺炎 疫情 斗争 再次 被 载入 史册".split(" ")));
KeyWordAnalyzer keyWordAnalyzer=new KeyWordAnalyzer();
keyWordAnalyzer.calculateTFIDF(docs);
//测试模型
HashSet<String> tokens = new HashSet<>(Arrays.asList("习近平 总书记 一直 亲自 指挥 亲自 部署 武汉 的 疫情 防控 工作".split(" ")));
List<String> keyword = tfidfAnalyzer.keyword(tokens, 5);
System.out.println(keyword);
}
打印结果如下：
[的, 武汉, 疫情, 习近平]
关键词提取的还行，但是缺点是提取到了 “的”，很显然不是关键词，这就需要对文本进行预处理，去除 停用词后，再提取关键词效果会更好。