onnx ranking feature

8 files changed, 64 insertions, 40 deletions

diff --git a/eval/src/tests/tensor/onnx_wrapper/onnx_wrapper_test.cpp b/eval/src/tests/tensor/onnx_wrapper/onnx_wrapper_test.cpp
index 07e065f9e39..28a4a34b2e4 100644
--- a/eval/src/tests/tensor/onnx_wrapper/onnx_wrapper_test.cpp
+++ b/eval/src/tests/tensor/onnx_wrapper/onnx_wrapper_test.cpp
@@ -19,33 +19,10 @@ 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";
 
-vespalib::string to_str(const std::vector<size_t> &dim_sizes) {
-    vespalib::string res;
-    for (size_t dim_size: dim_sizes) {
-        if (dim_size == 0) {
-            res += "[]";
-        } else {
-            res += fmt("[%zu]", dim_size);
-        }
-    }
-    return res;
-}
-
-vespalib::string to_str(OnnxWrapper::TensorInfo::ElementType element_type) {
-    if (element_type == OnnxWrapper::TensorInfo::ElementType::FLOAT) {
-        return "float";
-    }
-    if (element_type == OnnxWrapper::TensorInfo::ElementType::DOUBLE) {
-        return "double";
-    }
-    return "???";
-}
-
 void dump_info(const char *ctx, const std::vector<OnnxWrapper::TensorInfo> &info) {
     fprintf(stderr, "%s:\n", ctx);
     for (size_t i = 0; i < info.size(); ++i) {
-        fprintf(stderr, "  %s[%zu]: '%s' %s%s\n", ctx, i, info[i].name.c_str(),
-                to_str(info[i].elements).c_str(),to_str(info[i].dimensions).c_str());
+        fprintf(stderr, "  %s[%zu]: '%s' %s\n", ctx, i, info[i].name.c_str(), info[i].type_as_string().c_str());
     }
 }
 
@@ -57,21 +34,17 @@ TEST(OnnxWrapperTest, onnx_model_can_be_inspected)
     ASSERT_EQ(wrapper.inputs().size(), 3);
     ASSERT_EQ(wrapper.outputs().size(), 1);
     //-------------------------------------------------------------------------
-    EXPECT_EQ(       wrapper.inputs()[0].name,        "query_tensor");
-    EXPECT_EQ(to_str(wrapper.inputs()[0].dimensions), "[1][4]");
-    EXPECT_EQ(to_str(wrapper.inputs()[0].elements),   "float");
+    EXPECT_EQ(wrapper.inputs()[0].name,             "query_tensor");
+    EXPECT_EQ(wrapper.inputs()[0].type_as_string(), "float[1][4]");
     //-------------------------------------------------------------------------
-    EXPECT_EQ(       wrapper.inputs()[1].name,        "attribute_tensor");
-    EXPECT_EQ(to_str(wrapper.inputs()[1].dimensions), "[4][1]");
-    EXPECT_EQ(to_str(wrapper.inputs()[1].elements),   "float");
+    EXPECT_EQ(wrapper.inputs()[1].name,             "attribute_tensor");
+    EXPECT_EQ(wrapper.inputs()[1].type_as_string(), "float[4][1]");
     //-------------------------------------------------------------------------
-    EXPECT_EQ(       wrapper.inputs()[2].name,        "bias_tensor");
-    EXPECT_EQ(to_str(wrapper.inputs()[2].dimensions), "[1][1]");
-    EXPECT_EQ(to_str(wrapper.inputs()[2].elements),   "float");
+    EXPECT_EQ(wrapper.inputs()[2].name,             "bias_tensor");
+    EXPECT_EQ(wrapper.inputs()[2].type_as_string(), "float[1][1]");
     //-------------------------------------------------------------------------
-    EXPECT_EQ(       wrapper.outputs()[0].name,        "output");
-    EXPECT_EQ(to_str(wrapper.outputs()[0].dimensions), "[1][1]");
-    EXPECT_EQ(to_str(wrapper.outputs()[0].elements),   "float");
+    EXPECT_EQ(wrapper.outputs()[0].name,             "output");
+    EXPECT_EQ(wrapper.outputs()[0].type_as_string(), "float[1][1]");
 }
 
 TEST(OnnxWrapperTest, onnx_model_can_be_evaluated)
