naive_bayes.h

Go to the documentation of this file.

#ifndef META_NAIVE_BAYES_H_
#define META_NAIVE_BAYES_H_

#include <unordered_map>
#include "index/forward_index.h"
#include "classify/classifier/classifier.h"
#include "classify/classifier_factory.h"
#include "meta.h"
#include "stats/multinomial.h"
#include "util/sparse_vector.h"

namespace meta
{
namespace classify
{

class naive_bayes : public classifier
{
  public:
    const static constexpr double default_alpha = 0.1;
    const static constexpr double default_beta = 0.1;

    naive_bayes(std::shared_ptr<index::forward_index> idx,
                double alpha = default_alpha, double beta = default_beta);

    void train(const std::vector<doc_id>& docs) override;

    class_label classify(doc_id d_id) override;

    void reset() override;

    const static std::string id;

  private:
    util::sparse_vector<class_label, stats::multinomial<term_id>> term_probs_;

    stats::multinomial<class_label> class_probs_;
};

template <>
std::unique_ptr<classifier>
    make_classifier<naive_bayes>(const cpptoml::table& config,
                                 std::shared_ptr<index::forward_index> idx);

}
}
#endif