BM25.cpp

Go to the documentation of this file.

#include "BM25.h"
 
namespace krotos
{
 
BM25::BM25(const std::vector<String>& corpus) { initialise(corpus); }
 
void BM25::initialise(const std::vector<String>& corpus)
{
    for (std::size_t i = 0; i < corpus.size(); ++i)
    {
        const auto document = tokenize(corpus[i]);
        m_documentLength.push_back(document.size());
        for (const auto& word : document)
        {
            if (!m_termToDocument[word].contains(i))
            {
                m_termToDocument[word][i] = 0;
            }
            m_termToDocument[word][i] += 1;
        }
    }
    m_corpusSize = (int)m_documentLength.size();
    m_averageDocumentLength =
            (float)std::reduce(m_documentLength.begin(), m_documentLength.end()) / std::max(m_corpusSize, 1);
 
    std::vector<String> wordsWithNegativeIDF;
    for (const auto& [word, documents] : m_termToDocument)
    {
        const float idf = std::log(m_corpusSize - documents.size() + 0.5f) - std::log(documents.size() + 0.5f);
        m_inverseDocumentFrequency[word] = idf;
 
        if (idf < 0.0f)
        {
            wordsWithNegativeIDF.push_back(word);
        }
    }
 
    m_averageInverseDocumentFrequency = 0.0f;
    for (const auto& [key, value] : m_inverseDocumentFrequency)
    {
        m_averageInverseDocumentFrequency += value;
    }
    m_averageInverseDocumentFrequency /= std::max((int)m_inverseDocumentFrequency.size(), 1);
 
    const auto eps = m_epsilon * m_averageInverseDocumentFrequency;
    for (const auto& word : wordsWithNegativeIDF)
    {
        m_inverseDocumentFrequency[word] = eps;
    }
}
 
StringArray BM25::tokenize(const String& text)
{
    StringArray words;
    words.addTokens(text.trim().toLowerCase().removeCharacters(","), false);
    StringArray results;
    results.ensureStorageAllocated(words.size());
    for (auto word : words)
    {
        // check word isn't a stopword
        if (!m_stopwords.contains(word))
        {
            // apply stemmer
            word = stemmer(word);
            results.add(word);
        }
    }
    return results;
}
 
String BM25::stemmer(String word)
{
    // S-Stripper (a weak stemmer)
    // University of Otago at INEX 2010
    // X.-F. Jia, D. Alexander, V. Wood, A. Trotman (2010)
    // Proceedings of INEX 2010, pp. 250-268
 
    if (word.endsWith("ies"))
    {
        return word.replaceSection(word.length() - 3, 3, "y");
    }
    else if (word.endsWith("es"))
    {
        return word.replaceSection(word.length() - 2, 2, "");
    }
    else if (word.endsWith("s"))
    {
        return word.replaceSection(word.length() - 1, 1, "");
    }
    return word;
}
 
std::vector<float> BM25::getScores(String query)
{
    const auto queryTerms = tokenize(query);
 
    std::vector<float> scores(m_corpusSize, 0.0f);
    for (const auto& token : queryTerms)
    {
        if (m_termToDocument.contains(token))
        {
            for (const auto& [index, frequency] : m_termToDocument[token])
            {
                scores[index] +=
                        m_inverseDocumentFrequency[token] * frequency * (m_k1 + 1) /
                        (frequency + m_k1 * (1 - m_b + m_b * m_documentLength[index] / m_averageDocumentLength));
            }
        }
    }
    return scores;
}
 
std::vector<float> BM25::getBatchScores(String query, const std::vector<std::size_t>& ids)
{
    const auto queryTerms = tokenize(query);
 
    std::vector<float> scores(ids.size(), 0.0f);
    for (const auto& token : queryTerms)
    {
        if (m_termToDocument.contains(token))
        {
            for (std::size_t i = 0; i < ids.size(); ++i)
            {
                const auto index = ids[i];
                const auto frequency = m_termToDocument[token][index];
                scores[i] += m_inverseDocumentFrequency[token] * frequency * (m_k1 + 1) /
                             (frequency + m_k1 * (1 - m_b + m_b * m_documentLength[index] / m_averageDocumentLength));
            }
        }
    }
    return scores;
}
 
} // namespace krotos