[8.16] backporting fix for negative scores in text_similarity_ranker retriever (#121056)

Author: pmpailis

docs/changelog/120930.yaml

@@ -0,0 +1,6 @@
+pr: 120930
+summary: Normalize negative scores for `text_similarity_reranker` retriever
+area: Ranking
+type: bug
+issues:
+ - 120201

docs/reference/search/retriever.asciidoc

@@ -367,6 +367,23 @@ Currently you can:
 ** Then set up an <<inference-example-eland,{es} service inference endpoint>> with the `rerank` task type
 ** Refer to the <<text-similarity-reranker-retriever-example-eland,example>> on this page for a step-by-step guide.
 
+[IMPORTANT]
+====
+Scores from the re-ranking process are normalized using the following formula before returned to the user,
+to avoid having negative scores.
+[source,text]
+----
+score = max(score, 0) + min(exp(score), 1)
+----
+Using the above, any initially negative scores are projected to (0, 1) and positive scores to [1, infinity).
+To revert back if needed, one can use:
+[source, text]
+----
+score = score - 1, if score >= 0
+score = ln(score), if score < 0
+----
+====
+
 ===== Parameters
 `retriever`::
 (Required, <<retriever, retriever>>)

server/src/main/java/org/elasticsearch/search/retriever/rankdoc/RankDocsQuery.java

@@ -58,6 +58,11 @@ public static class TopQuery extends Query {
             this.queryNames = queryNames;
             this.segmentStarts = segmentStarts;
             this.contextIdentity = contextIdentity;
+            for (RankDoc doc : docs) {
+                if (false == doc.score >= 0) {
+                    throw new IllegalArgumentException("RankDoc scores must be positive values. Missing a normalization step?");
+                }
+            }
         }
 
         @Override
@@ -161,7 +166,11 @@ public float getMaxScore(int docId) {
 
                         @Override
                         public float score() {
-                            return docs[upTo].score;
+                            // We could still end up with a valid 0 score for a RankDoc
+                            // so here we want to differentiate between this and all the tailQuery matches
+                            // that would also produce a 0 score due to filtering, by setting the score to `Float.MIN_VALUE` instead for
+                            // RankDoc matches.
+                            return Math.max(docs[upTo].score, Float.MIN_VALUE);
                         }
 
                         @Override

server/src/test/java/org/elasticsearch/index/query/RankDocsQueryBuilderTests.java

@@ -251,4 +251,16 @@ public void testUnknownField() throws IOException {
     public void testValidOutput() throws IOException {
         // no-op since RankDocsQueryBuilder is an internal only API
     }
+
+    public void shouldThrowForNegativeScores() throws IOException {
+        try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
+            iw.addDocument(new Document());
+            try (IndexReader reader = iw.getReader()) {
+                SearchExecutionContext context = createSearchExecutionContext(newSearcher(reader));
+                RankDocsQueryBuilder queryBuilder = new RankDocsQueryBuilder(new RankDoc[] { new RankDoc(0, -1.0f, 0) }, null, false);
+                IllegalArgumentException ex = expectThrows(IllegalArgumentException.class, () -> queryBuilder.doToQuery(context));
+                assertEquals("RankDoc scores must be positive values. Missing a normalization step?", ex.getMessage());
+            }
+        }
+    }
 }

test/framework/src/main/java/org/elasticsearch/search/rank/rerank/AbstractRerankerIT.java

@@ -56,6 +56,10 @@ public enum ThrowingRankBuilderType {
 
     protected abstract Collection<Class<? extends Plugin>> pluginsNeeded();
 
+    protected boolean shouldCheckScores() {
+        return true;
+    }
+
     @Override
     protected Collection<Class<? extends Plugin>> nodePlugins() {
         return pluginsNeeded();
@@ -95,9 +99,11 @@ public void testRerankerNoExceptions() throws Exception {
                 int rank = 1;
                 for (SearchHit searchHit : response.getHits().getHits()) {
                     assertThat(searchHit, hasId(String.valueOf(5 - (rank - 1))));
-                    assertEquals(0.5f - ((rank - 1) * 0.1f), searchHit.getScore(), 1e-5f);
                     assertThat(searchHit, hasRank(rank));
                     assertNotNull(searchHit.getFields().get(searchField));
+                    if (shouldCheckScores()) {
+                        assertEquals(0.5f - ((rank - 1) * 0.1f), searchHit.getScore(), 1e-5f);
+                    }
                     rank++;
                 }
             }
@@ -140,9 +146,11 @@ public void testRerankerPagination() throws Exception {
                 int rank = 3;
                 for (SearchHit searchHit : response.getHits().getHits()) {
                     assertThat(searchHit, hasId(String.valueOf(5 - (rank - 1))));
-                    assertEquals(0.5f - ((rank - 1) * 0.1f), searchHit.getScore(), 1e-5f);
                     assertThat(searchHit, hasRank(rank));
                     assertNotNull(searchHit.getFields().get(searchField));
+                    if (shouldCheckScores()) {
+                        assertEquals(0.5f - ((rank - 1) * 0.1f), searchHit.getScore(), 1e-5f);
+                    }
                     rank++;
                 }
             }
@@ -222,9 +230,11 @@ public void testNotAllShardsArePresentInFetchPhase() throws Exception {
                 int rank = 1;
                 for (SearchHit searchHit : response.getHits().getHits()) {
                     assertThat(searchHit, hasId(String.valueOf(5 - (rank - 1))));
-                    assertEquals(0.5f - ((rank - 1) * 0.1f), searchHit.getScore(), 1e-5f);
                     assertThat(searchHit, hasRank(rank));
                     assertNotNull(searchHit.getFields().get(searchField));
+                    if (shouldCheckScores()) {
+                        assertEquals(0.5f - ((rank - 1) * 0.1f), searchHit.getScore(), 1e-5f);
+                    }
                     rank++;
                 }
             }

x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/AbstractTestInferenceService.java

@@ -25,9 +25,16 @@
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Random;
 
 public abstract class AbstractTestInferenceService implements InferenceService {
 
+    protected static final Random random = new Random(
+        System.getProperty("tests.seed") == null
+            ? System.currentTimeMillis()
+            : Long.parseUnsignedLong(System.getProperty("tests.seed").split(":")[0], 16)
+    );
+
     protected static int stringWeight(String input, int position) {
         int hashCode = input.hashCode();
         if (hashCode < 0) {

x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestRerankingServiceExtension.java

@@ -36,6 +36,7 @@
 import java.util.Map;
 
 public class TestRerankingServiceExtension implements InferenceServiceExtension {
+
     @Override
     public List<Factory> getInferenceServiceFactories() {
         return List.of(TestInferenceService::new);
@@ -122,9 +123,12 @@ public void chunkedInfer(
         private RankedDocsResults makeResults(List<String> input) {
             List<RankedDocsResults.RankedDoc> results = new ArrayList<>();
             int totalResults = input.size();
+            float minScore = random.nextFloat(-1f, 1f);
             float resultDiff = 0.2f;
             for (int i = 0; i < input.size(); i++) {
-                results.add(new RankedDocsResults.RankedDoc(totalResults - 1 - i, resultDiff * (totalResults - i), input.get(i)));
+                results.add(
+                    new RankedDocsResults.RankedDoc(totalResults - 1 - i, minScore + resultDiff * (totalResults - i), input.get(i))
+                );
             }
             return new RankedDocsResults(results);
         }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankFeaturePhaseRankCoordinatorContext.java

@@ -20,8 +20,8 @@
 import org.elasticsearch.xpack.inference.services.cohere.rerank.CohereRerankTaskSettings;
 import org.elasticsearch.xpack.inference.services.googlevertexai.rerank.GoogleVertexAiRerankTaskSettings;
 
+import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
 
@@ -130,10 +130,15 @@ protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[
      */
     @Override
     protected RankFeatureDoc[] preprocess(RankFeatureDoc[] originalDocs) {
-        return Arrays.stream(originalDocs)
-            .filter(doc -> minScore == null || doc.score >= minScore)
-            .sorted(Comparator.comparing((RankFeatureDoc doc) -> doc.score).reversed())
-            .toArray(RankFeatureDoc[]::new);
+        List<RankFeatureDoc> docs = new ArrayList<>();
+        for (RankFeatureDoc doc : originalDocs) {
+            if (minScore == null || doc.score >= minScore) {
+                doc.score = normalizeScore(doc.score);
+                docs.add(doc);
+            }
+        }
+        docs.sort(RankFeatureDoc::compareTo);
+        return docs.toArray(new RankFeatureDoc[0]);
     }
 
     protected InferenceAction.Request generateRequest(List<String> docFeatures) {
@@ -154,7 +159,15 @@ private float[] extractScoresFromRankedDocs(List<RankedDocsResults.RankedDoc> ra
         for (RankedDocsResults.RankedDoc rankedDoc : rankedDocs) {
             scores[rankedDoc.index()] = rankedDoc.relevanceScore();
         }
-
         return scores;
     }
+
+    private static float normalizeScore(float score) {
+        // As some models might produce negative scores, we want to ensure that all scores will be positive
+        // so we will make use of the following normalization formula:
+        // score = max(score, 0) + min(exp(score), 1)
+        // this will ensure that all positive scores lie in the [1, inf) range,
+        // while negative values (and 0) will be shifted to (0, 1]
+        return Math.max(score, 0) + Math.min((float) Math.exp(score), 1);
+    }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankRetrieverBuilder.java

@@ -156,6 +156,7 @@ protected RankDoc[] combineInnerRetrieverResults(List<ScoreDoc[]> rankResults) {
         TextSimilarityRankDoc[] textSimilarityRankDocs = new TextSimilarityRankDoc[scoreDocs.length];
         for (int i = 0; i < scoreDocs.length; i++) {
             ScoreDoc scoreDoc = scoreDocs[i];
+            assert scoreDoc.score >= 0;
             textSimilarityRankDocs[i] = new TextSimilarityRankDoc(scoreDoc.doc, scoreDoc.score, scoreDoc.shardIndex, inferenceId, field);
         }
         return textSimilarityRankDocs;

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankMultiNodeTests.java

@@ -50,4 +50,9 @@ public void testQueryPhaseShardThrowingAllShardsFail() throws Exception {
     public void testQueryPhaseCoordinatorThrowingAllShardsFail() throws Exception {
         // no-op
     }
+
+    @Override
+    protected boolean shouldCheckScores() {
+        return false;
+    }
 }