diff --git a/eval/src/vespa/eval/eval/tensor_spec.cpp b/eval/src/vespa/eval/eval/tensor_spec.cpp
index a24b45dc3f5..b4b2e3d3afc 100644
--- a/eval/src/vespa/eval/eval/tensor_spec.cpp
+++ b/eval/src/vespa/eval/eval/tensor_spec.cpp
@@ -1,6 +1,9 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
 #include "tensor_spec.h"
+#include "value.h"
+#include "tensor.h"
+#include "tensor_engine.h"
 #include <vespa/vespalib/util/stringfmt.h>
 #include <vespa/vespalib/data/slime/slime.h>
 #include <ostream>
@@ -94,6 +97,18 @@ TensorSpec::from_slime(const slime::Inspector &tensor)
     return spec;
 }
 
+TensorSpec
+TensorSpec::from_value(const eval::Value &value)
+{
+    if (const eval::Tensor *tensor = value.as_tensor()) {
+        return tensor->engine().to_spec(*tensor);
+    }
+    if (value.is_double()) {
+        return TensorSpec("double").add({}, value.as_double());
+    }
+    return TensorSpec("error");
+}
+
 bool
 operator==(const TensorSpec &lhs, const TensorSpec &rhs)
 {
diff --git a/eval/src/vespa/eval/eval/tensor_spec.h b/eval/src/vespa/eval/eval/tensor_spec.h
index 22aa47f5ddb..f4f116454d5 100644
--- a/eval/src/vespa/eval/eval/tensor_spec.h
+++ b/eval/src/vespa/eval/eval/tensor_spec.h
@@ -18,6 +18,8 @@ struct Inspector;
 
 namespace eval {
 
+class Value;
+
 /**
  * An implementation-independent specification of the type and
  * contents of a tensor.
@@ -78,6 +80,7 @@ public:
     vespalib::string to_string() const;
     void to_slime(slime::Cursor &tensor) const;
     static TensorSpec from_slime(const slime::Inspector &tensor);
+    static TensorSpec from_value(const eval::Value &value);
 };
 
 bool operator==(const TensorSpec &lhs, const TensorSpec &rhs);
diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor_view.h b/eval/src/vespa/eval/tensor/dense/dense_tensor_view.h
index 33183d267c1..93dd2dbedeb 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor_view.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_view.h
@@ -55,6 +55,7 @@ protected:
         : _typeRef(type_in),
           _cellsRef()
     {}
+    DenseTensorView(const DenseTensorView &rhs) : DenseTensorView(rhs._typeRef, rhs._cellsRef) {}
 
     void initCellsRef(TypedCells cells_in) {
         assert(_typeRef.cell_type() == cells_in.type);
diff --git a/eval/src/vespa/eval/tensor/dense/mutable_dense_tensor_view.h b/eval/src/vespa/eval/tensor/dense/mutable_dense_tensor_view.h
index db241ef6a2b..5e4a48462d7 100644
--- a/eval/src/vespa/eval/tensor/dense/mutable_dense_tensor_view.h
+++ b/eval/src/vespa/eval/tensor/dense/mutable_dense_tensor_view.h
@@ -17,6 +17,7 @@ private:
 
 public:
     MutableDenseTensorView(eval::ValueType type_in);
+    MutableDenseTensorView(MutableDenseTensorView &&) = default;
     void setCells(TypedCells cells_in) {
         initCellsRef(cells_in);
     }
diff --git a/eval/src/vespa/eval/tensor/dense/onnx_wrapper.cpp b/eval/src/vespa/eval/tensor/dense/onnx_wrapper.cpp
index fa0379473c9..125095ff23e 100644
--- a/eval/src/vespa/eval/tensor/dense/onnx_wrapper.cpp
+++ b/eval/src/vespa/eval/tensor/dense/onnx_wrapper.cpp
@@ -18,6 +18,16 @@ namespace vespalib::tensor {
 
 namespace {
 
+vespalib::string to_str(OnnxWrapper::TensorInfo::ElementType element_type) {
+    if (element_type == OnnxWrapper::TensorInfo::ElementType::FLOAT) {
+        return "float";
+    }
+    if (element_type == OnnxWrapper::TensorInfo::ElementType::DOUBLE) {
+        return "double";
+    }
+    return "???";
+}
+
 ValueType::CellType as_cell_type(OnnxWrapper::TensorInfo::ElementType type) {
     if (type == OnnxWrapper::TensorInfo::ElementType::FLOAT) {
         return ValueType::CellType::FLOAT;
@@ -134,6 +144,20 @@ OnnxWrapper::TensorInfo::make_compatible_type() const
     return ValueType::tensor_type(std::move(dim_list), as_cell_type(elements));
 }
 
+vespalib::string
+OnnxWrapper::TensorInfo::type_as_string() const
+{
+    vespalib::string res = to_str(elements);
+    for (size_t dim_size: dimensions) {
+        if (dim_size == 0) {
+            res += "[]";
+        } else {
+            res += fmt("[%zu]", dim_size);
+        }
+    }
+    return res;
+}
+
 OnnxWrapper::TensorInfo::~TensorInfo() = default;
 
 OnnxWrapper::Shared::Shared()
@@ -222,12 +246,14 @@ OnnxWrapper::OnnxWrapper(const vespalib::string &model_file, Optimize optimize)
 OnnxWrapper::~OnnxWrapper() = default;
 
 OnnxWrapper::Result
-OnnxWrapper::eval(const Params &params)
+OnnxWrapper::eval(const Params &params) const
 {
     assert(params.values.size() == _inputs.size());
     Ort::RunOptions run_opts(nullptr);
-    return Result(_session.Run(run_opts, _input_name_refs.data(), params.values.data(), _inputs.size(),
-                               _output_name_refs.data(), _outputs.size()));
+    // NB: Run requires non-const session
+    Ort::Session &session = const_cast<Ort::Session&>(_session);
+    return Result(session.Run(run_opts, _input_name_refs.data(), params.values.data(), _inputs.size(),
+                              _output_name_refs.data(), _outputs.size()));
 }
 
 }
diff --git a/eval/src/vespa/eval/tensor/dense/onnx_wrapper.h b/eval/src/vespa/eval/tensor/dense/onnx_wrapper.h
index 67a64f2d318..abe1da252c7 100644
--- a/eval/src/vespa/eval/tensor/dense/onnx_wrapper.h
+++ b/eval/src/vespa/eval/tensor/dense/onnx_wrapper.h
@@ -28,6 +28,7 @@ public:
         ElementType elements;
         bool is_compatible(const eval::ValueType &type) const;
         eval::ValueType make_compatible_type() const;
+        vespalib::string type_as_string() const;
         ~TensorInfo();
     };
 
@@ -48,6 +49,7 @@ public:
         std::vector<Ort::Value> values;
         Result(std::vector<Ort::Value> values_in) : values(std::move(values_in)) {}
     public:
+        static Result make_empty() { return Result({}); }
         size_t num_values() const { return values.size(); }
         void get(size_t idx, MutableDenseTensorView &dst);
     };
@@ -78,7 +80,7 @@ public:
     ~OnnxWrapper();
     const std::vector<TensorInfo> &inputs() const { return _inputs; }
     const std::vector<TensorInfo> &outputs() const { return _outputs; }
-    Result eval(const Params &params); // NB: Run requires non-const session
+    Result eval(const Params &params) const;
 };
 
 }
diff --git a/eval/src/vespa/eval/tensor/dense/typed_cells.h b/eval/src/vespa/eval/tensor/dense/typed_cells.h
index d1d2baa535e..6ea2b40689e 100644
--- a/eval/src/vespa/eval/tensor/dense/typed_cells.h
+++ b/eval/src/vespa/eval/tensor/dense/typed_cells.h
@@ -44,6 +44,9 @@ struct TypedCells {
         abort();
     }
 
+    TypedCells(TypedCells &&other) = default;
+    TypedCells(const TypedCells &other) = default;
+    TypedCells & operator= (TypedCells &&other) = default;
     TypedCells & operator= (const TypedCells &other) = default;
 };