自己手写BM25算法

python 实现 BM25 算法

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2020/7/30 20:50
# @Author  : fuGuoWen
# @Site    : 
# @File    : test06.py
# @Software: PyCharm
import math

class BM25(object):

    def __init__(self, docs):
        """

        :param docs: 分好词的list
        """
        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 = {}
        self.k1 = 1.5
        self.b = 0.75
        self.init()

    def init(self):
        for doc in self.docs:
            # 统计每一行单词出现的个数
            tmp = {}
            for word in doc:
                if not word in tmp:
                    tmp[word] = 0
                tmp[word] += 1
            self.f.append(tmp)
            for k, v in tmp.items():
                if k not in self.df:
                    self.df[k] = 0
                self.df[k] += 1
        for k, v in self.df.items():
            # self.idf[k] = math.log(self.D - v + 0.5) - math.log(v + 0.5)
            self.idf[k] = math.log(self.D +1) - math.log(v + 0.5)

    def sim(self, doc, index):
        """

        :param doc: 问题
        :param index: 训练数据的下标
        :return:
        """
        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, doc):
        """
        找出训练数据中所有相似的句子概率
        :param doc:  一句话的分词list
        :return:
        """
        scores = []
        for index in range(self.D):
            score = self.sim(doc, index)
            scores.append(score)
        return scores

if __name__ == '__main__':
    corpus = [
       ['hello','world','hello'],['hello','go']
    ]
    bm=BM25(corpus)
    # 统计每个小列表中出现的次数
    print(bm.f)
    # 每个单词在文档库中出现了在几个文档中
    print(bm.df)
    print(bm.idf)
    scores=bm.simall(['hello'])
    print(scores)