<commit_before>#ifndef DIMMWITTED_GIBBS_SAMPLER_H_
#define DIMMWITTED_GIBBS_SAMPLER_H_

#include "common.h"
#include "factor_graph.h"
#include "numa_nodes.h"
#include "timer.h"
#include <stdlib.h>
#include <thread>

namespace dd {

class GibbsSamplerThread;

/**
 * Class for a single NUMA node sampler
 */
class GibbsSampler {
 private:
  std::unique_ptr<CompactFactorGraph> pfg;
  std::unique_ptr<InferenceResult> pinfrs;
  NumaNodes numa_nodes_;
  std::vector<GibbsSamplerThread> workers;
  std::vector<std::thread> threads;

 public:
  CompactFactorGraph &fg;
  InferenceResult &infrs;

  // number of threads
  size_t nthread;
  // node id
  size_t nodeid;

  /**
   * Constructs a GibbsSampler given factor graph, number of threads, and
   * node id.
   */
  GibbsSampler(std::unique_ptr<CompactFactorGraph> pfg, const Weight weights[],
               const NumaNodes &numa_nodes, size_t nthread, size_t nodeid,
               const CmdParser &opts);

  /**
   * Performs sample
   */
  void sample(num_epochs_t i_epoch);

  /**
   * Performs SGD
   */
  void sample_sgd(double stepsize);

  /**
   * Waits for sample worker to finish
   */
  void wait();
};

/**
 * Class for single thread sampler
 */
class GibbsSamplerThread {
 private:
  // shard and variable id range assigned to this one
  variable_id_t start, end;

  // RNG seed
  unsigned short p_rand_seed[3];

  // potential for each proposals for categorical
  std::vector<double> varlen_potential_buffer_;

  // references and cached flags
  CompactFactorGraph &fg;
  InferenceResult &infrs;
  bool sample_evidence;
  bool learn_non_evidence;

 public:
  /**
   * Constructs a GibbsSamplerThread with given factor graph
   */
  GibbsSamplerThread(CompactFactorGraph &fg, InferenceResult &infrs,
                     size_t i_sharding, size_t n_sharding,
                     const CmdParser &opts);

  /**
   * Samples variables. The variables are divided into n_sharding equal
   * partitions
   * based on their ids. This function samples variables in the i_sharding-th
   * partition.
   */
  void sample();

  /**
   * Performs SGD with by sampling variables.  The variables are divided into
   * n_sharding equal partitions based on their ids. This function samples
   * variables
   * in the i_sharding-th partition.
   */
  void sample_sgd(double stepsize);

  /**
   * Performs SGD by sampling a single variable with id vid
   */
  inline void sample_sgd_single_variable(variable_id_t vid, double stepsize);

  /**
   * Samples a single variable with id vid
   */
  inline void sample_single_variable(variable_id_t vid);

  // sample a single variable
  inline variable_value_t draw_sample(Variable &variable,
                                      const variable_value_t assignments[],
                                      const double weight_values[]);

