Add Bounded Window to Inference Models for Rescoring to Ensure Positive Score Range (#125694) (#127345)

* apply bounded window inference model

* linting

* add unit tests

* [CI] Auto commit changes from spotless

* add additional tests

* remove unused constructor

---------

Co-authored-by: elasticsearchmachine <[email protected]>
(cherry picked from commit e77bf80)

Author: markjhoy

docs/changelog/125694.yaml

@@ -0,0 +1,5 @@
+pr: 125694
+summary: LTR score bounding
+area: Ranking
+type: bug
+issues: []

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/inference/BoundedInferenceModel.java

@@ -0,0 +1,14 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.core.ml.inference.trainedmodel.inference;
+
+public interface BoundedInferenceModel extends InferenceModel {
+    double getMinPredictedValue();
+
+    double getMaxPredictedValue();
+}

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/inference/BoundedWindowInferenceModel.java

@@ -0,0 +1,123 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.core.ml.inference.trainedmodel.inference;
+
+import org.elasticsearch.common.logging.LoggerMessageFormat;
+import org.elasticsearch.inference.InferenceResults;
+import org.elasticsearch.xpack.core.ml.inference.results.RegressionInferenceResults;
+import org.elasticsearch.xpack.core.ml.inference.trainedmodel.InferenceConfig;
+import org.elasticsearch.xpack.core.ml.inference.trainedmodel.TargetType;
+
+import java.util.Map;
+
+public class BoundedWindowInferenceModel implements BoundedInferenceModel {
+    public static final double DEFAULT_MIN_PREDICTED_VALUE = 0;
+
+    private final BoundedInferenceModel model;
+    private final double minPredictedValue;
+    private final double maxPredictedValue;
+    private final double adjustmentValue;
+
+    public BoundedWindowInferenceModel(BoundedInferenceModel model) {
+        this.model = model;
+        this.minPredictedValue = model.getMinPredictedValue();
+        this.maxPredictedValue = model.getMaxPredictedValue();
+
+        if (this.minPredictedValue < DEFAULT_MIN_PREDICTED_VALUE) {
+            this.adjustmentValue = DEFAULT_MIN_PREDICTED_VALUE - this.minPredictedValue;
+        } else {
+            this.adjustmentValue = 0.0;
+        }
+    }
+
+    @Override
+    public String[] getFeatureNames() {
+        return model.getFeatureNames();
+    }
+
+    @Override
+    public TargetType targetType() {
+        return model.targetType();
+    }
+
+    @Override
+    public InferenceResults infer(Map<String, Object> fields, InferenceConfig config, Map<String, String> featureDecoderMap) {
+        return boundInferenceResultScores(model.infer(fields, config, featureDecoderMap));
+    }
+
+    @Override
+    public InferenceResults infer(double[] features, InferenceConfig config) {
+        return boundInferenceResultScores(model.infer(features, config));
+    }
+
+    @Override
+    public boolean supportsFeatureImportance() {
+        return model.supportsFeatureImportance();
+    }
+
+    @Override
+    public String getName() {
+        return "bounded_window[" + model.getName() + "]";
+    }
+
+    @Override
+    public void rewriteFeatureIndices(Map<String, Integer> newFeatureIndexMapping) {
+        model.rewriteFeatureIndices(newFeatureIndexMapping);
+    }
+
+    @Override
+    public long ramBytesUsed() {
+        return model.ramBytesUsed();
+    }
+
+    @Override
+    public double getMinPredictedValue() {
+        return minPredictedValue;
+    }
+
+    @Override
+    public double getMaxPredictedValue() {
+        return maxPredictedValue;
+    }
+
+    private InferenceResults boundInferenceResultScores(InferenceResults inferenceResult) {
+        // if the min value < the default minimum, slide the values up by the adjustment value
+        if (inferenceResult instanceof RegressionInferenceResults regressionInferenceResults) {
+            double predictedValue = ((Number) regressionInferenceResults.predictedValue()).doubleValue();
+
+            predictedValue += this.adjustmentValue;
+
+            return new RegressionInferenceResults(
+                predictedValue,
+                inferenceResult.getResultsField(),
+                ((RegressionInferenceResults) inferenceResult).getFeatureImportance()
+            );
+        }
+
+        throw new IllegalStateException(
+            LoggerMessageFormat.format(
+                "Model used within a {} should return a {} but got {} instead",
+                BoundedWindowInferenceModel.class.getSimpleName(),
+                RegressionInferenceResults.class.getSimpleName(),
+                inferenceResult.getClass().getSimpleName()
+            )
+        );
+    }
+
+    @Override
+    public String toString() {
+        return "BoundedWindowInferenceModel{"
+            + "model="
+            + model
+            + ", minPredictedValue="
+            + getMinPredictedValue()
+            + ", maxPredictedValue="
+            + getMaxPredictedValue()
+            + '}';
+    }
+}

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/inference/EnsembleInferenceModel.java

@@ -11,6 +11,7 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.util.CachedSupplier;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Tuple;
 import org.elasticsearch.inference.InferenceResults;
@@ -36,6 +37,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
@@ -52,7 +54,7 @@
 import static org.elasticsearch.xpack.core.ml.inference.trainedmodel.ensemble.Ensemble.CLASSIFICATION_WEIGHTS;
 import static org.elasticsearch.xpack.core.ml.inference.trainedmodel.ensemble.Ensemble.TRAINED_MODELS;
 
-public class EnsembleInferenceModel implements InferenceModel {
+public class EnsembleInferenceModel implements InferenceModel, BoundedInferenceModel {
 
     public static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(EnsembleInferenceModel.class);
     private static final Logger LOGGER = LogManager.getLogger(EnsembleInferenceModel.class);
@@ -97,6 +99,7 @@ public static EnsembleInferenceModel fromXContent(XContentParser parser) {
     private final List<String> classificationLabels;
     private final double[] classificationWeights;
     private volatile boolean preparedForInference = false;
+    private final Supplier<double[]> predictedValuesBoundariesSupplier;
 
     private EnsembleInferenceModel(
         List<InferenceModel> models,
@@ -112,6 +115,7 @@ private EnsembleInferenceModel(
         this.classificationWeights = classificationWeights == null
             ? null
             : classificationWeights.stream().mapToDouble(Double::doubleValue).toArray();
+        this.predictedValuesBoundariesSupplier = CachedSupplier.wrap(this::initModelBoundaries);
     }
 
     @Override
@@ -328,21 +332,57 @@ public double[] getClassificationWeights() {
 
     @Override
     public String toString() {
-        return "EnsembleInferenceModel{"
-            + "featureNames="
-            + Arrays.toString(featureNames)
-            + ", models="
-            + models
-            + ", outputAggregator="
-            + outputAggregator
-            + ", targetType="
-            + targetType
-            + ", classificationLabels="
-            + classificationLabels
-            + ", classificationWeights="
-            + Arrays.toString(classificationWeights)
-            + ", preparedForInference="
-            + preparedForInference
-            + '}';
+        StringBuilder builder = new StringBuilder("EnsembleInferenceModel{");
+
+        builder.append("featureNames=")
+            .append(Arrays.toString(featureNames))
+            .append(", models=")
+            .append(models)
+            .append(", outputAggregator=")
+            .append(outputAggregator)
+            .append(", targetType=")
+            .append(targetType);
+
+        if (targetType == TargetType.CLASSIFICATION) {
+            builder.append(", classificationLabels=")
+                .append(classificationLabels)
+                .append(", classificationWeights=")
+                .append(Arrays.toString(classificationWeights));
+        } else if (targetType == TargetType.REGRESSION) {
+            builder.append(", minPredictedValue=")
+                .append(getMinPredictedValue())
+                .append(", maxPredictedValue=")
+                .append(getMaxPredictedValue());
+        }
+
+        builder.append(", preparedForInference=").append(preparedForInference);
+
+        return builder.append('}').toString();
+    }
+
+    @Override
+    public double getMinPredictedValue() {
+        return this.predictedValuesBoundariesSupplier.get()[0];
+    }
+
+    @Override
+    public double getMaxPredictedValue() {
+        return this.predictedValuesBoundariesSupplier.get()[1];
+    }
+
+    private double[] initModelBoundaries() {
+        double[] modelsMinBoundaries = new double[models.size()];
+        double[] modelsMaxBoundaries = new double[models.size()];
+        int i = 0;
+        for (InferenceModel model : models) {
+            if (model instanceof BoundedInferenceModel boundedInferenceModel) {
+                modelsMinBoundaries[i] = boundedInferenceModel.getMinPredictedValue();
+                modelsMaxBoundaries[i++] = boundedInferenceModel.getMaxPredictedValue();
+            } else {
+                throw new IllegalStateException("All submodels have to be bounded");
+            }
+        }
+
+        return new double[] { outputAggregator.aggregate(modelsMinBoundaries), outputAggregator.aggregate(modelsMaxBoundaries) };
     }
 }

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/inference/InferenceDefinition.java

@@ -14,6 +14,7 @@
 import org.elasticsearch.xpack.core.ml.inference.preprocessing.LenientlyParsedPreProcessor;
 import org.elasticsearch.xpack.core.ml.inference.preprocessing.PreProcessor;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.InferenceConfig;
+import org.elasticsearch.xpack.core.ml.inference.trainedmodel.LearningToRankConfig;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.TargetType;
 import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;
 
@@ -79,13 +80,21 @@ private void preProcess(Map<String, Object> fields) {
 
     public InferenceResults infer(Map<String, Object> fields, InferenceConfig config) {
         preProcess(fields);
+
+        InferenceModel inferenceModel = trainedModel;
+
+        if (config instanceof LearningToRankConfig) {
+            assert trainedModel instanceof BoundedInferenceModel;
+            inferenceModel = new BoundedWindowInferenceModel((BoundedInferenceModel) trainedModel);
+        }
+
         if (config.requestingImportance() && trainedModel.supportsFeatureImportance() == false) {
             throw ExceptionsHelper.badRequestException(
                 "Feature importance is not supported for the configured model of type [{}]",
                 trainedModel.getName()
             );
         }
-        return trainedModel.infer(fields, config, config.requestingImportance() ? getDecoderMap() : Collections.emptyMap());
+        return inferenceModel.infer(fields, config, config.requestingImportance() ? getDecoderMap() : Collections.emptyMap());
     }
 
     public TargetType getTargetType() {