计算两篇文档的余弦相似度（tfidf）

# -*- coding:utf-8 -*-

"""
@author: Linlifang
"""

import os
import jieba
import sys
import re
import string
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.feature_extraction.text import CountVectorizer
reload(sys)
sys.setdefaultencoding('utf-8')

'''
首先读取文件夹里的文档，然后通过结巴分词，将分词的结果存入文件，接着使用sklearn包计算每一篇文档的tfidf值
并保存在一个文件里，最后从这些文件中任选两个txt文件来计算他们的余弦相似度。
'''
def getFileList(path):
    filelist = []
    files = os.listdir(path)
    for f in files:
        if f[0] == '.':
            pass
        else:
            filelist.append(f)
    return filelist, path


def segment(filename, path, segPath):
    f = open(path + "/" + filename, 'r+')
    file_list = f.read()
    f.close()
#对文档进行分词处理
    if not os.path.exists(segPath):
        os.mkdir(segPath)
#对空格，换行符进行处理
    # Segmenting the document
    seg_list = jieba.cut(file_list, cut_all=False)
    # stopword = open('stopworda.txt').readlines()
    result = []
    for seg in seg_list:
        seg = ''.join(seg.split())
        reg = 'w+'
        r = re.search(reg, seg)
        if seg != '' and seg != ' = ' and seg != '[' and seg != ']' and seg != '(' and seg != ')' and not r:
            result.append(seg)
    finalresult = []
    stopword = open('stopworda.txt').read()
    for word in result: #去除停用词
        if word in stopword:
            continue
        else:
            if word >= u'\u4e00' and word <= u'\u9fa5': #判断是否是汉字
                finalresult.append(word)

            # 将分词后的结果用空格隔开，保存在本地
    f = open(segPath + "/" + filename + "-seg.txt", "w+")
    f.write(' '.join(finalresult))
    f.close()


# 读取已经分词好的文档，进行TFIDF计算
def Tfidf(filelist, sFilePath, path):
    corpus = []
    for ff in filelist:
        fname = path + "/" + ff
        f = open(fname + "-seg.txt", 'r+')
        content = f.read()
        f.close()
        corpus.append(content)

    vectorizer = CountVectorizer()
    transformer = TfidfTransformer()
    tfidf = transformer.fit_transform(vectorizer.fit_transform(corpus))
    word = vectorizer.get_feature_names()  # 全部文本关键字
    weight = tfidf.toarray()

    if not os.path.exists(sFilePath):
        os.mkdir(sFilePath)

    for i in range(len(weight)):
        print u'-writing all the tf-idf in the ', i, u'file into ', sFilePath + '/' + string.zfill(i, 2) + ".txt"
        f = open(sFilePath + "/" + string.zfill(i, 2) + ".txt", 'w+')
        for j in range(len(word)):
            f.write(word[j] + "  " + str(weight[i][j]) + " " + "\n")
        
def coutcos(file1,file2):
    cipin1 = open(file1).readlines()
    cipin2 = open(file2).readlines()
    list1 = []
    list2 = []
    for x in cipin1:
        y = x.split(' ')
        list1.append(y[2])
    for x in cipin2:
        y = x.split(' ')
        list2.append(y[2])
    dot_product = 0.0
    normA = 0.0
    normB = 0.0
    for a, b in zip(list1, list2):
        a = float(a)
        b = float(b)
        dot_product += a * b
        normA += a ** 2
        normB += b ** 2
    if normA == 0.0 or normB == 0.0:
        return None
    else:
        return dot_product / ((normA * normB) ** 0.5)

if __name__ == "__main__":
    # 保存TFIDF的计算结果到文件夹
    sFilePath = "C:/Users/llfang1/PycharmProjects/untitled2/corpus/tfidffile"
    # 保存分词的文件夹
    segPath = 'C:/Users/llfang1/PycharmProjects/untitled2/corpus/segfile'
    (allfile, path) = getFileList('C:/Users/llfang1/PycharmProjects/untitled2/corpus/allkeyword')
    for ff in allfile:
        print "Using jieba on " + ff
        segment(ff, path, segPath)
    Tfidf(allfile, sFilePath, segPath)
    file1 = sFilePath + "/" + "04.txt"
    file2 = sFilePath + "/" + "05.txt"
    similar = coutcos(file1,file2)
    print similar

注：此程序參考了一位同行的程序后进行了改动并添加一些内容