text: // Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
  //
  // Licensed under the Apache License, Version 2.0 (the "License");
  // you may not use this file except in compliance with the License.
  // You may obtain a copy of the License at
  //
  //     http://www.apache.org/licenses/LICENSE-2.0
  //
  // Unless required by applicable law or agreed to in writing, software
  // distributed under the License is distributed on an "AS IS" BASIS,
  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  // See the License for the specific language governing permissions and
  // limitations under the License.
  
  #pragma once
  #include <vector>
  
  #include "paddle/fluid/distributed/index_dataset/index_wrapper.h"
  #include "paddle/fluid/framework/data_feed.h"
  #include "paddle/fluid/framework/program_desc.h"
  #include "paddle/fluid/operators/math/sampler.h"
  #include "paddle/fluid/platform/enforce.h"
  
  namespace paddle {
  namespace distributed {
  
  class IndexSampler {
   public:
    virtual ~IndexSampler() {}
    IndexSampler() {}
  
    template <typename T>
    static std::shared_ptr<IndexSampler> Init(const std::string& name) {
      std::shared_ptr<IndexSampler> instance = nullptr;
      instance.reset(new T(name));
      return instance;
    }
  
    virtual void init_layerwise_conf(
        const std::vector<uint16_t>& layer_sample_counts,
        uint16_t start_sample_layer = 1, uint16_t seed = 0) {}
    virtual void init_beamsearch_conf(const int64_t k) {}
    virtual std::vector<std::vector<uint64_t>> sample(
        const std::vector<std::vector<uint64_t>>& user_inputs,
        const std::vector<uint64_t>& input_targets,
        bool with_hierarchy = false) = 0;
  
    virtual void sample_from_dataset(
        const uint16_t sample_slot,
        std::vector<paddle::framework::Record>* src_datas,
        std::vector<paddle::framework::Record>* sample_results) = 0;
  };
  
  class LayerWiseSampler : public IndexSampler {
   public:
    virtual ~LayerWiseSampler() {}
    explicit LayerWiseSampler(const std::string& name) {
      tree_ = IndexWrapper::GetInstance()->get_tree_index(name);
    }
  
    void init_layerwise_conf(const std::vector<uint16_t>& layer_sample_counts,
                             uint16_t start_sample_layer,
                             uint16_t seed) override {
      seed_ = seed;
      start_sample_layer_ = start_sample_layer;
  
      PADDLE_ENFORCE_GT(
          start_sample_layer_, 0,
          paddle::platform::errors::InvalidArgument(
              "start sampler layer = [%d], it should greater than 0.",
              start_sample_layer_));
      PADDLE_ENFORCE_LT(start_sample_layer_, tree_->Height(),
                        paddle::platform::errors::InvalidArgument(
                            "start sampler layer = [%d], it should less than "
                            "max_layer, which is [%d].",
                            start_sample_layer_, tree_->Height()));
  
      size_t i = 0;
      layer_counts_sum_ = 0;
      layer_counts_.clear();
      int cur_layer = start_sample_layer_;
      while (cur_layer < tree_->Height()) {
        int layer_sample_num = 1;
        if (i < layer_sample_counts.size()) {
          layer_sample_num = layer_sample_counts[i];
        }
        layer_counts_sum_ += layer_sample_num + 1;
        layer_counts_.push_back(layer_sample_num);
        VLOG(3) << "[INFO] level " << cur_layer
                << " sample_layer_counts.push_back: " << layer_sample_num;
        cur_layer += 1;
        i += 1;
      }
      reverse(layer_counts_.begin(), layer_counts_.end());
      VLOG(3) << "sample counts sum: " << layer_counts_sum_;
  
      auto max_layer = tree_->Height();
      sampler_vec_.clear();
      layer_ids_.clear();
  
      auto layer_index = max_layer - 1;
      size_t idx = 0;
      while (layer_index >= start_sample_layer_) {
        auto layer_codes = tree_->GetLayerCodes(layer_index);
        layer_ids_.push_back(tree_->GetNodes(layer_codes));
        auto sampler_temp =
            std::make_shared<paddle::operators::math::UniformSampler>(
                layer_ids_[idx].size() - 1, seed_);
        sampler_vec_.push_back(sampler_temp);
        layer_index--;
        idx++;
      }
    }
    std::vector<std::vector<uint64_t>> sample(
        const std::vector<std::vector<uint64_t>>& user_inputs,
        const std::vector<uint64_t>& target_ids, bool with_hierarchy) override;
  
    void sample_from_dataset(
        const uint16_t sample_slot,
        std::vector<paddle::framework::Record>* src_datas,
        std::vector<paddle::framework::Record>* sample_results) override;
  
   private:
    std::vector<int> layer_counts_;
    int64_t layer_counts_sum_{0};
    std::shared_ptr<TreeIndex> tree_{nullptr};
    int seed_{0};
    int start_sample_layer_{1};
    std::vector<std::shared_ptr<paddle::operators::math::Sampler>> sampler_vec_;
    std::vector<std::vector<IndexNode>> layer_ids_;
  };
  
  }  // end namespace distributed
  }  // end namespace paddle
  