#!/usr/bin/env python # -*- coding: utf-8 -*- from collections import defaultdict import math import unittest from gensim.models.bm25model import BM25ABC from gensim.models import OkapiBM25Model, LuceneBM25Model, AtireBM25Model from gensim.corpora import Dictionary class BM25Stub(BM25ABC): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def precompute_idfs(self, dfs, num_docs): return dict() def get_term_weights(self, num_tokens, term_frequencies, idfs): return term_frequencies class BM25ABCTest(unittest.TestCase): def setUp(self): self.documents = [['cat', 'dog', 'mouse'], ['cat', 'lion'], ['cat', 'lion']] self.dictionary = Dictionary(self.documents) self.expected_avgdl = sum(map(len, self.documents)) / len(self.documents) def test_avgdl_from_corpus(self): corpus = list(map(self.dictionary.doc2bow, self.documents)) model = BM25Stub(corpus=corpus) actual_avgdl = model.avgdl self.assertAlmostEqual(self.expected_avgdl, actual_avgdl) def test_avgdl_from_dictionary(self): model = BM25Stub(dictionary=self.dictionary) actual_avgdl = model.avgdl self.assertAlmostEqual(self.expected_avgdl, actual_avgdl) class OkapiBM25ModelTest(unittest.TestCase): def setUp(self): self.documents = [['cat', 'dog', 'mouse'], ['cat', 'lion'], ['cat', 'lion']] self.dictionary = Dictionary(self.documents) self.k1, self.b, self.epsilon = 1.5, 0.75, 0.25 def get_idf(word): frequency = sum(map(lambda document: word in document, self.documents)) return math.log((len(self.documents) - frequency + 0.5) / (frequency + 0.5)) dog_idf = get_idf('dog') cat_idf = get_idf('cat') mouse_idf = get_idf('mouse') lion_idf = get_idf('lion') average_idf = (dog_idf + cat_idf + mouse_idf + lion_idf) / len(self.dictionary) eps = self.epsilon * average_idf self.expected_dog_idf = dog_idf if dog_idf > 0 else eps self.expected_cat_idf = cat_idf if cat_idf > 0 else eps self.expected_mouse_idf = mouse_idf if mouse_idf > 0 else eps self.expected_lion_idf = lion_idf if lion_idf > 0 else eps def test_idfs_from_corpus(self): corpus = list(map(self.dictionary.doc2bow, self.documents)) model = OkapiBM25Model(corpus=corpus, k1=self.k1, b=self.b, epsilon=self.epsilon) actual_dog_idf = model.idfs[self.dictionary.token2id['dog']] actual_cat_idf = model.idfs[self.dictionary.token2id['cat']] actual_mouse_idf = model.idfs[self.dictionary.token2id['mouse']] actual_lion_idf = model.idfs[self.dictionary.token2id['lion']] self.assertAlmostEqual(self.expected_dog_idf, actual_dog_idf) self.assertAlmostEqual(self.expected_cat_idf, actual_cat_idf) self.assertAlmostEqual(self.expected_mouse_idf, actual_mouse_idf) self.assertAlmostEqual(self.expected_lion_idf, actual_lion_idf) def test_idfs_from_dictionary(self): model = OkapiBM25Model(dictionary=self.dictionary, k1=self.k1, b=self.b, epsilon=self.epsilon) actual_dog_idf = model.idfs[self.dictionary.token2id['dog']] actual_cat_idf = model.idfs[self.dictionary.token2id['cat']] actual_mouse_idf = model.idfs[self.dictionary.token2id['mouse']] actual_lion_idf = model.idfs[self.dictionary.token2id['lion']] self.assertAlmostEqual(self.expected_dog_idf, actual_dog_idf) self.assertAlmostEqual(self.expected_cat_idf, actual_cat_idf) self.assertAlmostEqual(self.expected_mouse_idf, actual_mouse_idf) self.assertAlmostEqual(self.expected_lion_idf, actual_lion_idf) def test_score(self): model = OkapiBM25Model(dictionary=self.dictionary, k1=self.k1, b=self.b, epsilon=self.epsilon) first_document = self.documents[0] first_bow = self.dictionary.doc2bow(first_document) weights = defaultdict(lambda: 0.0) weights.update(model[first_bow]) actual_dog_weight = weights[self.dictionary.token2id['dog']] actual_cat_weight = weights[self.dictionary.token2id['cat']] actual_mouse_weight = weights[self.dictionary.token2id['mouse']] actual_lion_weight = weights[self.dictionary.token2id['lion']] def get_expected_weight(word): idf = model.idfs[self.dictionary.token2id[word]] numerator = self.k1 + 1 denominator = 1 + self.k1 * (1 - self.b + self.b * len(first_document) / model.avgdl) return idf * numerator / denominator expected_dog_weight = get_expected_weight('dog') if 'dog' in first_document else 0.0 expected_cat_weight = get_expected_weight('cat') if 'cat' in first_document else 0.0 expected_mouse_weight = get_expected_weight('mouse') if 'mouse' in first_document else 0.0 expected_lion_weight = get_expected_weight('lion') if 'lion' in first_document else 0.0 self.assertAlmostEqual(expected_dog_weight, actual_dog_weight) self.assertAlmostEqual(expected_cat_weight, actual_cat_weight) self.assertAlmostEqual(expected_mouse_weight, actual_mouse_weight) self.assertAlmostEqual(expected_lion_weight, actual_lion_weight) class LuceneBM25ModelTest(unittest.TestCase): def setUp(self): self.documents = [['cat', 'dog', 'mouse'], ['cat', 'lion'], ['cat', 'lion']] self.dictionary = Dictionary(self.documents) self.k1, self.b = 1.5, 0.75 def get_idf(word): frequency = sum(map(lambda document: word in document, self.documents)) return math.log(1.0 + (len(self.documents) - frequency + 0.5) / (frequency + 0.5)) self.expected_dog_idf = get_idf('dog') self.expected_cat_idf = get_idf('cat') self.expected_mouse_idf = get_idf('mouse') self.expected_lion_idf = get_idf('lion') def test_idfs_from_corpus(self): corpus = list(map(self.dictionary.doc2bow, self.documents)) model = LuceneBM25Model(corpus=corpus, k1=self.k1, b=self.b) actual_dog_idf = model.idfs[self.dictionary.token2id['dog']] actual_cat_idf = model.idfs[self.dictionary.token2id['cat']] actual_mouse_idf = model.idfs[self.dictionary.token2id['mouse']] actual_lion_idf = model.idfs[self.dictionary.token2id['lion']] self.assertAlmostEqual(self.expected_dog_idf, actual_dog_idf) self.assertAlmostEqual(self.expected_cat_idf, actual_cat_idf) self.assertAlmostEqual(self.expected_mouse_idf, actual_mouse_idf) self.assertAlmostEqual(self.expected_lion_idf, actual_lion_idf) def test_idfs_from_dictionary(self): model = LuceneBM25Model(dictionary=self.dictionary, k1=self.k1, b=self.b) actual_dog_idf = model.idfs[self.dictionary.token2id['dog']] actual_cat_idf = model.idfs[self.dictionary.token2id['cat']] actual_mouse_idf = model.idfs[self.dictionary.token2id['mouse']] actual_lion_idf = model.idfs[self.dictionary.token2id['lion']] self.assertAlmostEqual(self.expected_dog_idf, actual_dog_idf) self.assertAlmostEqual(self.expected_cat_idf, actual_cat_idf) self.assertAlmostEqual(self.expected_mouse_idf, actual_mouse_idf) self.assertAlmostEqual(self.expected_lion_idf, actual_lion_idf) def test_score(self): model = LuceneBM25Model(dictionary=self.dictionary, k1=self.k1, b=self.b) first_document = self.documents[0] first_bow = self.dictionary.doc2bow(first_document) weights = defaultdict(lambda: 0.0) weights.update(model[first_bow]) actual_dog_weight = weights[self.dictionary.token2id['dog']] actual_cat_weight = weights[self.dictionary.token2id['cat']] actual_mouse_weight = weights[self.dictionary.token2id['mouse']] actual_lion_weight = weights[self.dictionary.token2id['lion']] def get_expected_weight(word): idf = model.idfs[self.dictionary.token2id[word]] denominator = 1 + self.k1 * (1 - self.b + self.b * len(first_document) / model.avgdl) return idf / denominator expected_dog_weight = get_expected_weight('dog') if 'dog' in first_document else 0.0 expected_cat_weight = get_expected_weight('cat') if 'cat' in first_document else 0.0 expected_mouse_weight = get_expected_weight('mouse') if 'mouse' in first_document else 0.0 expected_lion_weight = get_expected_weight('lion') if 'lion' in first_document else 0.0 self.assertAlmostEqual(expected_dog_weight, actual_dog_weight) self.assertAlmostEqual(expected_cat_weight, actual_cat_weight) self.assertAlmostEqual(expected_mouse_weight, actual_mouse_weight) self.assertAlmostEqual(expected_lion_weight, actual_lion_weight) class AtireBM25ModelTest(unittest.TestCase): def setUp(self): self.documents = [['cat', 'dog', 'mouse'], ['cat', 'lion'], ['cat', 'lion']] self.dictionary = Dictionary(self.documents) self.k1, self.b, self.epsilon = 1.5, 0.75, 0.25 def get_idf(word): frequency = sum(map(lambda document: word in document, self.documents)) return math.log(len(self.documents) / frequency) self.expected_dog_idf = get_idf('dog') self.expected_cat_idf = get_idf('cat') self.expected_mouse_idf = get_idf('mouse') self.expected_lion_idf = get_idf('lion') def test_idfs_from_corpus(self): corpus = list(map(self.dictionary.doc2bow, self.documents)) model = AtireBM25Model(corpus=corpus, k1=self.k1, b=self.b) actual_dog_idf = model.idfs[self.dictionary.token2id['dog']] actual_cat_idf = model.idfs[self.dictionary.token2id['cat']] actual_mouse_idf = model.idfs[self.dictionary.token2id['mouse']] actual_lion_idf = model.idfs[self.dictionary.token2id['lion']] self.assertAlmostEqual(self.expected_dog_idf, actual_dog_idf) self.assertAlmostEqual(self.expected_cat_idf, actual_cat_idf) self.assertAlmostEqual(self.expected_mouse_idf, actual_mouse_idf) self.assertAlmostEqual(self.expected_lion_idf, actual_lion_idf) def test_idfs_from_dictionary(self): model = AtireBM25Model(dictionary=self.dictionary, k1=self.k1, b=self.b) actual_dog_idf = model.idfs[self.dictionary.token2id['dog']] actual_cat_idf = model.idfs[self.dictionary.token2id['cat']] actual_mouse_idf = model.idfs[self.dictionary.token2id['mouse']] actual_lion_idf = model.idfs[self.dictionary.token2id['lion']] self.assertAlmostEqual(self.expected_dog_idf, actual_dog_idf) self.assertAlmostEqual(self.expected_cat_idf, actual_cat_idf) self.assertAlmostEqual(self.expected_mouse_idf, actual_mouse_idf) self.assertAlmostEqual(self.expected_lion_idf, actual_lion_idf) def test_score(self): model = AtireBM25Model(dictionary=self.dictionary, k1=self.k1, b=self.b) first_document = self.documents[0] first_bow = self.dictionary.doc2bow(first_document) weights = defaultdict(lambda: 0.0) weights.update(model[first_bow]) actual_dog_weight = weights[self.dictionary.token2id['dog']] actual_cat_weight = weights[self.dictionary.token2id['cat']] actual_mouse_weight = weights[self.dictionary.token2id['mouse']] actual_lion_weight = weights[self.dictionary.token2id['lion']] def get_expected_weight(word): idf = model.idfs[self.dictionary.token2id[word]] numerator = self.k1 + 1 denominator = 1 + self.k1 * (1 - self.b + self.b * len(first_document) / model.avgdl) return idf * numerator / denominator expected_dog_weight = get_expected_weight('dog') if 'dog' in first_document else 0.0 expected_cat_weight = get_expected_weight('cat') if 'cat' in first_document else 0.0 expected_mouse_weight = get_expected_weight('mouse') if 'mouse' in first_document else 0.0 expected_lion_weight = get_expected_weight('lion') if 'lion' in first_document else 0.0 self.assertAlmostEqual(expected_dog_weight, actual_dog_weight) self.assertAlmostEqual(expected_cat_weight, actual_cat_weight) self.assertAlmostEqual(expected_mouse_weight, actual_mouse_weight) self.assertAlmostEqual(expected_lion_weight, actual_lion_weight)