Line data    Source code

       1              : #include "triggeralgs/ProtoDUNEBSMWindow/CompiledModelInterface.hpp"
       2              : 
       3              : #include <iostream>
       4              : 
       5              : namespace triggeralgs {
       6              : 
       7            0 : CompiledModelInterface::CompiledModelInterface(int nbatch, bool is_pdvd) : num_batch(nbatch) {
       8            0 :   if (is_pdvd) {  
       9            0 :     model_ptr = std::make_unique<TreelitePDVDModel>();
      10              :   } else {
      11            0 :     model_ptr = std::make_unique<TreelitePDHDModel>();
      12              :   }
      13            0 : }
      14              : 
      15            0 : CompiledModelInterface::~CompiledModelInterface() {}
      16              :     
      17            0 : int CompiledModelInterface::GetNumFeatures() {
      18            0 :   return model_ptr->get_num_feature();
      19              : }
      20              : 
      21            0 : void CompiledModelInterface::Predict(Entry *input, float *result) {
      22            0 :   for (int rid = 0; rid < num_batch; ++rid) {
      23            0 :     model_ptr->predict(input, 0, result);
      24              :   }
      25            0 : }
      26              : 
      27            0 : bool CompiledModelInterface::Classify(const float *result, float bdt_threshold) {
      28            0 :   for (uint64_t rid = 0; rid < num_batch; rid++) {
      29            0 :     if (result[rid] > bdt_threshold) {
      30              :       return true;
      31              :     }
      32              :   }
      33              :   return false;
      34              : }
      35              : 
      36              : 
      37              : } // namespace triggeralgs