Add Hugging Face Rerank support

docs/changelog/127966.yaml

@@ -0,0 +1,5 @@
+pr: 127966
+summary: "[ML] Add Rerank support to the Inference Plugin"
+area: Machine Learning
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -179,6 +179,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion V_8_19_FIELD_CAPS_ADD_CLUSTER_ALIAS = def(8_841_0_32);
     public static final TransportVersion ESQL_HASH_OPERATOR_STATUS_OUTPUT_TIME_8_19 = def(8_841_0_34);
     public static final TransportVersion RERANKER_FAILURES_ALLOWED_8_19 = def(8_841_0_35);
+    public static final TransportVersion ML_INFERENCE_HUGGING_FACE_RERANK_ADDED_8_19 = def(8_841_0_36);
     public static final TransportVersion V_9_0_0 = def(9_000_0_09);
     public static final TransportVersion INITIAL_ELASTICSEARCH_9_0_1 = def(9_000_0_10);
     public static final TransportVersion INITIAL_ELASTICSEARCH_9_0_2 = def(9_000_0_11);
@@ -261,7 +262,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ESQL_HASH_OPERATOR_STATUS_OUTPUT_TIME = def(9_077_0_00);
     public static final TransportVersion ML_INFERENCE_HUGGING_FACE_CHAT_COMPLETION_ADDED = def(9_078_0_00);
     public static final TransportVersion NODES_STATS_SUPPORTS_MULTI_PROJECT = def(9_079_0_00);
