Add option to use default signature for onnx models

5 files changed, 27 insertions, 9 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 55979023119..5ccdcfba215 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
@@ -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 com.yahoo.vespa.model.ml;
 
+import ai.vespa.rankingexpression.importer.IntermediateGraph;
 import ai.vespa.rankingexpression.importer.configmodelview.ImportedMlFunction;
 import com.google.common.collect.ImmutableMap;
 import com.yahoo.collections.Pair;
@@ -160,13 +161,16 @@ public class ConvertedModel {
         ExpressionFunction expression = expressions.get(arguments.toName());
         if (expression != null) return expression;
 
-        if ( ! arguments.signature().isPresent()) {
+        expression = expressions.get(IntermediateGraph.defaultSignature() + "." + arguments.toName());
+        if (expression != null) return expression;
+
+        if (arguments.signature().isEmpty()) {
             if (expressions.size() > 1)
                 throw new IllegalArgumentException("Multiple candidate expressions " + missingExpressionMessageSuffix());
             return expressions.values().iterator().next();
         }
 
-        if ( ! arguments.output().isPresent()) {
+        if (arguments.output().isEmpty()) {
             List<Map.Entry<String, ExpressionFunction>> entriesWithTheRightPrefix =
                     expressions.entrySet().stream().filter(entry -> entry.getKey().startsWith(arguments.signature().get() + ".")).collect(Collectors.toList());
             if (entriesWithTheRightPrefix.size() < 1)
diff --git a/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithOnnxTestCase.java b/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithOnnxTestCase.java
index 754f161c70b..4387d19c474 100644
--- a/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithOnnxTestCase.java
+++ b/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithOnnxTestCase.java
@@ -128,6 +128,20 @@ public class RankingExpressionWithOnnxTestCase {
     }
 
     @Test
+    public void testOnnxReferenceWithSpecifiedOutput() {
+        RankProfileSearchFixture search = fixtureWith("tensor<float>(d0[2],d1[784])(0.0)",
+                "onnx('mnist_softmax.onnx', 'add')");
+        search.assertFirstPhaseExpression(vespaExpression, "my_profile");
+    }
+
+    @Test
+    public void testOnnxReferenceWithSpecifiedOutputAndSignature() {
+        RankProfileSearchFixture search = fixtureWith("tensor<float>(d0[2],d1[784])(0.0)",
+                "onnx('mnist_softmax.onnx', 'default.add')");
+        search.assertFirstPhaseExpression(vespaExpression, "my_profile");
+    }
+
+    @Test
     public void testOnnxReferenceMissingFunction() throws ParseException {
         try {
             RankProfileSearchFixture search = new RankProfileSearchFixture(
@@ -171,7 +185,7 @@ public class RankingExpressionWithOnnxTestCase {
     @Test
     public void testOnnxReferenceSpecifyingNonExistingOutput() {
         try {
-            RankProfileSearchFixture search = fixtureWith("tensor(d0[2],d1[784])(0.0)",
+            RankProfileSearchFixture search = fixtureWith("tensor<float>(d0[2],d1[784])(0.0)",
                     "onnx('mnist_softmax.onnx', 'y')");
             search.assertFirstPhaseExpression(vespaExpression, "my_profile");
             fail("Expecting exception");
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/IntermediateGraph.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/IntermediateGraph.java
index 14aa3ebf84e..16c1b913188 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/IntermediateGraph.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/IntermediateGraph.java
@@ -58,7 +58,7 @@ public class IntermediateGraph {
         return signatures.computeIfAbsent(signature, (k) -> new GraphSignature()).outputs;
     }
 
-    public String defaultSignature() {
+    public static String defaultSignature() {
         return "default";
     }
 
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/onnx/GraphImporter.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/onnx/GraphImporter.java
index d42338deaf8..75ed19c010c 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/onnx/GraphImporter.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/onnx/GraphImporter.java
@@ -151,7 +151,7 @@ class GraphImporter {
             OrderedTensorType type = TypeConverter.typeFrom(valueInfoProto.getType());
             operation = new Argument(intermediateGraph.name(), valueInfoProto.getName(), type);
 
-            intermediateGraph.inputs(intermediateGraph.defaultSignature())
+            intermediateGraph.inputs(IntermediateGraph.defaultSignature())
                     .put(IntermediateOperation.namePartOf(name), operation.vespaName());
 
         } else if (isConstantTensor(name, onnxGraph)) {
@@ -171,7 +171,7 @@ class GraphImporter {
             }
 
             if (isOutputNode(name, onnxGraph)) {
-                intermediateGraph.outputs(intermediateGraph.defaultSignature())
+                intermediateGraph.outputs(IntermediateGraph.defaultSignature())
                         .put(IntermediateOperation.namePartOf(name), operation.name());
             }
         }
@@ -238,7 +238,7 @@ class GraphImporter {
     }
 
     private static void verifyNoWarnings(IntermediateGraph intermediateGraph) {
-        for (java.util.Map.Entry<String, String> output : intermediateGraph.outputs(intermediateGraph.defaultSignature()).entrySet()) {
+        for (java.util.Map.Entry<String, String> output : intermediateGraph.outputs(IntermediateGraph.defaultSignature()).entrySet()) {
             IntermediateOperation operation = intermediateGraph.get(output.getValue());
             Set<String> warnings = getWarnings(operation);
             if (warnings.size() > 0) {
@@ -248,7 +248,7 @@ class GraphImporter {
     }
 
     private static void verifyOutputTypes(Onnx.GraphProto onnxGraph, IntermediateGraph intermediateGraph) {
-        for (java.util.Map.Entry<String, String> output : intermediateGraph.outputs(intermediateGraph.defaultSignature()).entrySet()) {
+        for (java.util.Map.Entry<String, String> output : intermediateGraph.outputs(IntermediateGraph.defaultSignature()).entrySet()) {
             IntermediateOperation operation = intermediateGraph.get(output.getValue());
             Onnx.ValueInfoProto onnxNode = getOutputNode(output.getKey(), onnxGraph);
             OrderedTensorType type = operation.type().orElseThrow(
diff --git a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/onnx/OnnxOperationsTestCase.java b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/onnx/OnnxOperationsTestCase.java
index 6954abe5157..9927f4df04f 100644
--- a/model-integration/src/test/java/ai/vespa/rankingexpression/importer/onnx/OnnxOperationsTestCase.java
+++ b/model-integration/src/test/java/ai/vespa/rankingexpression/importer/onnx/OnnxOperationsTestCase.java
@@ -391,7 +391,7 @@ public class OnnxOperationsTestCase {
     private TensorFunction optimizeAndRename(String opName, IntermediateOperation op) {
         IntermediateGraph graph = new IntermediateGraph(modelName);
         graph.put(opName, op);
-        graph.outputs(graph.defaultSignature()).put(opName, opName);
+        graph.outputs(IntermediateGraph.defaultSignature()).put(opName, opName);
         graph.optimize();
         return op.function().get();
     }