ONNXModelRunner.h

Go to the documentation of this file.

//=== MLModelRunner/ONNXModelRunner/ONNXModelRunner.h - C++ --------------===//
//
// Part of the MLCompilerBridge Project, under the Apache License v2.0 with LLVM
// Exceptions. See the LICENSE file for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//===----------------------------------------------------------------------===//
 
#ifndef ONNX_MODELRUNNER_H
#define ONNX_MODELRUNNER_H
 
#include "MLModelRunner/MLModelRunner.h"
#include "MLModelRunner/ONNXModelRunner/agent.h"
#include "MLModelRunner/ONNXModelRunner/environment.h"
 
namespace MLBridge {
 
class ONNXModelRunner : public MLModelRunner {
public:
  ONNXModelRunner(MLBridge::Environment *env,
                  std::map<std::string, Agent *> agents,
                  llvm::LLVMContext *Ctx = nullptr);
 
  void setEnvironment(MLBridge::Environment *_env) { env = _env; }
  MLBridge::Environment *getEnvironment() { return env; }
 
  void addAgent(Agent *agent, std::string name);
 
  void requestExit() override {}
 
private:
  MLBridge::Environment *env;
  std::map<std::string, Agent *> agents;
  void computeAction(Observation &obs);
  void *evaluateUntyped() override;
};
} // namespace MLBridge
#endif // ONNX_MODELRUNNER_H