<commit_before>#ifndef NLP_SAFT_OPENSOURCE_DRAGNN_CORE_BEAM_H_
#define NLP_SAFT_OPENSOURCE_DRAGNN_CORE_BEAM_H_

#include <algorithm>
#include <memory>
#include <vector>

#include "dragnn/core/interfaces/cloneable_transition_state.h"
#include "dragnn/core/interfaces/transition_state.h"
#include "tensorflow/core/platform/logging.h"

namespace syntaxnet {
namespace dragnn {

// The Beam class wraps the logic necessary to advance a set of transition
// states for an arbitrary Component. Because the Beam class is generic, it
// doesn't know how to act on the states it is provided - the instantiating
// Component is expected to provide it the three functions it needs to interact
// with that Component's TransitionState subclasses.

template <typename T>
class Beam {
 public:
  // Creates a new Beam which can grow up to max_size elements.
  explicit Beam(int max_size) : max_size_(max_size), num_steps_(0) {
    VLOG(2) << "Creating beam with max size " << max_size_;
    static_assert(
        std::is_base_of<CloneableTransitionState<T>, T>::value,
        "This class must be instantiated to use a CloneableTransitionState");
  }

  // Sets the Beam functions, as follows:
  // bool is_allowed(TransitionState *, int): Return true if transition 'int' is
  //   allowed for transition state 'TransitionState *'.
  // void perform_transition(TransitionState *, int): Performs transition 'int'
  //   on transition state 'TransitionState *'.
  // int oracle_function(TransitionState *): Returns the oracle-specified action
  //   for transition state 'TransitionState *'.
  void SetFunctions(std::function<bool(T *, int)> is_allowed,
                    std::function<bool(T *)> is_final,
                    std::function<void(T *, int)> perform_transition,
                    std::function<int(T *)> oracle_function) {
    is_allowed_ = is_allowed;
    is_final_ = is_final;
    perform_transition_ = perform_transition;
    oracle_function_ = oracle_function;
  }

  // Resets the Beam and initializes it with the given set of states. The Beam
  // takes ownership of these TransitionStates.
  void Init(std::vector<std::unique_ptr<T>> initial_states) {
    VLOG(2) << "Initializing beam. Beam max size is " << max_size_;
    CHECK_LE(initial_states.size(), max_size_)
        << "Attempted to initialize a beam with more states ("
        << initial_states.size() << ") than the max size " << max_size_;
    beam_ = std::move(initial_states);
    std::vector<int> previous_beam_indices(max_size_, -1);
    for (int i = 0; i < beam_.size(); ++i) {
      previous_beam_indices.at(i) = beam_[i]->ParentBeamIndex();
      beam_[i]->SetBeamIndex(i);
    }
    beam_index_history_.emplace_back(previous_beam_indices);
  }

