InferenceSession.h

#pragma once
#include <string>
#include <memory>
#include <unordered_map>
#include <vector>
#include <unordered_set>
#include <onnx/onnx-ml.pb.h>   
#include <google/protobuf/message.h>
#include <google/protobuf/io/zero_copy_stream_impl.h>
#include <google/protobuf/io/coded_stream.h>
#include <boost/asio/thread_pool.hpp> 
#include <boost/asio/post.hpp>
#include <latch>
#include "graph.h"

class Graph;

class InferenceSession {
public:
  explicit InferenceSession(std::string model_path, size_t num_threads = 4)
      : model_path_(std::move(model_path)), thread_pool_(num_threads) {};

  bool LoadModel();
  bool BuildGraph();        

  onnx::ModelProto& GetModel() { return model_; }
  const onnx::ModelProto& GetModel() const { return model_; }

  Graph* graph() { return graph_.get(); }
  const Graph* graph() const { return graph_.get(); }

  std::latch& get_remaining_nodes() { return *remaining_nodes_; }
  boost::asio::thread_pool& get_thread_pool() { return thread_pool_; }

  std::unordered_map<std::string, Tensor<float>>& get_values_map() {
    return values_map_;
  }
  std::unordered_map<std::string, Tensor<float>>& get_weights_map() {
    return weights_map_;
  }

  std::vector<Tensor<float>> run(std::unordered_map<std::string, Tensor<float>>& input_tensors);

private:
  std::string model_path_;
  onnx::ModelProto model_;
  std::unique_ptr<Graph> graph_;
  std::unordered_set<std::string> weights_;
  std::vector<std::string> graph_input_names_;
  std::vector<std::string> graph_output_names_;
  std::unordered_map<std::string, Tensor<float>> values_map_;
  std::unordered_map<std::string, Tensor<float>> weights_map_;
  boost::asio::thread_pool thread_pool_;
  std::unique_ptr<std::latch> remaining_nodes_;
};