NLTK实现TF-IDF，并结合余弦相似度进行文本相似度计算（附完整代码实现）
NLTK实现TF-IDF，并结合余弦相似度进行文本相似度计算
TF-IDF（词频-逆文件频率）
TF-IDF（term frequency–inverse document frequency，词频-逆文件频率）是一种常用的加权技术，一般被用来寻找文本中的关键词。
TF-IDF实际上可分为TF（term frequency，词频）和IDF（inverse document frequency，逆文件频率）：
T F i , j = n i , j ∑ k n k , j TF_{i,j}=\frac{n_{i,j}}{\sum _k n_{k,j}} TFi,j=∑knk,jni,j
其中 ni,j是该词在文件 dj 中出现的次数，分母则是文件 dj 中所有词汇出现的次数总和。
I D F i = log ∣ D ∣ ∣ { j : t i ∈ d j } ∣ IDF_i =\log \frac{|D|}{|\{j:t_i\in d_j\}|} IDFi=log∣{j:ti∈dj}∣∣D∣
其中，|D|是语料库中的文件总数。 |{j:ti∈dj}| 表示包含词语 ti 的文件数目（即 ni,j≠0 的文件数目）。
文档中某个词的TF-IDF值计算方法即为：
T F − I D F = T F × I D F TF-IDF=TF \times IDF TF−IDF=TF×IDF
Cosine（余弦相似度）
Cosine（余弦相似度）是通过测量两个向量的夹角的余弦值（也即向量内积）来度量它们之间的相似性，常用于文本相似度计算中，例如，向量a和b的余弦相似度计算公式即为：
cos ( θ ) = a ⋅ b ∣ ∣ a ∣ ∣ × ∣ ∣ b ∣ ∣ \cos(\theta)=\frac{a \cdot b}{||a|| \times ||b||} cos(θ)=∣∣a∣∣×∣∣b∣∣a⋅b
计算文本相似度时，文本中每个词项会被赋予不同的权重，并合并为向量。即，该向量各个维度上的值对应于文档中每个词在该文档中的权重。
文本相似度计算
文本相似度计算的大体思路为：
1.使用现有文本库，构建语料库
2.针对需要比较相似度的两条文本，分别计算其每一个单词的权重（TF-IDF）
3.根据单词权重，生成文本的特征向量
4.对两条文本的特征向量计算其余弦相似度，即为两条文本的文本相似度
文本相似度代码实现
代码使用到的第三方库包含：
import operator
import numpy as np
from nltk import word_tokenize
from nltk.corpus import stopwords
from nltk.text import TextCollection
from nltk.stem import WordNetLemmatizer
from numpy import dot
from numpy.linalg import norm
1、构建语料库
在使用TF-IDF生成首先需要使用NLTK对文本库中所有的文本进行分词，以构建语料库：
"""
yzuy
使用文本库中所有的文本构建语料库
:return: Corpus: 语料库
"""
WordList = []
for File in FileList: # FileList为文本文件列表
with open(File) as f:
Text = f.read()
WordList.append(Split_Word(Text)) # 对每一个文件进行分词等，得到词表
Corpus = TextCollection(WordList) # 使用词表构建语料库
其中的Split_Word()函数主要用来对文本进行分词、去标点符号、去停用词、词形还原：
Punctuation = ['~', '`', '``', '!', '@', '#', '$', '%', '^', '&', '*', '(', ')', '_', '+', '-', '=', '{', '}', '|', '[',']', '\\', ':', '\"', ';', '\'', '<', '>', '?', ',', '.', '/'] # 标点符号表
def Split_Word(Text):
"""
yzuy
对文本进行分词、去标点符号、去停用词、词形还原
:param Text: 文本库中的一段文本
:return: TextAfterLemmatize: 文本经过处理后的得到的词表
"""
TokenText = word_tokenize(Text.lower()) # 将文本转为小写（也可以不转），并使用nltk对其进行分词
TextWithoutPunc = [word for word in TokenText if word not in Punctuation] # 去除标点符号项
stops = set(stopwords.words("english")) # 加载NLTK英文停用词表
TextWithoutDeac = [word for word in TextWithoutPunc if word not in stops] # 去停用词
TextAfterLemmatize = []
lemmatizer = WordNetLemmatizer()
for word in TextWithoutDeac:
TextAfterLemmatize.append(lemmatizer.lemmatize(lemmatizer.lemmatize(word, pos='v'), pos='n')) # 词形还原
return TextAfterLemmatize
2、计算TF-IDF值
此步骤可以直接使用NLTK自带的TF-IDF函数实现：
def TF_IDF(Corpus, WordList):
"""
yzuy
计算文本中每个单词的TF-IDF值
:param Corpus: 语料库
:param WordList: 目标文本使用上一步的Split_Word(Text)函数处理得到的词表
:return: TF_IDF_Dict: 文本中每一个单词及其TF-IDF值构成的字典
"""
TF_IDF_Dict = {}
for Word in WordList:
TF_IDF_Dict[Word] = Corpus.tf_idf(Word, WordList) # 使用nltk自带的TF-IDF函数对词表中每个词计算其TF-IDF值
return TF_IDF_Dict
3、生成文本特征向量
本步骤统计两条目标文本列表的并集，并针对每一条文本生成其特征向量：
def Construct_Union(WordListA, WordListB):
"""
yzuy
针对两条目标文本列表获取其并集
:param WordListA: 目标文本列表A
:param WordListB: 目标文本列表B
:return: Union_List: 文本列表的并集
"""
Union_List = list(set(WordListA).union(set(WordListB)))
return Union_List
def Vectorization(Union_List, SortedList):
"""
yzuy
生成目标文本的特征向量
:param Union_List: 两条目标文本列表的并集
:param SortedList: 本条文本中单词对应的权重值列表
:return: Vector: 本条文本的特征向量
"""
Vector = []
for Word in Union_List:
Num = 0.0
for Tuple in SortedList: # SortedList格式为:[('mini', 0.44), ('medium', 0.29)]
if Tuple[0] == Word:
Num = Tuple[1] # 取该单词的权重
break
Vector.append(Num)
return Vector
4、计算文本相似度
针对两条文本的特征向量计算其文本相似度：
def Cosine(VectorA, VectorB):
"""
yzuy
计算两向量的余弦相似度
:param VectorA: 向量A
:param VectorB: 向量B
:return: Similarity: 向量余弦相似度
"""
Similarity = dot(VectorA, VectorB) / (norm(VectorA) * norm(VectorB))
if np.isnan(Similarity): # 出现分母为0情况，直接返回其相似度为0
return 0.0
return Similarity
def CountSimilarity(Corpus, TextA, TextB):
"""
yzuy
计算两条文本的相似度
:param Corpus: 语料库
:param TextA: 目标文本A
:param TextB: 目标文本B
:return: Vector: 本条文本的特征向量
"""
ListA = Split_Word(TextA)
ListB = Split_Word(TextB)
Union_List = Construct_Union(ListA, ListB) # 针对两条目标文本列表获取其并集
TF_IDF_A = TF_IDF(Corpus, ListA) # 计算文本列表A的TF-IDF值
TF_IDF_B = TF_IDF(Corpus, ListB)
Sorted_A = sorted(TF_IDF_A.items(), key=operator.itemgetter(1), reverse=True) # 根据单词特征值大小进行排序
Sorted_B = sorted(TF_IDF_B.items(), key=operator.itemgetter(1), reverse=True)
Vector_A = Vectorization(Union_List, Sorted_A) # 生成目标文本的特征向量
Vector_B = Vectorization(Union_List, Sorted_B)
return Cosine(Vector_A, Vector_B) # 计算两向量的余弦相似度并返回，作为两条目标文本的相似度
整体实现
import os
import operator
import numpy as np
from nltk import word_tokenize
from nltk.corpus import stopwords
from nltk.text import TextCollection
from nltk.stem import WordNetLemmatizer
from numpy import dot
from numpy.linalg import norm
Punctuation = ['~', '`', '``', '!', '@', '#', '$', '%', '^', '&', '*', '(', ')', '_', '+', '-', '=', '{', '}', '|', '[',
']', '\\', ':', '\"', ';', '\'', '<', '>', '?', ',', '.', '/']
Text1 = 'The headerVerifyInfo function in lib/header.c in RPM before 4.9.1.3 allows remote attackers to cause a denial of service (crash) and possibly execute arbitrary code via a negative value in a region offset of a package header, which is not properly handled in a numeric range comparison.'
Text2 = 'The php_register_variable_ex function in php_variables.c in PHP 5.3.9 allows remote attackers to execute arbitrary code via a request containing a large number of variables, related to improper handling of array variables. NOTE: this vulnerability exists because of an incorrect fix for CVE-2011-4885.'
def Split_Word(Text):
"""
yzuy
对文本进行分词、去标点符号、去停用词、词形还原
:param Text: 文本库中的一段文本
:return: TextAfterLemmatize: 文本经过处理后的得到的词表
"""
TokenText = word_tokenize(Text.lower()) # 将文本转为小写（也可以不转），并使用nltk对其进行分词
TextWithoutPunc = [word for word in TokenText if word not in Punctuation] # 去除标点符号项
stops = set(stopwords.words("english")) # 加载NLTK英文停用词表
TextWithoutDeac = [word for word in TextWithoutPunc if word not in stops] # 去停用词
TextAfterLemmatize = []
lemmatizer = WordNetLemmatizer()
for word in TextWithoutDeac:
TextAfterLemmatize.append(lemmatizer.lemmatize(lemmatizer.lemmatize(word, pos='v'), pos='n')) # 词形还原
return TextAfterLemmatize
def Traverse_Path(Path):
"""
yzuy
遍历路径，获取路径下的文件列表
:param Path: 目标路径
:return: FileList: 路径下的文件列表
"""
FileList = []
for roots, dirs, files in os.walk(Path):
for file in files:
FilePath = os.path.join(roots, file)
FileList.append(FilePath)
return FileList
def TF_IDF(Corpus, WordList):
"""
yzuy
计算文本中每个单词的TF-IDF值
:param Corpus: 语料库
:param WordList: 目标文本使用上一步的Split_Word(Text)函数处理得到的词表
:return: TF_IDF_Dict: 文本中每一个单词及其TF-IDF值构成的字典
"""
TF_IDF_Dict = {}
for Word in WordList:
TF_IDF_Dict[Word] = Corpus.tf_idf(Word, WordList) # 使用nltk自带的TF-IDF函数对词表中每个词计算其TF-IDF值
return TF_IDF_Dict
def Construct_Union(WordListA, WordListB):
"""
yzuy
针对两条目标文本列表获取其并集
:param WordListA: 目标文本列表A
:param WordListB: 目标文本列表B
:return: Union_List: 文本列表的并集
"""
Union_List = list(set(WordListA).union(set(WordListB)))
return Union_List
def Vectorization(Union_List, SortedList):
"""
yzuy
生成目标文本的特征向量
:param Union_List: 两条目标文本列表的并集
:param SortedList: 本条文本中单词对应的权重值列表
:return: Vector: 本条文本的特征向量
"""
Vector = []
for Word in Union_List:
Num = 0.0
for Tuple in SortedList: # SortedList格式为:[('mini', 0.44), ('medium', 0.29)]
if Tuple[0] == Word:
Num = Tuple[1] # 取该单词的权重
break
Vector.append(Num)
return Vector
def Cosine(VectorA, VectorB):
"""
yzuy
计算两向量的余弦相似度
:param VectorA: 向量A
:param VectorB: 向量B
:return: Similarity: 向量余弦相似度
"""
Similarity = dot(VectorA, VectorB) / (norm(VectorA) * norm(VectorB))
if np.isnan(Similarity): # 出现分母为0情况，直接返回其相似度为0
return 0.0
return Similarity
def CountSimilarity(Corpus, TextA, TextB):
"""
yzuy
计算两条文本的相似度
:param Corpus: 语料库
:param TextA: 目标文本A
:param TextB: 目标文本B
:return: Vector: 本条文本的特征向量
"""
ListA = Split_Word(TextA)
ListB = Split_Word(TextB)
Union_List = Construct_Union(ListA, ListB) # 针对两条目标文本列表获取其并集
TF_IDF_A = TF_IDF(Corpus, ListA) # 计算文本列表A的TF-IDF值
TF_IDF_B = TF_IDF(Corpus, ListB)
Sorted_A = sorted(TF_IDF_A.items(), key=operator.itemgetter(1), reverse=True) # 根据单词特征值大小进行排序
Sorted_B = sorted(TF_IDF_B.items(), key=operator.itemgetter(1), reverse=True)
Vector_A = Vectorization(Union_List, Sorted_A) # 生成目标文本的特征向量
Vector_B = Vectorization(Union_List, Sorted_B)
return Cosine(Vector_A, Vector_B) # 计算两向量的余弦相似度并返回，作为两条目标文本的相似度
if __name__ == '__main__':
WordList = []
OriginPath = './dataset' # 待测文本存储位置
FileList = Traverse_Path(OriginPath)
for File in FileList: # FileList为文本文件列表
with open(File) as f:
Text = f.read()
WordList.append(Split_Word(Text)) # 对每一个文件进行分词等，得到词表
Corpus = TextCollection(WordList) # 使用词表构建语料库
Similarity = CountSimilarity(Corpus, Text1, Text2)
整体代码及示例文本库已整理发布在Github：
直达链接