-
+    public static final TransportVersion ML_INFERENCE_HUGGING_FACE_RERANK_ADDED = def(9_080_0_00);
     /*
      * STOP! READ THIS FIRST! No, really,
      *        ____ _____ ___  ____  _        ____  _____    _    ____    _____ _   _ ___ ____    _____ ___ ____  ____ _____ _

server/src/test/java/org/elasticsearch/inference/SettingsConfigurationTestUtils.java

@@ -20,9 +20,8 @@
 public class SettingsConfigurationTestUtils {
 
     public static SettingsConfiguration getRandomSettingsConfigurationField() {
-        return new SettingsConfiguration.Builder(EnumSet.of(TaskType.TEXT_EMBEDDING, TaskType.SPARSE_EMBEDDING)).setDefaultValue(
-            randomAlphaOfLength(10)
-        )
+        return new SettingsConfiguration.Builder(EnumSet.of(TaskType.TEXT_EMBEDDING, TaskType.SPARSE_EMBEDDING, TaskType.RERANK))
+            .setDefaultValue(randomAlphaOfLength(10))
             .setDescription(randomAlphaOfLength(10))
             .setLabel(randomAlphaOfLength(10))
             .setRequired(randomBoolean())

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/inference/results/RankedDocsResultsTests.java

@@ -82,4 +82,13 @@ private List<RankedDocsResults.RankedDoc> rankedDocsNullStringToEmpty(List<Ranke
     protected RankedDocsResults doParseInstance(XContentParser parser) throws IOException {
         return RankedDocsResults.createParser(true).apply(parser, null);
     }
+
+    public record RerankExpectation(Map<String, Object> rankedDocFields) {}
+
+    public static Map<String, Object> buildExpectationRerank(List<RerankExpectation> rerank) {
+        return Map.of(
+            RankedDocsResults.RERANK,
+            rerank.stream().map(rerankExpectation -> Map.of(RankedDocsResults.RankedDoc.NAME, rerankExpectation.rankedDocFields)).toList()
+        );
+    }
 }

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceBaseRestTest.java

@@ -171,6 +171,20 @@ static String mockDenseServiceModelConfig() {
             """;
     }
 
+    static String mockRerankServiceModelConfig() {
+        return """
+            {
+              "service": "test_reranking_service",
+              "service_settings": {
+                 "model_id": "my_model",
+                 "api_key": "abc64"
+              },
+              "task_settings": {
+              }
+            }
+            """;
+    }
+
     static void deleteModel(String modelId) throws IOException {
         var request = new Request("DELETE", "_inference/" + modelId);
         var response = client().performRequest(request);
@@ -484,6 +498,10 @@ private String jsonBody(List<String> input, @Nullable String query) {
     @SuppressWarnings("unchecked")
     protected void assertNonEmptyInferenceResults(Map<String, Object> resultMap, int expectedNumberOfResults, TaskType taskType) {
         switch (taskType) {
+            case RERANK -> {
+                var results = (List<Map<String, Object>>) resultMap.get(TaskType.RERANK.toString());
+                assertThat(results, hasSize(expectedNumberOfResults));
+            }
             case SPARSE_EMBEDDING -> {
                 var results = (List<Map<String, Object>>) resultMap.get(TaskType.SPARSE_EMBEDDING.toString());
                 assertThat(results, hasSize(expectedNumberOfResults));

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceCrudIT.java

@@ -53,9 +53,12 @@ public void testCRUD() throws IOException {
         for (int i = 0; i < 4; i++) {
             putModel("te_model_" + i, mockDenseServiceModelConfig(), TaskType.TEXT_EMBEDDING);
         }
+        for (int i = 0; i < 3; i++) {
+            putModel("re-model-" + i, mockRerankServiceModelConfig(), TaskType.RERANK);
+        }
 
         var getAllModels = getAllModels();
-        int numModels = 12;
+        int numModels = 15;
         assertThat(getAllModels, hasSize(numModels));
 
         var getSparseModels = getModels("_all", TaskType.SPARSE_EMBEDDING);
@@ -71,6 +74,13 @@ public void testCRUD() throws IOException {
         for (var denseModel : getDenseModels) {
             assertEquals("text_embedding", denseModel.get("task_type"));
         }
+
+        var getRerankModels = getModels("_all", TaskType.RERANK);
+        int numRerankModels = 4;
+        assertThat(getRerankModels, hasSize(numRerankModels));
+        for (var denseModel : getRerankModels) {
+            assertEquals("rerank", denseModel.get("task_type"));
+        }
         String oldApiKey;
         {
             var singleModel = getModels("se_model_1", TaskType.SPARSE_EMBEDDING);
@@ -100,6 +110,9 @@ public void testCRUD() throws IOException {
         for (int i = 0; i < 4; i++) {
             deleteModel("te_model_" + i, TaskType.TEXT_EMBEDDING);
         }
+        for (int i = 0; i < 3; i++) {
+            deleteModel("re-model-" + i, TaskType.RERANK);
+        }
     }
 
     public void testGetModelWithWrongTaskType() throws IOException {

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetServicesIT.java

@@ -101,7 +101,7 @@ public void testGetServicesWithTextEmbeddingTaskType() throws IOException {
 
     public void testGetServicesWithRerankTaskType() throws IOException {
         List<Object> services = getServices(TaskType.RERANK);
-        assertThat(services.size(), equalTo(7));
+        assertThat(services.size(), equalTo(8));
 
         var providers = providers(services);
 
@@ -115,7 +115,8 @@ public void testGetServicesWithRerankTaskType() throws IOException {
                     "googlevertexai",
                     "jinaai",
                     "test_reranking_service",
-                    "voyageai"
+                    "voyageai",
+                    "hugging_face"
                 ).toArray()
             )
         );

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/MockRerankInferenceServiceIT.java

@@ -0,0 +1,67 @@
+/*
+ * 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.inference;
+
+import org.elasticsearch.inference.TaskType;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+public class MockRerankInferenceServiceIT extends InferenceBaseRestTest {
+
+    @SuppressWarnings("unchecked")
+    public void testMockService() throws IOException {
+        String inferenceEntityId = "test-mock";
+        var putModel = putModel(inferenceEntityId, mockRerankServiceModelConfig(), TaskType.RERANK);
+        var model = getModels(inferenceEntityId, TaskType.RERANK).get(0);
+
+        for (var modelMap : List.of(putModel, model)) {
+            assertEquals(inferenceEntityId, modelMap.get("inference_id"));
+            assertEquals(TaskType.RERANK, TaskType.fromString((String) modelMap.get("task_type")));
+            assertEquals("test_reranking_service", modelMap.get("service"));
+        }
+
+        List<String> input = List.of(randomAlphaOfLength(10));
+        var inference = infer(inferenceEntityId, input);
+        assertNonEmptyInferenceResults(inference, 1, TaskType.RERANK);
+        assertEquals(inference, infer(inferenceEntityId, input));
+        assertNotEquals(inference, infer(inferenceEntityId, randomValueOtherThan(input, () -> List.of(randomAlphaOfLength(10)))));
+    }
+
+    public void testMockServiceWithMultipleInputs() throws IOException {
+        String inferenceEntityId = "test-mock-with-multi-inputs";
+        putModel(inferenceEntityId, mockRerankServiceModelConfig(), TaskType.RERANK);
+        var queryParams = Map.of("timeout", "120s");
+
+        var inference = infer(
+            inferenceEntityId,
+            TaskType.RERANK,
+            List.of(randomAlphaOfLength(5), randomAlphaOfLength(10)),
+            "What if?",
+            queryParams
+        );
+
+        assertNonEmptyInferenceResults(inference, 2, TaskType.RERANK);
+    }
+
+    @SuppressWarnings("unchecked")
+    public void testMockService_DoesNotReturnSecretsInGetResponse() throws IOException {
+        String inferenceEntityId = "test-mock";
+        var putModel = putModel(inferenceEntityId, mockRerankServiceModelConfig(), TaskType.RERANK);
+        var model = getModels(inferenceEntityId, TaskType.RERANK).get(0);
+
+        var serviceSettings = (Map<String, Object>) model.get("service_settings");
+        assertNull(serviceSettings.get("api_key"));
+        assertNotNull(serviceSettings.get("model_id"));
+
+        var putServiceSettings = (Map<String, Object>) putModel.get("service_settings");
+        assertNull(putServiceSettings.get("api_key"));
+        assertNotNull(putServiceSettings.get("model_id"));
+    }
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceNamedWriteablesProvider.java

@@ -80,6 +80,8 @@
 import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceServiceSettings;
 import org.elasticsearch.xpack.inference.services.huggingface.completion.HuggingFaceChatCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.huggingface.elser.HuggingFaceElserServiceSettings;
+import org.elasticsearch.xpack.inference.services.huggingface.rerank.HuggingFaceRerankServiceSettings;
+import org.elasticsearch.xpack.inference.services.huggingface.rerank.HuggingFaceRerankTaskSettings;
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.embeddings.IbmWatsonxEmbeddingsServiceSettings;
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.rerank.IbmWatsonxRerankServiceSettings;
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.rerank.IbmWatsonxRerankTaskSettings;
@@ -365,6 +367,16 @@ private static void addHuggingFaceNamedWriteables(List<NamedWriteableRegistry.En
                 HuggingFaceChatCompletionServiceSettings::new
             )
         );
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(TaskSettings.class, HuggingFaceRerankTaskSettings.NAME, HuggingFaceRerankTaskSettings::new)
+        );
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                ServiceSettings.class,
+                HuggingFaceRerankServiceSettings.NAME,
+                HuggingFaceRerankServiceSettings::new
+            )
+        );
     }
 
     private static void addGoogleAiStudioNamedWritables(List<NamedWriteableRegistry.Entry> namedWriteables) {

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

@@ -19,6 +19,7 @@
 import org.elasticsearch.xpack.core.inference.results.RankedDocsResults;
 import org.elasticsearch.xpack.inference.services.cohere.rerank.CohereRerankTaskSettings;
 import org.elasticsearch.xpack.inference.services.googlevertexai.rerank.GoogleVertexAiRerankTaskSettings;
+import org.elasticsearch.xpack.inference.services.huggingface.rerank.HuggingFaceRerankTaskSettings;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -91,7 +92,10 @@ protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[
             } else if (r.getEndpoints().isEmpty() == false
                 && r.getEndpoints().get(0).getTaskSettings() instanceof GoogleVertexAiRerankTaskSettings googleVertexAiTaskSettings) {
                     configuredTopN = googleVertexAiTaskSettings.topN();
-                }
+                } else if (r.getEndpoints().isEmpty() == false
+                    && r.getEndpoints().get(0).getTaskSettings() instanceof HuggingFaceRerankTaskSettings huggingFaceRerankTaskSettings) {
+                        configuredTopN = huggingFaceRerankTaskSettings.getTopNDocumentsOnly();
+                    }
             if (configuredTopN != null && configuredTopN < rankWindowSize) {
                 l.onFailure(
                     new IllegalArgumentException(

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceBaseService.java

@@ -57,6 +57,7 @@ public void parseRequestConfig(
     ) {
         try {
             Map<String, Object> serviceSettingsMap = removeFromMapOrThrowIfNull(config, ModelConfigurations.SERVICE_SETTINGS);
+            Map<String, Object> taskSettingsMap = removeFromMapOrDefaultEmpty(config, ModelConfigurations.TASK_SETTINGS);
 
             ChunkingSettings chunkingSettings = null;
             if (TaskType.TEXT_EMBEDDING.equals(taskType)) {
@@ -66,17 +67,21 @@ public void parseRequestConfig(
             }
 
             var model = createModel(
-                inferenceEntityId,
-                taskType,
-                serviceSettingsMap,
-                chunkingSettings,
-                serviceSettingsMap,
-                TaskType.unsupportedTaskTypeErrorMsg(taskType, name()),
-                ConfigurationParseContext.REQUEST
+                new HuggingFaceModelParameters(
+                    inferenceEntityId,
+                    taskType,
+                    serviceSettingsMap,
+                    taskSettingsMap,
+                    chunkingSettings,
+                    serviceSettingsMap,
+                    TaskType.unsupportedTaskTypeErrorMsg(taskType, name()),
+                    ConfigurationParseContext.REQUEST
+                )
             );
 
             throwIfNotEmptyMap(config, name());
             throwIfNotEmptyMap(serviceSettingsMap, name());
+            throwIfNotEmptyMap(taskSettingsMap, name());
 
             parsedModelListener.onResponse(model);
         } catch (Exception e) {
@@ -92,6 +97,7 @@ public HuggingFaceModel parsePersistedConfigWithSecrets(
         Map<String, Object> secrets
     ) {
         Map<String, Object> serviceSettingsMap = removeFromMapOrThrowIfNull(config, ModelConfigurations.SERVICE_SETTINGS);
+        Map<String, Object> taskSettingsMap = removeFromMapOrDefaultEmpty(config, ModelConfigurations.TASK_SETTINGS);
         Map<String, Object> secretSettingsMap = removeFromMapOrThrowIfNull(secrets, ModelSecrets.SECRET_SETTINGS);
 
         ChunkingSettings chunkingSettings = null;
@@ -100,45 +106,44 @@ public HuggingFaceModel parsePersistedConfigWithSecrets(
         }
 
         return createModel(
-            inferenceEntityId,
-            taskType,
-            serviceSettingsMap,
-            chunkingSettings,
-            secretSettingsMap,
-            parsePersistedConfigErrorMsg(inferenceEntityId, name()),
-            ConfigurationParseContext.PERSISTENT
+            new HuggingFaceModelParameters(
+                inferenceEntityId,
+                taskType,
+                serviceSettingsMap,
+                taskSettingsMap,
+                chunkingSettings,
+                secretSettingsMap,
+                parsePersistedConfigErrorMsg(inferenceEntityId, name()),
+                ConfigurationParseContext.PERSISTENT
+            )
         );
     }
 
     @Override
     public HuggingFaceModel parsePersistedConfig(String inferenceEntityId, TaskType taskType, Map<String, Object> config) {
         Map<String, Object> serviceSettingsMap = removeFromMapOrThrowIfNull(config, ModelConfigurations.SERVICE_SETTINGS);
+        Map<String, Object> taskSettingsMap = removeFromMapOrDefaultEmpty(config, ModelConfigurations.TASK_SETTINGS);
 
         ChunkingSettings chunkingSettings = null;
         if (TaskType.TEXT_EMBEDDING.equals(taskType)) {
             chunkingSettings = ChunkingSettingsBuilder.fromMap(removeFromMap(config, ModelConfigurations.CHUNKING_SETTINGS));
         }
 
         return createModel(
-            inferenceEntityId,
-            taskType,
-            serviceSettingsMap,
-            chunkingSettings,
-            null,
-            parsePersistedConfigErrorMsg(inferenceEntityId, name()),
-            ConfigurationParseContext.PERSISTENT
+            new HuggingFaceModelParameters(
+                inferenceEntityId,
+                taskType,
+                serviceSettingsMap,
+                taskSettingsMap,
+                chunkingSettings,
+                null,
+                parsePersistedConfigErrorMsg(inferenceEntityId, name()),
+                ConfigurationParseContext.PERSISTENT
+            )
         );
     }
 
-    protected abstract HuggingFaceModel createModel(
-        String inferenceEntityId,
-        TaskType taskType,
-        Map<String, Object> serviceSettings,
-        ChunkingSettings chunkingSettings,
-        Map<String, Object> secretSettings,
-        String failureMessage,
-        ConfigurationParseContext context
-    );
+    protected abstract HuggingFaceModel createModel(HuggingFaceModelParameters input);
 
     @Override
     public void doInfer(

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceModel.java

@@ -11,6 +11,7 @@
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.inference.ModelConfigurations;
 import org.elasticsearch.inference.ModelSecrets;
+import org.elasticsearch.inference.TaskSettings;
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.services.RateLimitGroupingModel;
 import org.elasticsearch.xpack.inference.services.ServiceUtils;
@@ -35,6 +36,13 @@ public HuggingFaceModel(
         apiKey = ServiceUtils.apiKey(apiKeySecrets);
     }
 
+    protected HuggingFaceModel(HuggingFaceModel model, TaskSettings taskSettings) {
+        super(model, taskSettings);
+
+        rateLimitServiceSettings = model.rateLimitServiceSettings();
+        apiKey = model.apiKey();
+    }
+
     public HuggingFaceRateLimitServiceSettings rateLimitServiceSettings() {
         return rateLimitServiceSettings;
     }