信息檢索之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])