  /**
    * Resets RNG seed to given values
    */
  void set_random_seed(unsigned short s0, unsigned short s1, unsigned short s2);
};

inline void GibbsSamplerThread::sample_sgd_single_variable(variable_id_t vid,
                                                           double stepsize) {
  // stochastic gradient ascent
  // gradient of weight = E[f|D] - E[f], where D is evidence variables,
  // f is the factor function, E[] is expectation. Expectation is calculated
  // using a sample of the variable.

  Variable &variable = fg.variables[vid];
  if (variable.is_observation) return;

  if (!learn_non_evidence && !variable.is_evid) return;

  variable_value_t proposal = 0;

  // sample the variable with evidence unchanged
  proposal = variable.is_evid
                 ? variable.assignment_evid
                 : draw_sample(variable, infrs.assignments_evid.get(),
                               infrs.weight_values.get());

  infrs.assignments_evid[variable.id] = proposal;

  // sample the variable regardless of whether it's evidence
  proposal = draw_sample(variable, infrs.assignments_free.get(),
                         infrs.weight_values.get());
  infrs.assignments_free[variable.id] = proposal;

  fg.update_weight(variable, infrs, stepsize);
}

inline void GibbsSamplerThread::sample_single_variable(variable_id_t vid) {
  // this function uses the same sampling technique as in
  // sample_sgd_single_variable

  Variable &variable = fg.variables[vid];
  if (variable.is_observation) return;

  if (!variable.is_evid || sample_evidence) {
    variable_value_t proposal = draw_sample(
        variable, infrs.assignments_evid.get(), infrs.weight_values.get());
    infrs.assignments_evid[variable.id] = proposal;

    // bookkeep aggregates for computing marginals
    ++infrs.agg_nsamples[variable.id];
    switch (variable.domain_type) {
      case DTYPE_BOOLEAN:
        infrs.agg_means[variable.id] += proposal;
        break;<commit_msg>Fix broken Gibbs chain in SGD<commit_after>#ifndef DIMMWITTED_GIBBS_SAMPLER_H_
#define DIMMWITTED_GIBBS_SAMPLER_H_

#include "common.h"
#include "factor_graph.h"
#include "numa_nodes.h"
#include "timer.h"
#include <stdlib.h>
#include <thread>

namespace dd {

class GibbsSamplerThread;

/**
 * Class for a single NUMA node sampler
 */
class GibbsSampler {
 private:
  std::unique_ptr<CompactFactorGraph> pfg;
  std::unique_ptr<InferenceResult> pinfrs;
  NumaNodes numa_nodes_;
  std::vector<GibbsSamplerThread> workers;
  std::vector<std::thread> threads;

 public:
  CompactFactorGraph &fg;
  InferenceResult &infrs;

  // number of threads
  size_t nthread;
  // node id
  size_t nodeid;

  /**
   * Constructs a GibbsSampler given factor graph, number of threads, and
   * node id.
   */
  GibbsSampler(std::unique_ptr<CompactFactorGraph> pfg, const Weight weights[],
               const NumaNodes &numa_nodes, size_t nthread, size_t nodeid,
               const CmdParser &opts);

  /**
   * Performs sample
   */
  void sample(num_epochs_t i_epoch);

  /**
   * Performs SGD
   */
  void sample_sgd(double stepsize);

  /**
   * Waits for sample worker to finish
   */
  void wait();
};

/**
 * Class for single thread sampler
 */
class GibbsSamplerThread {
 private:
  // shard and variable id range assigned to this one
  variable_id_t start, end;

  // RNG seed
  unsigned short p_rand_seed[3];

  // potential for each proposals for categorical
  std::vector<double> varlen_potential_buffer_;

  // references and cached flags
  CompactFactorGraph &fg;
  InferenceResult &infrs;
  bool sample_evidence;
  bool learn_non_evidence;

 public:
  /**
   * Constructs a GibbsSamplerThread with given factor graph
   */
  GibbsSamplerThread(CompactFactorGraph &fg, InferenceResult &infrs,
                     size_t i_sharding, size_t n_sharding,
                     const CmdParser &opts);

  /**
   * Samples variables. The variables are divided into n_sharding equal
   * partitions
   * based on their ids. This function samples variables in the i_sharding-th
   * partition.
   */
  void sample();

  /**
   * Performs SGD with by sampling variables.  The variables are divided into
   * n_sharding equal partitions based on their ids. This function samples
   * variables
   * in the i_sharding-th partition.
   */
  void sample_sgd(double stepsize);

  /**
   * Performs SGD by sampling a single variable with id vid
   */
  inline void sample_sgd_single_variable(variable_id_t vid, double stepsize);

  /**
   * Samples a single variable with id vid
   */
  inline void sample_single_variable(variable_id_t vid);

  // sample a single variable
  inline variable_value_t draw_sample(Variable &variable,
                                      const variable_value_t assignments[],
                                      const double weight_values[]);

  /**
    * Resets RNG seed to given values
    */
  void set_random_seed(unsigned short s0, unsigned short s1, unsigned short s2);
};

inline void GibbsSamplerThread::sample_sgd_single_variable(variable_id_t vid,
                                                           double stepsize) {
  // stochastic gradient ascent
  // gradient of weight = E[f|D] - E[f], where D is evidence variables,
  // f is the factor function, E[] is expectation. Expectation is calculated
  // using a sample of the variable.

  Variable &variable = fg.variables[vid];
  if (variable.is_observation) return;

  // sample the variable
  variable_value_t proposal = draw_sample(
      variable, infrs.assignments_evid.get(), infrs.weight_values.get());

  infrs.assignments_free[variable.id] = proposal;

  infrs.assignments_evid[variable.id] = (
    variable.is_evid ? variable.assignment_evid : proposal);

  if (!learn_non_evidence && !variable.is_evid) return;

  fg.update_weight(variable, infrs, stepsize);
}

inline void GibbsSamplerThread::sample_single_variable(variable_id_t vid) {
  // this function uses the same sampling technique as in
  // sample_sgd_single_variable

  Variable &variable = fg.variables[vid];
  if (variable.is_observation) return;

  if (!variable.is_evid || sample_evidence) {
    variable_value_t proposal = draw_sample(
        variable, infrs.assignments_evid.get(), infrs.weight_values.get());
    infrs.assignments_evid[variable.id] = proposal;

    // bookkeep aggregates for computing marginals
    ++infrs.agg_nsamples[variable.id];
    switch (variable.domain_type) {
      case DTYPE_BOOLEAN:
        infrs.agg_means[variable.id] += proposal;
        break;
{'fraction_non_alphanumeric': 0.08894359413464625, 'fraction_numerical': 0.0009632880231189126, 'mean_word_length': 3.121746801940891, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '71036', 'n_tokens_mistral': 2949, 'n_tokens_neox': 2831, 'n_words': 782}