Handle small and large Constants in a uniform way.

4 files changed, 5 insertions, 5 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 d84bccc6424..edccc5544af 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
@@ -85,7 +85,7 @@ public class ImportedModel implements ImportedMlModel {
      * 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, String> smallConstants() { return asStrings(smallConstants); }
+    public Map<String, Tensor> smallConstants() { return Map.copyOf(smallConstants); }
 
     boolean hasSmallConstant(String name) { return smallConstants.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 8f1b93fdec2..84e227e29ad 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,7 +23,7 @@ public interface ImportedMlModel {
     ModelType modelType();
 
     Optional<String> inputTypeSpec(String input);
-    Map<String, String> smallConstants();
+    Map<String, Tensor> smallConstants();
     Map<String, Tensor> largeConstants();
     Map<String, String> functions();
     List<ImportedMlFunction> outputExpressions();
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 2e79e227d88..2682ee40d76 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
@@ -64,7 +64,7 @@ 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.from(tensor))));
+        result.smallConstants().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 c165e1afa8c..60e0df8ba90 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
@@ -39,8 +39,8 @@ public class VespaImportTestCase {
         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"));
-        assertEquals("tensor():{3.0}", model.smallConstants().get("constant2"));
+        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(1, model.largeConstants().size());
         assertEquals("tensor(x[3]):[0.5, 1.5, 2.5]", model.largeConstants().get("constant1asLarge").toString());