Update model-integration for supporting BERT-type models

1 files changed, 7 insertions, 1 deletions

diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Softmax.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Softmax.java
index 83086926316..d03827f4c72 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Softmax.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Softmax.java
@@ -28,6 +28,10 @@ public class Softmax extends IntermediateOperation {
     @Override
     protected OrderedTensorType lazyGetType() {
         if ( ! allInputTypesPresent(1)) return null;
+
+        // input is referenced twice due to overflow avoidance, so make this it's own function.
+        inputs.get(0).exportAsRankingFunction = true;
+
         return inputs.get(0).type().get();
     }
 
@@ -50,7 +54,9 @@ public class Softmax extends IntermediateOperation {
         }
 
         TensorFunction input = inputs.get(0).function().get();
-        TensorFunction exp = new Map(input, ScalarFunctions.exp());
+        TensorFunction max = new Reduce(input, Reduce.Aggregator.max, reduceDimensions);
+        TensorFunction cap = new Join(input, max, ScalarFunctions.subtract());  // to avoid overflow
+        TensorFunction exp = new Map(cap, ScalarFunctions.exp());
         TensorFunction sum = new Reduce(exp, Reduce.Aggregator.sum, reduceDimensions);
         TensorFunction div = new Join(exp, sum, ScalarFunctions.divide());