add special handling of TensorFunctionNode containing Generate function

4 files changed, 159 insertions, 0 deletions

diff --git a/config-model/src/test/derived/globalphase_token_functions/files/m.py b/config-model/src/test/derived/globalphase_token_functions/files/m.py
new file mode 100644
index 00000000000..004135b32eb
--- /dev/null
+++ b/config-model/src/test/derived/globalphase_token_functions/files/m.py
@@ -0,0 +1,51 @@
+# imports
+
+from onnx import TensorProto
+from onnx.helper import (
+    make_model, make_node, make_graph,
+    make_tensor_value_info, make_value_info)
+from onnx.checker import check_model
+
+# inputs
+
+# TensorProto.DOUBLE is the element type, [128] the shape
+A = make_tensor_value_info('input_ids',      TensorProto.DOUBLE, [128])
+B = make_tensor_value_info('attention_mask', TensorProto.DOUBLE, [128])
+C = make_tensor_value_info('token_type_ids', TensorProto.DOUBLE, [128])
+
+# outputs, the shape is defined
+Y = make_tensor_value_info('vector_Y', TensorProto.DOUBLE, [128])
+S = make_tensor_value_info('score', TensorProto.DOUBLE, [1])
+
+# Creates node defined by the operator type, inputs, outputs, and possibly options
+node1 = make_node('Mul', ['input_ids', 'attention_mask'], ['masked'])
+node2 = make_node('Add', ['masked', 'token_type_ids'], ['vector_Y'])
+node3 = make_node('ReduceSum', inputs=['vector_Y'], outputs=['score'], keepdims=1)
+
+# from nodes to graph
+# the graph is built from the list of nodes, the list of inputs,
+# the list of outputs and a name.
+
+graph = make_graph([node1, node2, node3],  # nodes
+                    'ranking_model',  # a name
+                    [A, B, C],  # inputs
+                    [S])  # outputs
+
+# onnx graph to model
+onnx_model = make_model(graph)
+
+# ensure we do not get too new opset version:
+del onnx_model.opset_import[:]
+opset = onnx_model.opset_import.add()
+opset.version = 17
+
+# Let's check the model is consistent, this function is described in
+# section Checker and Shape Inference.
+check_model(onnx_model)
+
+# The serialization
+with open("ranking_model.onnx", "wb") as f:
+    f.write(onnx_model.SerializeToString())
+
+# the work is done, let's display it...
+print(onnx_model)
diff --git a/config-model/src/test/derived/globalphase_token_functions/files/ranking_model.onnx b/config-model/src/test/derived/globalphase_token_functions/files/ranking_model.onnx
new file mode 100644
index 00000000000..274fbba3fdb
--- /dev/null
+++ b/config-model/src/test/derived/globalphase_token_functions/files/ranking_model.onnx
@@ -0,0 +1,23 @@
+:þ