  // Advances the Beam from the given transition matrix.
  void AdvanceFromPrediction(const float transition_matrix[], int matrix_length,
                             int num_actions) {
    // Ensure that the transition matrix is the correct size. All underlying
    // states should have the same transition profile, so using the one at 0
    // should be safe.
    CHECK_EQ(matrix_length, max_size_ * num_actions)
        << "Transition matrix size does not match max beam size * number of "
           "state transitions!";

    if (max_size_ == 1) {
      // In the case where beam size is 1, we can advance by simply finding the
      // highest score and advancing the beam state in place.
      VLOG(2) << "Beam size is 1. Using fast beam path.";
      int best_action = -1;
      float best_score = -INFINITY;
      auto &state = beam_[0];
      for (int action_idx = 0; action_idx < num_actions; ++action_idx) {
        if (is_allowed_(state.get(), action_idx) &&
            transition_matrix[action_idx] > best_score) {
          best_score = transition_matrix[action_idx];
          best_action = action_idx;
        }
      }
      CHECK_GE(best_action, 0) << "Num actions: " << num_actions
                               << " score[0]: " << transition_matrix[0];
      perform_transition_(state.get(), best_action);
      const float new_score = state->GetScore() + best_score;
      state->SetScore(new_score);
      state->SetBeamIndex(0);
    } else {
      // Create the vector of all possible transitions, along with their scores.
      std::vector<Transition> candidates;

      // Iterate through all beams, examining all actions for each beam.
      for (int beam_idx = 0; beam_idx < beam_.size(); ++beam_idx) {
        const auto &state = beam_[beam_idx];
        for (int action_idx = 0; action_idx < num_actions; ++action_idx) {
          // If the action is allowed, calculate the proposed new score and add
          // the candidate action to the vector of all actions at this state.
          if (is_allowed_(state.get(), action_idx)) {
            Transition candidate;

            // The matrix is laid out by beam index, with a linear set of
            // actions for that index - so beam N's actions start at [nr. of
            // actions]*[N].
            const int matrix_idx = action_idx + beam_idx * num_actions;
            CHECK_LT(matrix_idx, matrix_length)
                << "Matrix index out of bounds!";
            const double score_delta = transition_matrix[matrix_idx];
            CHECK(!isnan(score_delta));
            candidate.source_idx = beam_idx;
            candidate.action = action_idx;<commit_msg>SyntaxNet does not compile, because isnan is unknown. (#1229)<commit_after>#ifndef NLP_SAFT_OPENSOURCE_DRAGNN_CORE_BEAM_H_
#define NLP_SAFT_OPENSOURCE_DRAGNN_CORE_BEAM_H_

#include <algorithm>
#include <cmath>
#include <memory>
#include <vector>

#include "dragnn/core/interfaces/cloneable_transition_state.h"
#include "dragnn/core/interfaces/transition_state.h"
#include "tensorflow/core/platform/logging.h"

namespace syntaxnet {
namespace dragnn {

// The Beam class wraps the logic necessary to advance a set of transition
// states for an arbitrary Component. Because the Beam class is generic, it
// doesn't know how to act on the states it is provided - the instantiating
// Component is expected to provide it the three functions it needs to interact
// with that Component's TransitionState subclasses.

template <typename T>
class Beam {
 public:
  // Creates a new Beam which can grow up to max_size elements.
  explicit Beam(int max_size) : max_size_(max_size), num_steps_(0) {
    VLOG(2) << "Creating beam with max size " << max_size_;
    static_assert(
        std::is_base_of<CloneableTransitionState<T>, T>::value,
        "This class must be instantiated to use a CloneableTransitionState");
  }

  // Sets the Beam functions, as follows:
  // bool is_allowed(TransitionState *, int): Return true if transition 'int' is
  //   allowed for transition state 'TransitionState *'.
  // void perform_transition(TransitionState *, int): Performs transition 'int'
  //   on transition state 'TransitionState *'.
  // int oracle_function(TransitionState *): Returns the oracle-specified action
  //   for transition state 'TransitionState *'.
  void SetFunctions(std::function<bool(T *, int)> is_allowed,
                    std::function<bool(T *)> is_final,
                    std::function<void(T *, int)> perform_transition,
                    std::function<int(T *)> oracle_function) {
    is_allowed_ = is_allowed;
    is_final_ = is_final;
    perform_transition_ = perform_transition;
    oracle_function_ = oracle_function;
  }

  // Resets the Beam and initializes it with the given set of states. The Beam
  // takes ownership of these TransitionStates.
  void Init(std::vector<std::unique_ptr<T>> initial_states) {
    VLOG(2) << "Initializing beam. Beam max size is " << max_size_;
    CHECK_LE(initial_states.size(), max_size_)
        << "Attempted to initialize a beam with more states ("
        << initial_states.size() << ") than the max size " << max_size_;
    beam_ = std::move(initial_states);
    std::vector<int> previous_beam_indices(max_size_, -1);
    for (int i = 0; i < beam_.size(); ++i) {
      previous_beam_indices.at(i) = beam_[i]->ParentBeamIndex();
      beam_[i]->SetBeamIndex(i);
    }
    beam_index_history_.emplace_back(previous_beam_indices);
  }

  // Advances the Beam from the given transition matrix.
  void AdvanceFromPrediction(const float transition_matrix[], int matrix_length,
                             int num_actions) {
    // Ensure that the transition matrix is the correct size. All underlying
    // states should have the same transition profile, so using the one at 0
    // should be safe.
    CHECK_EQ(matrix_length, max_size_ * num_actions)
        << "Transition matrix size does not match max beam size * number of "
           "state transitions!";

    if (max_size_ == 1) {
      // In the case where beam size is 1, we can advance by simply finding the
      // highest score and advancing the beam state in place.
      VLOG(2) << "Beam size is 1. Using fast beam path.";
      int best_action = -1;
      float best_score = -INFINITY;
      auto &state = beam_[0];
      for (int action_idx = 0; action_idx < num_actions; ++action_idx) {
        if (is_allowed_(state.get(), action_idx) &&
            transition_matrix[action_idx] > best_score) {
          best_score = transition_matrix[action_idx];
          best_action = action_idx;
        }
      }
      CHECK_GE(best_action, 0) << "Num actions: " << num_actions
                               << " score[0]: " << transition_matrix[0];
      perform_transition_(state.get(), best_action);
      const float new_score = state->GetScore() + best_score;
      state->SetScore(new_score);
      state->SetBeamIndex(0);
    } else {
      // Create the vector of all possible transitions, along with their scores.
      std::vector<Transition> candidates;

      // Iterate through all beams, examining all actions for each beam.
      for (int beam_idx = 0; beam_idx < beam_.size(); ++beam_idx) {
        const auto &state = beam_[beam_idx];
        for (int action_idx = 0; action_idx < num_actions; ++action_idx) {
          // If the action is allowed, calculate the proposed new score and add
          // the candidate action to the vector of all actions at this state.
          if (is_allowed_(state.get(), action_idx)) {
            Transition candidate;

            // The matrix is laid out by beam index, with a linear set of
            // actions for that index - so beam N's actions start at [nr. of
            // actions]*[N].
            const int matrix_idx = action_idx + beam_idx * num_actions;
            CHECK_LT(matrix_idx, matrix_length)
                << "Matrix index out of bounds!";
            const double score_delta = transition_matrix[matrix_idx];
            CHECK(!std::isnan(score_delta));
            candidate.source_idx = beam_idx;
            candidate.action = action_idx;
{'fraction_non_alphanumeric': 0.09867207672445592, 'fraction_numerical': 0.0038731095536702325, 'mean_word_length': 2.9493809176984707, 'pattern_counts': {'":': 0, '<': 100, '<?xml version=': 0, '>': 44, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1436759', 'n_tokens_mistral': 3144, 'n_tokens_neox': 3064, 'n_words': 995}