#!/usr/bin/env python # -*- coding: utf-8 -*- """This module implements functionality related to the `Okapi Best Matching `_ class of bag-of-words vector space models. Robertson and Zaragoza [1]_ describe the original algorithm and its modifications. .. [1] Robertson S., Zaragoza H. (2015). `The Probabilistic Relevance Framework: BM25 and Beyond, `_. """ from abc import ABCMeta, abstractmethod from collections import defaultdict import logging import math from gensim import interfaces, utils import numpy as np logger = logging.getLogger(__name__) class BM25ABC(interfaces.TransformationABC, metaclass=ABCMeta): """Objects of this abstract class realize the transformation between word-document co-occurrence matrix (int) into a BM25 matrix (positive floats). Concrete subclasses of this abstract class implement different BM25 scoring functions. """ def __init__(self, corpus=None, dictionary=None): r"""Pre-compute the average length of a document and inverse term document frequencies, which will be used to weight term frequencies for the documents. Parameters ---------- corpus : iterable of iterable of (int, int) or None, optional An input corpus, which will be used to compute the average length of a document and inverse term document frequencies. If None, then `dictionary` will be used to compute the statistics. If both `corpus` and `dictionary` are None, the statistics will be left unintialized. Default is None. dictionary : :class:`~gensim.corpora.Dictionary` An input dictionary, which will be used to compute the average length of a document and inverse term document frequencies. If None, then `corpus` will be used to compute the statistics. If both `corpus` and `dictionary` are None, the statistics will be left unintialized. Default is None. Attributes ---------- avgdl : float The average length of a document. idfs : dict of (int, float) A mapping from term ids to inverse term document frequencies. """ self.avgdl, self.idfs = None, None if dictionary: if corpus: logger.warning("constructor received both corpus and dictionary; ignoring the corpus") num_tokens = sum(dictionary.cfs.values()) self.avgdl = num_tokens / dictionary.num_docs self.idfs = self.precompute_idfs(dictionary.dfs, dictionary.num_docs) elif corpus: dfs = defaultdict(lambda: 0) num_tokens = 0 num_docs = 0 for bow in corpus: num_tokens += len(bow) for term_id in set(term_id for term_id, _ in bow): dfs[term_id] += 1 num_docs += 1 self.avgdl = num_tokens / num_docs self.idfs = self.precompute_idfs(dfs, num_docs) else: pass @abstractmethod def precompute_idfs(self, dfs, num_docs): """Precompute inverse term document frequencies, which will be used to weight term frequencies for the documents. Parameters ---------- dfs : dict of (int, int) A mapping from term ids to term document frequencies. num_docs : int The total number of documents in the training corpus. Returns ------- idfs : dict of (int, float) A mapping from term ids to inverse term document frequencies. """ pass @abstractmethod def get_term_weights(self, num_tokens, term_frequencies, idfs): """Compute vector space weights for a set of terms in a document. Parameters ---------- num_tokens : int The number of tokens in the document. term_frequencies : ndarray 1D array of term frequencies. idfs : ndarray 1D array of inverse term document frequencies. Returns ------- term_weights : ndarray 1D array of vector space weights. """ pass def __getitem__(self, bow): is_corpus, bow = utils.is_corpus(bow) if is_corpus: return self._apply(bow) num_tokens = sum(freq for term_id, freq in bow) term_ids, term_frequencies, idfs = [], [], [] for term_id, term_frequency in bow: term_ids.append(term_id) term_frequencies.append(term_frequency) idfs.append(self.idfs.get(term_id) or 0.0) term_frequencies, idfs = np.array(term_frequencies), np.array(idfs) term_weights = self.get_term_weights(num_tokens, term_frequencies, idfs) vector = [ (term_id, float(weight)) for term_id, weight in zip(term_ids, term_weights) ] return vector class OkapiBM25Model(BM25ABC): """The original Okapi BM25 scoring function of Robertson et al. [2]_. Examples -------- .. sourcecode:: pycon >>> from gensim.corpora import Dictionary >>> from gensim.models import OkapiBM25Model >>> from gensim.test.utils import common_texts >>> >>> dictionary = Dictionary(common_texts) # fit dictionary >>> model = OkapiBM25Model(dictionary=dictionary) # fit model >>> >>> corpus = [dictionary.doc2bow(line) for line in common_texts] # convert corpus to BoW format >>> vector = model[corpus[0]] # apply model to the first corpus document References ---------- .. [2] Robertson S. E., Walker S., Jones S., Hancock-Beaulieu M. M., Gatford M. (1995). `Okapi at TREC-3 `_. *NIST Special Publication 500-226*. """ def __init__(self, corpus=None, dictionary=None, k1=1.5, b=0.75, epsilon=0.25): r"""Pre-compute the average length of a document and inverse term document frequencies, which will be used to weight term frequencies for the documents. Parameters ---------- corpus : iterable of iterable of (int, int) or None, optional An input corpus, which will be used to compute the average length of a document and inverse term document frequencies. If None, then `dictionary` will be used to compute the statistics. If both `corpus` and `dictionary` are None, the statistics will be left unintialized. Default is None. dictionary : :class:`~gensim.corpora.Dictionary` An input dictionary, which will be used to compute the average length of a document and inverse term document frequencies. If None, then `corpus` will be used to compute the statistics. If both `corpus` and `dictionary` are None, the statistics will be left unintialized. Default is None. k1 : float A positive tuning parameter that determines the impact of the term frequency on its BM25 weight. Singhal [5]_ suggests to set `k1` between 1.0 and 2.0. Default is 1.5. b : float A tuning parameter between 0.0 and 1.0 that determines the document length normalization: 1.0 corresponds to full document normalization, while 0.0 corresponds to no length normalization. Singhal [5]_ suggests to set `b` to 0.75, which is the default. epsilon : float A positive tuning parameter that lower-bounds an inverse document frequency. Defaults to 0.25. Attributes ---------- k1 : float A positive tuning parameter that determines the impact of the term frequency on its BM25 weight. Singhal [3]_ suggests to set `k1` between 1.0 and 2.0. Default is 1.5. b : float A tuning parameter between 0.0 and 1.0 that determines the document length normalization: 1.0 corresponds to full document normalization, while 0.0 corresponds to no length normalization. Singhal [3]_ suggests to set `b` to 0.75, which is the default. epsilon : float A positive tuning parameter that lower-bounds an inverse document frequency. Defaults to 0.25. References ---------- .. [3] Singhal, A. (2001). `Modern information retrieval: A brief overview `_. *IEEE Data Eng. Bull.*, 24(4), 35–43. """ self.k1, self.b, self.epsilon = k1, b, epsilon super().__init__(corpus, dictionary) def precompute_idfs(self, dfs, num_docs): idf_sum = 0 idfs = dict() negative_idfs = [] for term_id, freq in dfs.items(): idf = math.log(num_docs - freq + 0.5) - math.log(freq + 0.5) idfs[term_id] = idf idf_sum += idf if idf < 0: negative_idfs.append(term_id) average_idf = idf_sum / len(idfs) eps = self.epsilon * average_idf for term_id in negative_idfs: idfs[term_id] = eps return idfs def get_term_weights(self, num_tokens, term_frequencies, idfs): term_weights = idfs * (term_frequencies * (self.k1 + 1) / (term_frequencies + self.k1 * (1 - self.b + self.b * num_tokens / self.avgdl))) return term_weights class LuceneBM25Model(BM25ABC): """The scoring function of Apache Lucene 8 [4]_. Examples -------- .. sourcecode:: pycon >>> from gensim.corpora import Dictionary >>> from gensim.models import LuceneBM25Model >>> from gensim.test.utils import common_texts >>> >>> dictionary = Dictionary(common_texts) # fit dictionary >>> corpus = [dictionary.doc2bow(line) for line in common_texts] # convert corpus to BoW format >>> >>> model = LuceneBM25Model(dictionary=dictionary) # fit model >>> vector = model[corpus[0]] # apply model to the first corpus document References ---------- .. [4] Kamphuis, C., de Vries, A. P., Boytsov, L., Lin, J. (2020). Which BM25 Do You Mean? `A Large-Scale Reproducibility Study of Scoring Variants `_. In: Advances in Information Retrieval. 28–34. """ def __init__(self, corpus=None, dictionary=None, k1=1.5, b=0.75): r"""Pre-compute the average length of a document and inverse term document frequencies, which will be used to weight term frequencies for the documents. Parameters ---------- corpus : iterable of iterable of (int, int) or None, optional An input corpus, which will be used to compute the average length of a document and inverse term document frequencies. If None, then `dictionary` will be used to compute the statistics. If both `corpus` and `dictionary` are None, the statistics will be left unintialized. Default is None. dictionary : :class:`~gensim.corpora.Dictionary` An input dictionary, which will be used to compute the average length of a document and inverse term document frequencies. If None, then `corpus` will be used to compute the statistics. If both `corpus` and `dictionary` are None, the statistics will be left unintialized. Default is None. k1 : float A positive tuning parameter that determines the impact of the term frequency on its BM25 weight. Singhal [5]_ suggests to set `k1` between 1.0 and 2.0. Default is 1.5. b : float A tuning parameter between 0.0 and 1.0 that determines the document length normalization: 1.0 corresponds to full document normalization, while 0.0 corresponds to no length normalization. Singhal [5]_ suggests to set `b` to 0.75, which is the default. Attributes ---------- k1 : float A positive tuning parameter that determines the impact of the term frequency on its BM25 weight. Singhal [3]_ suggests to set `k1` between 1.0 and 2.0. Default is 1.5. b : float A tuning parameter between 0.0 and 1.0 that determines the document length normalization: 1.0 corresponds to full document normalization, while 0.0 corresponds to no length normalization. Singhal [3]_ suggests to set `b` to 0.75, which is the default. """ self.k1, self.b = k1, b super().__init__(corpus, dictionary) def precompute_idfs(self, dfs, num_docs): idfs = dict() for term_id, freq in dfs.items(): idf = math.log(num_docs + 1.0) - math.log(freq + 0.5) idfs[term_id] = idf return idfs def get_term_weights(self, num_tokens, term_frequencies, idfs): term_weights = idfs * (term_frequencies / (term_frequencies + self.k1 * (1 - self.b + self.b * num_tokens / self.avgdl))) return term_weights class AtireBM25Model(BM25ABC): """The scoring function of Trotman et al. [5]_. Examples -------- .. sourcecode:: pycon >>> from gensim.corpora import Dictionary >>> from gensim.models import AtireBM25Model >>> from gensim.test.utils import common_texts >>> >>> dictionary = Dictionary(common_texts) # fit dictionary >>> corpus = [dictionary.doc2bow(line) for line in common_texts] # convert corpus to BoW format >>> >>> model = AtireBM25Model(dictionary=dictionary) # fit model >>> vector = model[corpus[0]] # apply model to the first corpus document References ---------- .. [5] Trotman, A., Jia X., Crane M., `Towards an Efficient and Effective Search Engine `_, In: SIGIR 2012 Workshop on Open Source Information Retrieval. 40–47. """ def __init__(self, corpus=None, dictionary=None, k1=1.5, b=0.75): r"""Pre-compute the average length of a document and inverse term document frequencies, which will be used to weight term frequencies for the documents. Parameters ---------- corpus : iterable of iterable of (int, int) or None, optional An input corpus, which will be used to compute the average length of a document and inverse term document frequencies. If None, then `dictionary` will be used to compute the statistics. If both `corpus` and `dictionary` are None, the statistics will be left unintialized. Default is None. dictionary : :class:`~gensim.corpora.Dictionary` An input dictionary, which will be used to compute the average length of a document and inverse term document frequencies. If None, then `corpus` will be used to compute the statistics. If both `corpus` and `dictionary` are None, the statistics will be left unintialized. Default is None. k1 : float A positive tuning parameter that determines the impact of the term frequency on its BM25 weight. Singhal [5]_ suggests to set `k1` between 1.0 and 2.0. Default is 1.5. b : float A tuning parameter between 0.0 and 1.0 that determines the document length normalization: 1.0 corresponds to full document normalization, while 0.0 corresponds to no length normalization. Singhal [5]_ suggests to set `b` to 0.75, which is the default. Attributes ---------- k1 : float A positive tuning parameter that determines the impact of the term frequency on its BM25 weight. Singhal [3]_ suggests to set `k1` between 1.0 and 2.0. Default is 1.5. b : float A tuning parameter between 0.0 and 1.0 that determines the document length normalization: 1.0 corresponds to full document normalization, while 0.0 corresponds to no length normalization. Singhal [3]_ suggests to set `b` to 0.75, which is the default. """ self.k1, self.b = k1, b super().__init__(corpus, dictionary) def precompute_idfs(self, dfs, num_docs): idfs = dict() for term_id, freq in dfs.items(): idf = math.log(num_docs) - math.log(freq) idfs[term_id] = idf return idfs def get_term_weights(self, num_tokens, term_frequencies, idfs): term_weights = idfs * (term_frequencies * (self.k1 + 1) / (term_frequencies + self.k1 * (1 - self.b + self.b * num_tokens / self.avgdl))) return term_weights