+(
+	input_ids
+attention_maskmasked"Mul
+'
+masked
+token_type_idsvector_Y"Add
+-
+vector_Yscore"	ReduceSum*
+keepdims
 

ranking_modelZ
+	input_ids
+	
+€
Z
+attention_mask
+	
+€
Z
+token_type_ids
+	
+€
b
+score
+
+
+
B
\ No newline at end of file
diff --git a/config-model/src/test/derived/globalphase_token_functions/rank-profiles.cfg b/config-model/src/test/derived/globalphase_token_functions/rank-profiles.cfg
new file mode 100644
index 00000000000..33381060178
--- /dev/null
+++ b/config-model/src/test/derived/globalphase_token_functions/rank-profiles.cfg
@@ -0,0 +1,43 @@
+rankprofile[].name "default"
+rankprofile[].fef.property[].name "vespa.type.attribute.tokens"
+rankprofile[].fef.property[].value "tensor(d0[128])"
+rankprofile[].name "unranked"
+rankprofile[].fef.property[].name "vespa.rank.firstphase"
+rankprofile[].fef.property[].value "value(0)"
+rankprofile[].fef.property[].name "vespa.hitcollector.heapsize"
+rankprofile[].fef.property[].value "0"
+rankprofile[].fef.property[].name "vespa.hitcollector.arraysize"
+rankprofile[].fef.property[].value "0"
+rankprofile[].fef.property[].name "vespa.dump.ignoredefaultfeatures"
+rankprofile[].fef.property[].value "true"
+rankprofile[].fef.property[].name "vespa.type.attribute.tokens"
+rankprofile[].fef.property[].value "tensor(d0[128])"
+rankprofile[].name "using_model"
+rankprofile[].fef.property[].name "rankingExpression(__token_length@1019197748).rankingScript"
+rankprofile[].fef.property[].value "reduce(map(query(input), f(x)(x > 0)), sum)"
+rankprofile[].fef.property[].name "rankingExpression(__token_length@-812590320).rankingScript"
+rankprofile[].fef.property[].value "reduce(map(attribute(tokens), f(x)(x > 0)), sum)"
+rankprofile[].fef.property[].name "rankingExpression(input_ids).rankingScript"
+rankprofile[].fef.property[].value "tensor<float>(d0[1],d1[128])((if (d1 < 1.0, 101.0, if (d1 < 1.0 + rankingExpression(__token_length@1019197748), query(input){d0:(d1 - (1.0))}, if (d1 < (1.0 + rankingExpression(__token_length@1019197748) + 1.0), 102.0, if (d1 < (1.0 + rankingExpression(__token_length@1019197748) + 1.0 + rankingExpression(__token_length@-812590320)), attribute(tokens){d0:(d1 - (1.0 + rankingExpression(__token_length@1019197748) + 1.0))}, if (d1 < (1.0 + rankingExpression(__token_length@1019197748) + 1.0 + rankingExpression(__token_length@-812590320) + 1.0), 102.0, 0.0)))))))"
+rankprofile[].fef.property[].name "rankingExpression(input_ids).type"
+rankprofile[].fef.property[].value "tensor<float>(d0[1],d1[128])"
+rankprofile[].fef.property[].name "rankingExpression(token_type_ids).rankingScript"
+rankprofile[].fef.property[].value "tensor<float>(d0[1],d1[128])((if (d1 < (1.0 + rankingExpression(__token_length@1019197748) + 1.0), 0.0, if (d1 < (1.0 + rankingExpression(__token_length@1019197748) + 1.0 + rankingExpression(__token_length@-812590320) + 1.0), 1.0, 0.0))))"
+rankprofile[].fef.property[].name "rankingExpression(token_type_ids).type"
+rankprofile[].fef.property[].value "tensor<float>(d0[1],d1[128])"
+rankprofile[].fef.property[].name "rankingExpression(attention_mask).rankingScript"
+rankprofile[].fef.property[].value "tensor<float>(d0[1],d1[128])((if (d1 < (1.0 + rankingExpression(__token_length@1019197748) + 1.0 + rankingExpression(__token_length@-812590320) + 1.0), 1.0, 0.0)))"
+rankprofile[].fef.property[].name "rankingExpression(attention_mask).type"
+rankprofile[].fef.property[].value "tensor<float>(d0[1],d1[128])"
+rankprofile[].fef.property[].name "vespa.rank.globalphase"
+rankprofile[].fef.property[].value "rankingExpression(globalphase)"
+rankprofile[].fef.property[].name "rankingExpression(globalphase).rankingScript"
+rankprofile[].fef.property[].value "onnx(my_ranking_model).score{d0:0}"
+rankprofile[].fef.property[].name "vespa.match.feature"
+rankprofile[].fef.property[].value "attribute(tokens)"
+rankprofile[].fef.property[].name "vespa.globalphase.rerankcount"
+rankprofile[].fef.property[].value "1000"
+rankprofile[].fef.property[].name "vespa.type.attribute.tokens"
+rankprofile[].fef.property[].value "tensor(d0[128])"
+rankprofile[].fef.property[].name "vespa.type.query.input"
+rankprofile[].fef.property[].value "tensor(d0[32])"
diff --git a/config-model/src/test/derived/globalphase_token_functions/test.sd b/config-model/src/test/derived/globalphase_token_functions/test.sd
new file mode 100644
index 00000000000..f8cc8863ad1
--- /dev/null
+++ b/config-model/src/test/derived/globalphase_token_functions/test.sd
@@ -0,0 +1,42 @@
+# Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+schema test {
+    document test {
+        field title type string {
+            indexing: index | summary
+        }
+        field tokens type tensor(d0[128]) {
+            indexing: attribute
+        }
+    }
+    fieldset default {
+        fields: title
+    }
+
+    onnx-model my_ranking_model {
+        file: files/ranking_model.onnx
+        input input_ids: input_ids
+        input attention_mask: attention_mask
+        input token_type_ids: token_type_ids
+    }
+
+    rank-profile using_model {
+        inputs {
+            query(input) tensor(d0[32])
+        }
+        function input_ids() {
+            expression: tokenInputIds(128, query(input), attribute(tokens))
+        }
+        function token_type_ids() {
+            expression: tokenTypeIds(128, query(input), attribute(tokens))
+        }
+        function attention_mask() {
+            expression: tokenAttentionMask(128, query(input), attribute(tokens))
+        }
+        global-phase {
+            rerank-count: 1000
+            expression: onnx(my_ranking_model).score{d0:0}
+        }
+    }
+
+}