Avoid converting tensor to strings and back again.

2 files changed, 6 insertions, 4 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 4e7710aa449..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); }
 
@@ -95,7 +95,7 @@ public class ImportedModel implements ImportedMlModel {
      * For TensorFlow this corresponds to Variable files stored separately.
      */
     @Override
-    public Map<String, String> largeConstants() { return asStrings(largeConstants); }
+    public Map<String, Tensor> largeConstants() { return Map.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 a2626818f87..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
@@ -1,6 +1,8 @@
 // Copyright Yahoo. 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;
@@ -21,8 +23,8 @@ public interface ImportedMlModel {
     ModelType modelType();
 
     Optional<String> inputTypeSpec(String input);
-    Map<String, String> smallConstants();
-    Map<String, String> largeConstants();
+    Map<String, Tensor> smallConstants();
+    Map<String, Tensor> largeConstants();
     Map<String, String> functions();
     List<ImportedMlFunction> outputExpressions();