Add method to fetch constants as tensors

8 files changed, 23 insertions, 14 deletions

diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/ImportedModel.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/ImportedModel.java
index 47fe66dd424..f58357cb874 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/ImportedModel.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/ImportedModel.java
@@ -82,6 +82,9 @@ public class ImportedModel implements ImportedMlModel {
     @Override
     public Map<String, String> smallConstants() { return asStrings(smallConstants); }
 
+    @Override
+    public Map<String, Tensor> smallConstantValues() { return ImmutableMap.copyOf(smallConstants); }
+
     boolean hasSmallConstant(String name) { return smallConstants.containsKey(name); }
 
     /**
@@ -92,6 +95,9 @@ public class ImportedModel implements ImportedMlModel {
     @Override
     public Map<String, String> largeConstants() { return asStrings(largeConstants); }
 
+    @Override
+    public Map<String, Tensor> largeConstantValues() { return ImmutableMap.copyOf(largeConstants); }
+
     boolean hasLargeConstant(String name) { return largeConstants.containsKey(name); }
 
     /**
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/configmodelview/ImportedMlModel.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/configmodelview/ImportedMlModel.java
index e40a06af042..8e53d477050 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/configmodelview/ImportedMlModel.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/configmodelview/ImportedMlModel.java
@@ -1,6 +1,7 @@
 // Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.rankingexpression.importer.configmodelview;
 
+import com.yahoo.tensor.Tensor;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -16,7 +17,9 @@ public interface ImportedMlModel {
     String source();
     Optional<String> inputTypeSpec(String input);
     Map<String, String> smallConstants();
+    Map<String, Tensor> smallConstantValues();
     Map<String, String> largeConstants();
+    Map<String, Tensor> largeConstantValues();
     Map<String, String> functions();
     List<ImportedMlFunction> outputExpressions();
 
diff --git a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/onnx/OnnxMnistSoftmaxImportTestCase.java b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/onnx/OnnxMnistSoftmaxImportTestCase.java
index 35c853bd746..81851f5658c 100644
--- a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/onnx/OnnxMnistSoftmaxImportTestCase.java
+++ b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/onnx/OnnxMnistSoftmaxImportTestCase.java
@@ -29,13 +29,13 @@ public class OnnxMnistSoftmaxImportTestCase {
         // Check constants
         assertEquals(2, model.largeConstants().size());
 
-        Tensor constant0 = Tensor.from(model.largeConstants().get("test_Variable"));
+        Tensor constant0 = model.largeConstantValues().get("test_Variable");
         assertNotNull(constant0);
         assertEquals(new TensorType.Builder(TensorType.Value.FLOAT).indexed("d2", 784).indexed("d1", 10).build(),
                      constant0.type());
         assertEquals(7840, constant0.size());
 
-        Tensor constant1 = Tensor.from(model.largeConstants().get("test_Variable_1"));
+        Tensor constant1 = model.largeConstantValues().get("test_Variable_1");
         assertNotNull(constant1);
         assertEquals(new TensorType.Builder(TensorType.Value.FLOAT).indexed("d1", 10).build(), constant1.type());
         assertEquals(10, constant1.size());
@@ -84,8 +84,8 @@ public class OnnxMnistSoftmaxImportTestCase {
 
     private Context contextFrom(ImportedModel result) {
         MapContext context = new MapContext();
-        result.largeConstants().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(Tensor.from(tensor))));
-        result.smallConstants().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(Tensor.from(tensor))));
+        result.largeConstantValues().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(tensor)));
+        result.smallConstantValues().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(tensor)));
         return context;
     }
 
diff --git a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/onnx/TestableModel.java b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/onnx/TestableModel.java
index c5355ebdf6f..0283244664f 100644
--- a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/onnx/TestableModel.java
+++ b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/onnx/TestableModel.java
@@ -103,8 +103,8 @@ public class TestableModel {
 
     static Context contextFrom(ImportedModel result) {
         TestableModelContext context = new TestableModelContext();
-        result.largeConstants().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(Tensor.from(tensor))));
-        result.smallConstants().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(Tensor.from(tensor))));
+        result.largeConstantValues().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(tensor)));
+        result.smallConstantValues().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(tensor)));
         return context;
     }
 
diff --git a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/tensorflow/RegressionTestCase.java b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/tensorflow/RegressionTestCase.java
index 46ced6f42ad..02feba9614b 100644
--- a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/tensorflow/RegressionTestCase.java
+++ b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/tensorflow/RegressionTestCase.java
@@ -25,13 +25,13 @@ public class RegressionTestCase {
         // Check constants
         Assert.assertEquals(2, model.get().largeConstants().size());
 
-        Tensor constant0 = Tensor.from(model.get().largeConstants().get("test_Variable_read"));
+        Tensor constant0 = model.get().largeConstantValues().get("test_Variable_read");
         assertNotNull(constant0);
         assertEquals(new TensorType.Builder().indexed("d2", 1536).indexed("d1", 14).build(),
                      constant0.type());
         assertEquals(21504, constant0.size());
 
-        Tensor constant1 = Tensor.from(model.get().largeConstants().get("test_Variable_1_read"));
+        Tensor constant1 = model.get().largeConstantValues().get("test_Variable_1_read");
         assertNotNull(constant1);
         assertEquals(new TensorType.Builder().indexed("d1", 14).build(), constant1.type());
         assertEquals(14, constant1.size());
diff --git a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/tensorflow/TensorFlowMnistSoftmaxImportTestCase.java b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/tensorflow/TensorFlowMnistSoftmaxImportTestCase.java
index 50e24f20972..d7572bd7b6b 100644
--- a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/tensorflow/TensorFlowMnistSoftmaxImportTestCase.java
+++ b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/tensorflow/TensorFlowMnistSoftmaxImportTestCase.java
@@ -25,13 +25,13 @@ public class TensorFlowMnistSoftmaxImportTestCase {
         // Check constants
         Assert.assertEquals(2, model.get().largeConstants().size());
 
-        Tensor constant0 = Tensor.from(model.get().largeConstants().get("test_Variable_read"));
+        Tensor constant0 = model.get().largeConstantValues().get("test_Variable_read");
         assertNotNull(constant0);
         assertEquals(new TensorType.Builder().indexed("d2", 784).indexed("d1", 10).build(),
                      constant0.type());
         assertEquals(7840, constant0.size());
 
-        Tensor constant1 = Tensor.from(model.get().largeConstants().get("test_Variable_1_read"));
+        Tensor constant1 = model.get().largeConstantValues().get("test_Variable_1_read");
         assertNotNull(constant1);
         assertEquals(new TensorType.Builder().indexed("d1", 10).build(),
                      constant1.type());
diff --git a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/tensorflow/TestableTensorFlowModel.java b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/tensorflow/TestableTensorFlowModel.java
index 41f343dbdaa..6579101fc6b 100644
--- a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/tensorflow/TestableTensorFlowModel.java
+++ b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/tensorflow/TestableTensorFlowModel.java
@@ -103,8 +103,8 @@ public class TestableTensorFlowModel {
 
     static Context contextFrom(ImportedModel result) {
         TestableModelContext context = new TestableModelContext();
-        result.largeConstants().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(Tensor.from(tensor))));
-        result.smallConstants().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(Tensor.from(tensor))));
+        result.largeConstantValues().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(tensor)));
+        result.smallConstantValues().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(tensor)));
         return context;
     }
 
diff --git a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/tensorflow/Tf2OnnxImportTestCase.java b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/tensorflow/Tf2OnnxImportTestCase.java
index 0510a433dd9..aed3f349240 100644
--- a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/tensorflow/Tf2OnnxImportTestCase.java
+++ b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/tensorflow/Tf2OnnxImportTestCase.java
@@ -44,8 +44,8 @@ public class Tf2OnnxImportTestCase {
 
     private Context contextFrom(ImportedModel result) {
         MapContext context = new MapContext();
-        result.largeConstants().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(Tensor.from(tensor))));
-        result.smallConstants().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(Tensor.from(tensor))));
+        result.largeConstantValues().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(tensor)));
+        result.smallConstantValues().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(tensor)));
         return context;
     }