redhongx
52 lines
1.4 KiB
Python
52 lines
1.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
from __future__ import unicode_literals
|
|
|
|
import math
|
|
|
|
|
|
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 = {}
|
|
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)
|
|
|
|
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, doc):
|
|
scores = []
|
|
for index in range(self.D):
|
|
score = self.sim(doc, index)
|
|
scores.append(score)
|
|
return scores
|