Leave the existing interface as is and add 2 accessors for accessing the tensors directly.

6 files changed, 38 insertions, 19 deletions

diff --git a/config-model/src/main/java/com/yahoo/vespa/model/ml/ConvertedModel.java b/config-model/src/main/java/com/yahoo/vespa/model/ml/ConvertedModel.java
index 2344d3d15e0..96a6b39dc1a 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/ml/ConvertedModel.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/ml/ConvertedModel.java
@@ -199,8 +199,8 @@ public class ConvertedModel {
                                                                    ModelStore store) {
         // Add constants
         Set<String> constantsReplacedByFunctions = new HashSet<>();
-        model.smallConstants().forEach((k, v) -> transformSmallConstant(store, profile, k, v));
-        model.largeConstants().forEach((k, v) -> transformLargeConstant(store, profile, queryProfiles,
+        model.smallConstantTensors().forEach((k, v) -> transformSmallConstant(store, profile, k, v));
+        model.largeConstantTensors().forEach((k, v) -> transformLargeConstant(store, profile, queryProfiles,
                                                                         constantsReplacedByFunctions, k, v));
 
         // Add functions
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 edccc5544af..bb2484c0c28 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
@@ -81,21 +81,36 @@ public class ImportedModel implements ImportedMlModel {
     }
 
     /**
+     * Returns an immutable map of the small constants of this.
+     * These should have sizes up to a few kb at most, and correspond to constant values given in the source model.
+     */
+    @Override
+    public Map<String, Tensor> smallConstantTensors() { return Map.copyOf(smallConstants); }
+    /**
      * Returns an immutable map of the small constants of this, represented as strings on the standard tensor form.
      * These should have sizes up to a few kb at most, and correspond to constant values given in the source model.
+     * @deprecated Use smallConstantTensors instead
      */
     @Override
-    public Map<String, Tensor> smallConstants() { return Map.copyOf(smallConstants); }
+    @SuppressWarnings("removal")
+    public Map<String, String> smallConstants() { return asStrings(smallConstants); }
 
     boolean hasSmallConstant(String name) { return smallConstants.containsKey(name); }
 
     /**
      * Returns an immutable map of the large constants of this.
      * These can have sizes in gigabytes and must be distributed to nodes separately from configuration.
-     * For TensorFlow this corresponds to Variable files stored separately.
      */
     @Override
-    public Map<String, Tensor> largeConstants() { return Map.copyOf(largeConstants); }
+    public Map<String, Tensor> largeConstantTensors() { return Map.copyOf(largeConstants); }
+    /**
+     * Returns an immutable map of the large constants of this, represented as strings on the standard tensor form.
+     * These can have sizes in gigabytes and must be distributed to nodes separately from configuration.
+     * @deprecated Use largeConstantTensors instead
+     */
+    @Override
+    @SuppressWarnings("removal")
+    public Map<String, String> largeConstants() { return asStrings(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 84e227e29ad..8c8fc5c4b11 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
@@ -23,8 +23,12 @@ public interface ImportedMlModel {
     ModelType modelType();
 
     Optional<String> inputTypeSpec(String input);
-    Map<String, Tensor> smallConstants();
-    Map<String, Tensor> largeConstants();
+    @Deprecated(forRemoval = true)
+    Map<String, String> smallConstants();
+    @Deprecated(forRemoval = true)
+    Map<String, String> largeConstants();
+    Map<String, Tensor> smallConstantTensors();
+    Map<String, Tensor> largeConstantTensors();
     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 c1c798117b3..7fb167ee6f1 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
@@ -21,15 +21,15 @@ public class OnnxMnistSoftmaxImportTestCase {
         ImportedModel model = new OnnxImporter().importModel("test", "src/test/models/onnx/mnist_softmax/mnist_softmax.onnx").asNative();
 
         // Check constants
-        assertEquals(2, model.largeConstants().size());
+        assertEquals(2, model.largeConstantTensors().size());
 
-        Tensor constant0 = model.largeConstants().get("test_Variable");
+        Tensor constant0 = model.largeConstantTensors().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 = model.largeConstants().get("test_Variable_1");
+        Tensor constant1 = model.largeConstantTensors().get("test_Variable_1");
         assertNotNull(constant1);
         assertEquals(new TensorType.Builder(TensorType.Value.FLOAT).indexed("d1", 10).build(), constant1.type());
         assertEquals(10, constant1.size());
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 2682ee40d76..f78150d9875 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
@@ -63,8 +63,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)));
-        result.smallConstants().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(tensor)));
+        result.largeConstantTensors().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(tensor)));
+        result.smallConstantTensors().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(tensor)));
         return context;
     }
 
diff --git a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/vespa/VespaImportTestCase.java b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/vespa/VespaImportTestCase.java
index 60e0df8ba90..d9c7e67c946 100644
--- a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/vespa/VespaImportTestCase.java
+++ b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/vespa/VespaImportTestCase.java
@@ -38,12 +38,12 @@ public class VespaImportTestCase {
         assertEquals("tensor(name{},x[3])", model.inputs().get("input1").toString());
         assertEquals("tensor(x[3])", model.inputs().get("input2").toString());
 
-        assertEquals(2, model.smallConstants().size());
-        assertEquals("tensor(x[3]):[0.5, 1.5, 2.5]", model.smallConstants().get("constant1").toString());
-        assertEquals("tensor():{3.0}", model.smallConstants().get("constant2").toString());
+        assertEquals(2, model.smallConstantTensors().size());
+        assertEquals("tensor(x[3]):[0.5, 1.5, 2.5]", model.smallConstantTensors().get("constant1").toString());
+        assertEquals("tensor():{3.0}", model.smallConstantTensors().get("constant2").toString());
 
-        assertEquals(1, model.largeConstants().size());
-        assertEquals("tensor(x[3]):[0.5, 1.5, 2.5]", model.largeConstants().get("constant1asLarge").toString());
+        assertEquals(1, model.largeConstantTensors().size());
+        assertEquals("tensor(x[3]):[0.5, 1.5, 2.5]", model.largeConstantTensors().get("constant1asLarge").toString());
 
         assertEquals(2, model.expressions().size());
         assertEquals("reduce(reduce(input1 * input2, sum, name) * constant1, max, x) * constant2",
@@ -71,8 +71,8 @@ public class VespaImportTestCase {
         assertTrue(model.expressions().isEmpty());
         assertTrue(model.functions().isEmpty());
         assertTrue(model.inputs().isEmpty());
-        assertTrue(model.largeConstants().isEmpty());
-        assertTrue(model.smallConstants().isEmpty());
+        assertTrue(model.largeConstantTensors().isEmpty());
+        assertTrue(model.smallConstantTensors().isEmpty());
     }
 
     @Test