信息检索之BM25算法

BM25算法参考文档，搜索引擎中的相关应用。

import math
import jieba


class BM25(object):

    def __init__(self, docs):
        self.D = len(docs)  # 训练集的文本数量
        self.avgdl = sum([len(doc)+0.0 for doc in docs])/self.D # 平均每个文本的长度
        self.docs = docs
        self.f = []  # 一个文档中每个词的出现次数
        self.df = {} # 存储每个词及出现了该词的文档数量
        self.idf = {} # 存储每个词的idf值
        self.k1 = 1.5
        self.b = 0.75
        self.init()

    def init(self):
        for doc in self.docs:
            tmp = {}
            for word in doc:
                tmp[word] = tmp.get(word, 0) + 1  # 存储每个文档中每个词的出现次数
            self.f.append(tmp)
            for k in tmp.keys():
                self.df[k] = self.df.get(k, 0) + 1
        for k, v in self.df.items():
            self.idf[k] = math.log(self.D-v+0.5)-math.log(v+0.5)

        for key, value in self.__dict__.items():
            print(key, value)

    def sim(self, doc, index):
        score = 0
        for word in doc:
            if word not in self.f[index]:
                continue
            d = len(self.docs[index])
            score += (self.idf[word]*self.f[index][word]*(self.k1+1)/(self.f[index][word]+self.k1*(1-self.b+self.b*d/self.avgdl)))
        return score

    def simall(self, query_list):
        sim_dict = {}

        # 遍历所有的候选文章，计算每个文章与querey的score
        for index in range(self.D):
            # 查询的每个语素，计算每个语素与该句子的score
            score = 0
            for query in query_list:
                # 每个语素的idf
                if query in self.idf:
                    quer_idf = self.idf[query]
                else:
                    quer_idf = 0
                # 计算R
                if query in self.f[index]:
                    f_i = self.f[index][query]
                else:
                    f_i = 0
                R = (f_i*(self.k1+1))/(f_i+self.k1*(1-self.b+self.b*len(self.docs[index])/self.avgdl))
                # 计算idf*R
                sim = quer_idf * R
                score += sim
            if score != 0:
                sim_dict[score] = index

        sim_dict = sorted(sim_dict.items(),key=lambda item:item[0], reverse=True)
        return sim_dict

if __name__ == '__main__':

    sents = []
    sents.append("自然语言自然语言处理是计算机科学领域与人工智能领域中的一个重要方向。")
    sents.append("它研究能实现人与计算机之间用自然语言进行有效通信的各种理论和方法。")
    sents.append("自然语言处理是一门融语言学、计算机科学、数学于一体的科学。")
    sents.append("因此，这一领域的研究将涉及自然语言，即人们日常使用的语言，")
    sents.append("所以它与语言学的研究有着密切的联系，但又有重要的区别。")
    sents.append("自然语言处理并不是一般地研究自然语言，")
    sents.append("而在于研制能有效地实现自然语言通信的计算机系统，")
    sents.append("特别是其中的软件系统。因而它是计算机科学的一部分。")

    # 分词并剔除停用词
    docs = []
    for sent in sents:
        words = list(jieba.cut(sent))
        new_words = []
        for word in words:
            if word not in ["，", "。", "的", "而", "是", "与", "中", "它", "能"]:
                new_words.append(word)
        docs.append(new_words)
    s = BM25(docs)

    print()
    list1 = ['自然语言', '计算机科学', '人工智能', '领域']   # 分词之后的list
    sim_dict = s.simall(list1)
    for score, index in sim_dict:
        print(score, sents[index])