summaryrefslogtreecommitdiffstats
path: root/searchlib
diff options
context:
space:
mode:
authorHåvard Pettersen <havardpe@oath.com>2020-08-20 12:09:57 +0000
committerHåvard Pettersen <havardpe@oath.com>2020-08-21 11:03:42 +0000
commit7e4d00f04c5cd5ca229a5559f05a069908f6144a (patch)
treeaa6de03f758780ff331bf3bb1658eebd85601aa8 /searchlib
parent67e528443cca68cc527e50c2714ad1717563c458 (diff)
onnx ranking feature
Diffstat (limited to 'searchlib')
-rw-r--r--searchlib/CMakeLists.txt1
-rw-r--r--searchlib/src/tests/features/onnx_feature/CMakeLists.txt9
-rw-r--r--searchlib/src/tests/features/onnx_feature/onnx_feature_test.cpp100
-rw-r--r--searchlib/src/vespa/searchlib/features/CMakeLists.txt1
-rw-r--r--searchlib/src/vespa/searchlib/features/onnx_feature.cpp119
-rw-r--r--searchlib/src/vespa/searchlib/features/onnx_feature.h31
-rw-r--r--searchlib/src/vespa/searchlib/features/setup.cpp5
7 files changed, 264 insertions, 2 deletions
diff --git a/searchlib/CMakeLists.txt b/searchlib/CMakeLists.txt
index 1db9018bd21..4e8d93e3f81 100644
--- a/searchlib/CMakeLists.txt
+++ b/searchlib/CMakeLists.txt
@@ -139,6 +139,7 @@ vespa_define_module(
src/tests/features/native_dot_product
src/tests/features/nns_closeness
src/tests/features/nns_distance
+ src/tests/features/onnx_feature
src/tests/features/ranking_expression
src/tests/features/raw_score
src/tests/features/subqueries
diff --git a/searchlib/src/tests/features/onnx_feature/CMakeLists.txt b/searchlib/src/tests/features/onnx_feature/CMakeLists.txt
new file mode 100644
index 00000000000..8657d8987da
--- /dev/null
+++ b/searchlib/src/tests/features/onnx_feature/CMakeLists.txt
@@ -0,0 +1,9 @@
+# Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(searchlib_onnx_feature_test_app TEST
+ SOURCES
+ onnx_feature_test.cpp
+ DEPENDS
+ searchlib
+ GTest::GTest
+)
+vespa_add_test(NAME searchlib_onnx_feature_test_app COMMAND searchlib_onnx_feature_test_app)
diff --git a/searchlib/src/tests/features/onnx_feature/onnx_feature_test.cpp b/searchlib/src/tests/features/onnx_feature/onnx_feature_test.cpp
new file mode 100644
index 00000000000..cc6b8e0ce29
--- /dev/null
+++ b/searchlib/src/tests/features/onnx_feature/onnx_feature_test.cpp
@@ -0,0 +1,100 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/vespalib/stllike/string.h>
+#include <vespa/vespalib/util/stringfmt.h>
+#include <vespa/searchlib/features/rankingexpressionfeature.h>
+#include <vespa/searchlib/features/onnx_feature.h>
+#include <vespa/searchlib/fef/blueprintfactory.h>
+#include <vespa/searchlib/fef/indexproperties.h>
+#include <vespa/searchlib/fef/matchdatalayout.h>
+#include <vespa/searchlib/fef/test/indexenvironment.h>
+#include <vespa/searchlib/fef/test/queryenvironment.h>
+#include <vespa/searchlib/fef/rank_program.h>
+#include <vespa/searchlib/fef/test/test_features.h>
+#include <vespa/vespalib/gtest/gtest.h>
+
+using namespace search::fef;
+using namespace search::fef::test;
+using namespace search::features;
+using vespalib::make_string_short::fmt;
+using vespalib::eval::TensorSpec;
+
+std::string get_source_dir() {
+ const char *dir = getenv("SOURCE_DIRECTORY");
+ return (dir ? dir : ".");
+}
+std::string source_dir = get_source_dir();
+std::string vespa_dir = source_dir + "/" + "../../../../..";
+std::string simple_model = vespa_dir + "/" + "model-integration/src/test/models/onnx/simple/simple.onnx";
+
+uint32_t default_docid = 1;
+
+vespalib::string expr_feature(const vespalib::string &name) {
+ return fmt("rankingExpression(%s)", name.c_str());
+}
+
+vespalib::string onnx_feature(const vespalib::string &name) {
+ return fmt("onnxModel(%s)", name.c_str());
+}
+
+struct OnnxFeatureTest : ::testing::Test {
+ BlueprintFactory factory;
+ IndexEnvironment indexEnv;
+ BlueprintResolver::SP resolver;
+ Properties overrides;
+ MatchData::UP match_data;
+ RankProgram program;
+ OnnxFeatureTest() : factory(), indexEnv(), resolver(new BlueprintResolver(factory, indexEnv)),
+ overrides(), match_data(), program(resolver)
+ {
+ factory.addPrototype(std::make_shared<DocidBlueprint>());
+ factory.addPrototype(std::make_shared<RankingExpressionBlueprint>());
+ factory.addPrototype(std::make_shared<OnnxBlueprint>());
+ }
+ void add_expr(const vespalib::string &name, const vespalib::string &expr) {
+ vespalib::string feature_name = expr_feature(name);
+ vespalib::string expr_name = feature_name + ".rankingScript";
+ indexEnv.getProperties().add(expr_name, expr);
+ }
+ void add_onnx(const vespalib::string &name, const vespalib::string &file) {
+ vespalib::string feature_name = onnx_feature(name);
+ vespalib::string file_name = feature_name + ".fileref";
+ indexEnv.getProperties().add(file_name, file);
+ }
+ void compile(const vespalib::string &seed) {
+ resolver->addSeed(seed);
+ ASSERT_TRUE(resolver->compile());
+ MatchDataLayout mdl;
+ QueryEnvironment queryEnv(&indexEnv);
+ match_data = mdl.createMatchData();
+ program.setup(*match_data, queryEnv, overrides);
+ }
+ TensorSpec get(const vespalib::string &feature, uint32_t docid) {
+ auto result = program.get_all_features(false);
+ for (size_t i = 0; i < result.num_features(); ++i) {
+ if (result.name_of(i) == feature) {
+ return TensorSpec::from_value(result.resolve(i).as_object(docid));
+ }
+ }
+ return TensorSpec("error");
+ }
+ TensorSpec get(uint32_t docid) {
+ auto result = program.get_seeds(false);
+ EXPECT_EQ(1u, result.num_features());
+ return TensorSpec::from_value(result.resolve(0).as_object(docid));
+ }
+};
+
+TEST_F(OnnxFeatureTest, simple_onnx_model_can_be_calculated) {
+ add_expr("query_tensor", "tensor<float>(a[1],b[4]):[[docid,2,3,4]]");
+ add_expr("attribute_tensor", "tensor<float>(a[4],b[1]):[[5],[6],[7],[8]]");
+ add_expr("bias_tensor", "tensor<float>(a[1],b[1]):[[9]]");
+ add_onnx("simple", simple_model);
+ compile(onnx_feature("simple"));
+ EXPECT_EQ(get(1), TensorSpec("tensor<float>(d0[1],d1[1])").add({{"d0",0},{"d1",0}}, 79.0));
+ EXPECT_EQ(get("onnxModel(simple).output", 1), TensorSpec("tensor<float>(d0[1],d1[1])").add({{"d0",0},{"d1",0}}, 79.0));
+ EXPECT_EQ(get(2), TensorSpec("tensor<float>(d0[1],d1[1])").add({{"d0",0},{"d1",0}}, 84.0));
+ EXPECT_EQ(get(3), TensorSpec("tensor<float>(d0[1],d1[1])").add({{"d0",0},{"d1",0}}, 89.0));
+}
+
+GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/searchlib/src/vespa/searchlib/features/CMakeLists.txt b/searchlib/src/vespa/searchlib/features/CMakeLists.txt
index 215b6ade9fd..93fead713f4 100644
--- a/searchlib/src/vespa/searchlib/features/CMakeLists.txt
+++ b/searchlib/src/vespa/searchlib/features/CMakeLists.txt
@@ -40,6 +40,7 @@ vespa_add_library(searchlib_features OBJECT
nativeproximityfeature.cpp
nativerankfeature.cpp
nowfeature.cpp
+ onnx_feature.cpp
proximityfeature.cpp
querycompletenessfeature.cpp
queryfeature.cpp
diff --git a/searchlib/src/vespa/searchlib/features/onnx_feature.cpp b/searchlib/src/vespa/searchlib/features/onnx_feature.cpp
new file mode 100644
index 00000000000..f6d5c37b61d
--- /dev/null
+++ b/searchlib/src/vespa/searchlib/features/onnx_feature.cpp
@@ -0,0 +1,119 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "onnx_feature.h"
+#include <vespa/searchlib/fef/properties.h>
+#include <vespa/searchlib/fef/featureexecutor.h>
+#include <vespa/eval/tensor/dense/onnx_wrapper.h>
+#include <vespa/eval/tensor/dense/dense_tensor_view.h>
+#include <vespa/eval/tensor/dense/mutable_dense_tensor_view.h>
+#include <vespa/vespalib/util/stringfmt.h>
+#include <vespa/vespalib/util/stash.h>
+
+#include <vespa/log/log.h>
+LOG_SETUP(".features.onnx_feature");
+
+using search::fef::Blueprint;
+using search::fef::FeatureExecutor;
+using search::fef::FeatureType;
+using search::fef::IIndexEnvironment;
+using search::fef::IQueryEnvironment;
+using search::fef::ParameterList;
+using vespalib::Stash;
+using vespalib::eval::ValueType;
+using vespalib::make_string_short::fmt;
+using vespalib::tensor::DenseTensorView;
+using vespalib::tensor::MutableDenseTensorView;
+using vespalib::tensor::OnnxWrapper;
+
+namespace search::features {
+
+/**
+ * Feature executor that evaluates an onnx model
+ */
+class OnnxFeatureExecutor : public FeatureExecutor
+{
+private:
+ const OnnxWrapper &_model;
+ OnnxWrapper::Params _params;
+ OnnxWrapper::Result _result;
+ std::vector<MutableDenseTensorView> _views;
+
+public:
+ OnnxFeatureExecutor(const OnnxWrapper &model)
+ : _model(model), _params(), _result(OnnxWrapper::Result::make_empty()), _views()
+ {
+ _views.reserve(_model.outputs().size());
+ for (const auto &output: _model.outputs()) {
+ _views.emplace_back(output.make_compatible_type());
+ }
+ }
+ bool isPure() override { return true; }
+ void execute(uint32_t) override {
+ _params = OnnxWrapper::Params();
+ for (size_t i = 0; i < _model.inputs().size(); ++i) {
+ _params.bind(i, static_cast<const DenseTensorView&>(inputs().get_object(i).get()));
+ }
+ _result = _model.eval(_params);
+ for (size_t i = 0; i < _model.outputs().size(); ++i) {
+ _result.get(i, _views[i]);
+ outputs().set_object(i, _views[i]);
+ }
+ }
+};
+
+OnnxBlueprint::OnnxBlueprint()
+ : Blueprint("onnxModel"),
+ _model(nullptr)
+{
+}
+
+OnnxBlueprint::~OnnxBlueprint() = default;
+
+bool
+OnnxBlueprint::setup(const IIndexEnvironment &env,
+ const ParameterList &params)
+{
+ auto optimize = (env.getFeatureMotivation() == env.FeatureMotivation::VERIFY_SETUP)
+ ? OnnxWrapper::Optimize::DISABLE
+ : OnnxWrapper::Optimize::ENABLE;
+
+ // Note: Using the fileref property with the model name as
+ // fallback to get a file name. This needs to be replaced with an
+ // actual file reference obtained through config when available.
+ vespalib::string file_name = env.getProperties().lookup(getName(), "fileref").get(params[0].getValue());
+ try {
+ _model = std::make_unique<OnnxWrapper>(file_name, optimize);
+ } catch (std::exception &ex) {
+ return fail("Model setup failed: %s", ex.what());
+ }
+ for (size_t i = 0; i < _model->inputs().size(); ++i) {
+ const auto &model_input = _model->inputs()[i];
+ if (auto maybe_input = defineInput(fmt("rankingExpression(\"%s\")", model_input.name.c_str()), AcceptInput::OBJECT)) {
+ const FeatureType &feature_input = maybe_input.value();
+ assert(feature_input.is_object());
+ if (!model_input.is_compatible(feature_input.type())) {
+ return fail("incompatible type for input '%s': %s -> %s", model_input.name.c_str(),
+ feature_input.type().to_spec().c_str(), model_input.type_as_string().c_str());
+ }
+ }
+ }
+ for (size_t i = 0; i < _model->outputs().size(); ++i) {
+ const auto &model_output = _model->outputs()[i];
+ ValueType output_type = model_output.make_compatible_type();
+ if (output_type.is_error()) {
+ return fail("unable to make compatible type for output '%s': %s -> error",
+ model_output.name.c_str(), model_output.type_as_string().c_str());
+ }
+ describeOutput(model_output.name, "output from onnx model", FeatureType::object(output_type));
+ }
+ return true;
+}
+
+FeatureExecutor &
+OnnxBlueprint::createExecutor(const IQueryEnvironment &, Stash &stash) const
+{
+ assert(_model);
+ return stash.create<OnnxFeatureExecutor>(*_model);
+}
+
+}
diff --git a/searchlib/src/vespa/searchlib/features/onnx_feature.h b/searchlib/src/vespa/searchlib/features/onnx_feature.h
new file mode 100644
index 00000000000..eb6e368ffbd
--- /dev/null
+++ b/searchlib/src/vespa/searchlib/features/onnx_feature.h
@@ -0,0 +1,31 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <vespa/searchlib/fef/blueprint.h>
+
+namespace vespalib::tensor { class OnnxWrapper; }
+
+namespace search::features {
+
+/**
+ * Blueprint for the ranking feature used to evaluate an onnx model.
+ **/
+class OnnxBlueprint : public fef::Blueprint {
+private:
+ std::unique_ptr<vespalib::tensor::OnnxWrapper> _model;
+public:
+ OnnxBlueprint();
+ ~OnnxBlueprint() override;
+ void visitDumpFeatures(const fef::IIndexEnvironment &, fef::IDumpFeatureVisitor &) const override {}
+ fef::Blueprint::UP createInstance() const override {
+ return Blueprint::UP(new OnnxBlueprint());
+ }
+ fef::ParameterDescriptions getDescriptions() const override {
+ return fef::ParameterDescriptions().desc().string();
+ }
+ bool setup(const fef::IIndexEnvironment &env, const fef::ParameterList &params) override;
+ fef::FeatureExecutor &createExecutor(const fef::IQueryEnvironment &env, vespalib::Stash &stash) const override;
+};
+
+}
diff --git a/searchlib/src/vespa/searchlib/features/setup.cpp b/searchlib/src/vespa/searchlib/features/setup.cpp
index bd79f1d4fb5..c97d2d68812 100644
--- a/searchlib/src/vespa/searchlib/features/setup.cpp
+++ b/searchlib/src/vespa/searchlib/features/setup.cpp
@@ -23,6 +23,7 @@
#include "flow_completeness_feature.h"
#include "foreachfeature.h"
#include "freshnessfeature.h"
+#include "global_sequence_feature.h"
#include "item_raw_score_feature.h"
#include "jarowinklerdistancefeature.h"
#include "matchcountfeature.h"
@@ -34,6 +35,7 @@
#include "nativeproximityfeature.h"
#include "nativerankfeature.h"
#include "nowfeature.h"
+#include "onnx_feature.h"
#include "proximityfeature.h"
#include "querycompletenessfeature.h"
#include "queryfeature.h"
@@ -53,7 +55,6 @@
#include "termfeature.h"
#include "terminfofeature.h"
#include "text_similarity_feature.h"
-#include "global_sequence_feature.h"
#include "valuefeature.h"
#include "max_reduce_prod_join_replacer.h"
@@ -123,7 +124,7 @@ void setup_search_features(fef::IBlueprintRegistry & registry)
registry.addPrototype(std::make_shared<TermFieldMdBlueprint>());
registry.addPrototype(std::make_shared<ConstantBlueprint>());
registry.addPrototype(std::make_shared<GlobalSequenceBlueprint>());
-
+ registry.addPrototype(std::make_shared<OnnxBlueprint>());
// Ranking Expression
auto replacers = std::make_unique<ListExpressionReplacer>();