Avoid converting tensor to strings and back again.

6 files changed, 15 insertions, 16 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 9ba87fd24bf..2344d3d15e0 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
@@ -294,8 +294,7 @@ public class ConvertedModel {
     }
 
     private static void transformSmallConstant(ModelStore store, RankProfile profile, String constantName,
-                                               String constantValueString) {
-        Tensor constantValue = Tensor.from(constantValueString);
+                                               Tensor constantValue) {
         store.writeSmallConstant(constantName, constantValue);
         Reference name = FeatureNames.asConstantFeature(constantName);
         profile.add(new RankProfile.Constant(name, constantValue));
@@ -306,8 +305,7 @@ public class ConvertedModel {
                                                QueryProfileRegistry queryProfiles,
                                                Set<String> constantsReplacedByFunctions,
                                                String constantName,
-                                               String constantValueString) {
-        Tensor constantValue = Tensor.from(constantValueString);
+                                               Tensor constantValue) {
         RankProfile.RankingExpressionFunction rankingExpressionFunctionOverridingConstant = profile.getFunctions().get(constantName);
         if (rankingExpressionFunctionOverridingConstant != null) {
             TensorType functionType = rankingExpressionFunctionOverridingConstant.function().getBody().type(profile.typeContext(queryProfiles));
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();
 
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 b6b63912c52..c1c798117b3 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
@@ -23,13 +23,13 @@ public class OnnxMnistSoftmaxImportTestCase {
         // Check constants
         assertEquals(2, model.largeConstants().size());
 
-        Tensor constant0 = Tensor.from(model.largeConstants().get("test_Variable"));
+        Tensor constant0 = model.largeConstants().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.largeConstants().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 5d13697df06..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
@@ -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.from(tensor))));
-        result.smallConstants().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(Tensor.from(tensor))));
+        result.largeConstants().forEach((name, tensor) -> context.put("constant(" + name + ")", new TensorValue(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 25c51a75b0b..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
@@ -4,7 +4,6 @@ package ai.vespa.rankingexpression.importer.vespa;
 import ai.vespa.rankingexpression.importer.ImportedModel;
 import ai.vespa.rankingexpression.importer.configmodelview.ImportedMlFunction;
 import com.yahoo.searchlib.rankingexpression.RankingExpression;
-import com.yahoo.searchlib.rankingexpression.evaluation.Context;
 import com.yahoo.searchlib.rankingexpression.evaluation.MapContext;
 import com.yahoo.searchlib.rankingexpression.evaluation.TensorValue;
 import com.yahoo.searchlib.rankingexpression.parser.ParseException;
@@ -40,11 +39,11 @@ 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"));
+        assertEquals("tensor(x[3]):[0.5, 1.5, 2.5]", model.largeConstants().get("constant1asLarge").toString());
 
         assertEquals(2, model.expressions().size());
         assertEquals("reduce(reduce(input1 * input2, sum, name) * constant1, max, x) * constant2",