BM25.h

Go to the documentation of this file.

#pragma once
 
namespace krotos
{
class BM25
{
  public:
    BM25(const std::vector<String>& corpus);
    std::vector<float> getScores(String query);
    std::vector<float> getBatchScores(String query, const std::vector<std::size_t>& ids);
 
  private:
    void initialise(const std::vector<String>& corpus);
    StringArray tokenize(const String& text);
    String stemmer(String word);
 
    const float m_k1{1.5f};
    const float m_b{0.75f};
    const float m_epsilon{0.25f};
    float m_averageDocumentLength{1.0f};
    int m_corpusSize{0};
    float m_averageInverseDocumentFrequency{0.0f};
 
    std::unordered_map<String, std::unordered_map<int, int>> m_termToDocument;
    std::unordered_map<String, float> m_inverseDocumentFrequency;
    std::vector<int> m_documentLength;
 
    // TODO: validate that removing stopwords is useful, and try a smaller set of stopwords
    StringArray m_stopwords = {
            "i",          "me",       "my",         "myself",    "we",       "our",      "ours",       "ourselves",
            "you",        "you\'re",  "you\'ve",    "you\'ll",   "you\'d",   "your",     "yours",      "yourself",
            "yourselves", "he",       "him",        "his",       "himself",  "she",      "she\'s",     "her",
            "hers",       "herself",  "it",         "it\'s",     "its",      "itself",   "they",       "them",
            "their",      "theirs",   "themselves", "what",      "which",    "who",      "whom",       "this",
            "that",       "that\'ll", "these",      "those",     "am",       "is",       "are",        "was",
            "were",       "be",       "been",       "being",     "have",     "has",      "had",        "having",
            "do",         "does",     "did",        "doing",     "a",        "an",       "the",        "and",
            "but",        "if",       "or",         "because",   "as",       "until",    "while",      "of",
            "at",         "by",       "for",        "with",      "about",    "against",  "between",    "into",
            "through",    "during",   "before",     "after",     "above",    "below",    "to",         "from",
            "up",         "down",     "in",         "out",       "on",       "off",      "over",       "under",
            "again",      "further",  "then",       "once",      "here",     "there",    "when",       "where",
            "why",        "how",      "all",        "any",       "both",     "each",     "few",        "more",
            "most",       "other",    "some",       "such",      "no",       "nor",      "not",        "only",
            "own",        "same",     "so",         "than",      "too",      "very",     "s",          "t",
            "can",        "will",     "just",       "don",       "don\'t",   "should",   "should\'ve", "now",
            "d",          "ll",       "m",          "o",         "re",       "ve",       "y",          "ain",
            "aren",       "aren\'t",  "couldn",     "couldn\'t", "didn",     "didn\'t",  "doesn",      "doesn\'t",
            "hadn",       "hadn\'t",  "hasn",       "hasn\'t",   "haven",    "haven\'t", "isn",        "isn\'t",
            "ma",         "mightn",   "mightn\'t",  "mustn",     "mustn\'t", "needn",    "needn\'t",   "shan",
            "shan\'t",    "shouldn",  "shouldn\'t", "wasn",      "wasn\'t",  "weren",    "weren\'t",   "won",
            "won\'t",     "wouldn",   "wouldn\'t"};
};
 
} // namespace krotos