Add method to fetch constants as tensors

2 files changed, 9 insertions, 0 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();