Add Hugging Face Rerank support #127966

Evgenii-Kazannik · 2025-05-09T11:37:59Z

Have you signed the contributor license agreement?
Have you followed the contributor guidelines?
If submitting code, have you built your formula locally prior to submission with gradle check?
If submitting code, is your pull request against main? Unless there is a good reason otherwise, we prefer pull requests against main and will backport as needed.
If submitting code, have you checked that your submission is for an OS and architecture that we support?
If you are submitting this code for a class then read our policy for that.

CA have been signed.
Used the following with success:
gradlew :x-pack:plugin:inference:check
gradlew.bat :x-pack:plugin:inference:spotlessApply

Tested via api:

PUT {{base-url}}/_inference/rerank/bge-reranker-base-mkn
{
"service": "hugging_face",
"service_settings": {
"api_key": "{{hf-api-key}}",
"url": "{{hf-bge-reranker-url}}"
},
"task_settings": {
"return_text": false,
"top_n": 4
}
}

POST {{base-url}}/_inference/rerank/bge-reranker-base-mkn
{
"input": ["luke", "like", "leia", "chewy","r2d2", "star", "wars"],
"query": "star wars main character",
"top_n": 2,
"return_documents": false
}

{
"rerank": [
{
"index": 6,
"relevance_score": 0.50955844
},
{
"index": 5,
"relevance_score": 0.084341794
}
]
}

POST {{base-url}}/_inference/rerank/bge-reranker-base-mkn
{
"input": ["luke", "like", "leia", "chewy","r2d2", "star", "wars"],
"query": "star wars main character",
"top_n": 3,
"return_documents": true
}

{
"rerank": [
{
"index": 6,
"relevance_score": 0.5089636,
"text": "wars"
},
{
"index": 5,
"relevance_score": 0.08449275,
"text": "star"
},
{
"index": 3,
"relevance_score": 0.0045032725,
"text": "chewy"
}
]
}

Also there were the following HF task settings integrated additionally:
raw_scores, truncate, truncation_direction
For now removed from the PR saving into a distinct branch,
Just for a case if we decide to make those a part of the inference api

@jonathan-buttner @Jan-Kazlouski-elastic
Apologies for the delay. I meant to create this much sooner.
Thanks for your patience

Jan-Kazlouski-elastic · 2025-05-09T12:24:52Z

...ice-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceBaseRestTest.java

@@ -171,6 +171,22 @@ static String mockDenseServiceModelConfig() {
            """;
    }

+    static String mockRerankServiceModelConfig() {


I'm wondering if methods you've added to this class are actually used somewhere. Methods you've taken for reference are being called. The ones you've added - are not.

Thanks for noticing. It's used now

Jan-Kazlouski-elastic · 2025-05-09T12:25:48Z

...ice-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceBaseRestTest.java

@@ -484,6 +500,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 -> {


It looks like this method is not called with TaskType.RERANK param anywhere. meaning assertion isn't triggered.

Jan-Kazlouski-elastic · 2025-05-09T12:52:02Z

...main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceBaseService.java

@@ -92,14 +98,15 @@ public HuggingFaceModel parsePersistedConfigWithSecrets(
        Map<String, Object> secrets
    ) {
        Map<String, Object> serviceSettingsMap = removeFromMapOrThrowIfNull(config, ModelConfigurations.SERVICE_SETTINGS);
+        Map<String, Object> taskSettingsMap = removeFromMapOrThrowIfNull(config, ModelConfigurations.TASK_SETTINGS);


Correct me if I'm wrong. but won't that throw an exception if there are no task settings in config? If so, doesn't that affect other integrations that don't require TASK_SETTINGS to be present?

I added Rerank type check to ensure the method isn't used for other tasks

Jan-Kazlouski-elastic · 2025-05-09T12:59:03Z

...main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceBaseService.java

    }

    @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 = removeFromMapOrThrowIfNull(config, ModelConfigurations.TASK_SETTINGS);


Same question as above

Added type check before using the methos.Thanks

Jan-Kazlouski-elastic

Left a few comments.

Jan-Kazlouski-elastic · 2025-05-09T13:22:18Z

...sticsearch/xpack/inference/services/huggingface/request/rerank/HuggingFaceRerankRequest.java

+
+    @Override
+    public boolean[] getTruncationInfo() {
+        return null;


Can we have a comment here, explaining why null is returned?

Yeah truncation is only used in some services that support text embedding. Just say something like "Not applicable for rerank, only used in text embedding requests".

Added as suggested: Not applicable for rerank, only used in text embedding requests
Thanks all

Jan-Kazlouski-elastic · 2025-05-09T13:33:58Z

...sticsearch/xpack/inference/services/huggingface/rerank/HuggingFaceRerankServiceSettings.java

+
+    @Override
+    public TransportVersion getMinimalSupportedVersion() {
+        return TransportVersions.V_8_12_0;


Please check comments related to TransportVersions left by @jonathan-buttner to this PR: #127254
They would apply here as well.

Did it, read it through, updated. I'm going to update the versions once more before the merge

Jan-Kazlouski-elastic · 2025-05-09T13:34:41Z

...elasticsearch/xpack/inference/services/huggingface/rerank/HuggingFaceRerankTaskSettings.java

+
+    @Override
+    public TransportVersion getMinimalSupportedVersion() {
+        return TransportVersions.V_8_14_0;


Same thing here related to comments for TranportVersion

Thanks, applied the change

jonathan-buttner

Looking good! I left a few suggestions.

jonathan-buttner · 2025-05-09T19:32:40Z

...sticsearch/xpack/inference/services/huggingface/request/rerank/HuggingFaceRerankRequest.java

+        this.returnDocuments = returnDocuments;
+        this.topN = topN;
+        taskSettings = model.getTaskSettings();
+        this.model = model;


Since we're saving a reference to the model how about we remove the taskSettings and inferenceEntityId references and just use the model.

Thank you. Used them from the model

jonathan-buttner · 2025-05-09T19:34:11Z

.../main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceModelInput.java

+
+import java.util.Map;
+
+public class HuggingFaceModelInput {


How about we make this a record and maybe rename it to HuggingFaceModelParameters

Yep. The record fits better. Thanks. Done

jonathan-buttner · 2025-05-09T19:39:43Z

.../main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceModelInput.java

+    private final String failureMessage;
+    private final ConfigurationParseContext context;
+
+    public HuggingFaceModelInput(Builder builder) {


Should we make this private? We probably want the instantiation done through the builder.

The builder was replaced with the record as suggested so not needed anymore.
Though thank you for pointing that out

jonathan-buttner · 2025-05-09T19:42:45Z

...main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceBaseService.java

@@ -128,17 +140,13 @@ public HuggingFaceModel parsePersistedConfig(String inferenceEntityId, TaskType
            parsePersistedConfigErrorMsg(inferenceEntityId, name()),
            ConfigurationParseContext.PERSISTENT
        );
+
+        return createModel(
+            TaskType.RERANK.equals(taskType) ? modelBuilder.withTaskSettings(taskSettingsMap).build() : modelBuilder.build()


Looks like the builder accepts null task settings so how about we just pass in the task settings map, regardless of it being null or not. That way we don't need to check for rerank here.

The models accept the task settings map as is now. Thanks

jonathan-buttner · 2025-05-09T19:50:21Z

...sticsearch/xpack/inference/services/huggingface/request/rerank/HuggingFaceRerankRequest.java

+
+    @Override
+    public boolean[] getTruncationInfo() {
+        return null;


Yeah truncation is only used in some services that support text embedding. Just say something like "Not applicable for rerank, only used in text embedding requests".

jonathan-buttner · 2025-05-09T20:30:52Z

...sticsearch/xpack/inference/services/huggingface/rerank/HuggingFaceRerankServiceSettings.java

+        return RERANK_TOKEN_LIMIT;
+    }
+
+    // model is not defined in the service settings.


Since we encountered situations where the model id was required for chat completion, have we done any testing to see if the serverless style endpoint requires the model id?

The thing is that HF currently does not provide serverless for Rerank models. We cannot test it now

jonathan-buttner · 2025-05-09T20:31:24Z

...sticsearch/xpack/inference/services/huggingface/rerank/HuggingFaceRerankServiceSettings.java

+    @Override
+    protected XContentBuilder toXContentFragmentOfExposedFields(XContentBuilder builder, Params params) throws IOException {
+        builder.field(URL, uri.toString());
+        builder.field(MAX_INPUT_TOKENS, RERANK_TOKEN_LIMIT);


Let's remove this, since we don't use it.

Removed. Thank you

jonathan-buttner · 2025-05-09T20:36:05Z

...ticsearch/xpack/inference/services/huggingface/response/HuggingFaceRerankResponseEntity.java

@@ -0,0 +1,123 @@
+/*


We're trying to move away from this style of parsing and instead use an ObjectParser or ConstructingObjectParser. How about we switch this implementation to use ConstructingObjectParser? Here's an example: https://github.com/elastic/elasticsearch/blob/main/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/response/OpenAiEmbeddingsResponseEntity.java

jonathan-buttner · 2025-05-09T20:36:50Z

...ticsearch/xpack/inference/services/huggingface/response/HuggingFaceRerankResponseEntity.java

+import static org.elasticsearch.common.xcontent.XContentParserUtils.throwUnknownToken;
+import static org.elasticsearch.xpack.inference.external.response.XContentUtils.moveToFirstToken;
+
+public class HuggingFaceRerankResponseEntity extends ErrorResponse {


Hmm, typically we separate the valid response from the error response. Does the HuggingFaceErrorResponseEntity suffice?

jonathan-buttner · 2025-05-09T20:39:06Z

.../org/elasticsearch/xpack/inference/services/huggingface/action/HuggingFaceActionCreator.java

+        "Failed to send Hugging Face %s request from inference entity id [%s]";
+    static final ResponseHandler RERANK_HANDLER = new HuggingFaceResponseHandler(
+        "hugging face rerank",
+        (request, response) -> HuggingFaceRerankResponseEntity.fromResponse((HuggingFaceRerankRequest) request, response)


It'd be unlikely but can we do an instanceof check for request being a HuggingFaceRerankRequest? And throw an IllegalArgumentException if it's invalid.

Good point. Thank you Jonathan. Explicit check was added

# Conflicts: # server/src/main/java/org/elasticsearch/TransportVersions.java

elasticsearchmachine · 2025-05-13T11:47:49Z

Pinging @elastic/ml-core (Team:ML)

Jan-Kazlouski-elastic · 2025-05-14T18:31:12Z

.../org/elasticsearch/xpack/inference/services/huggingface/action/HuggingFaceActionCreator.java

+    private static final String FAILED_TO_SEND_REQUEST_ERROR_MESSAGE =
+        "Failed to send Hugging Face %s request from inference entity id [%s]";
+    private static final String INVALID_REQUEST_TYPE_MESSAGE = "Invalid request type: expected HuggingFace %s request but got %s";
+    static final ResponseHandler RERANK_HANDLER = new HuggingFaceResponseHandler("hugging face rerank", (request, response) -> {


Why is it package private and not just private?

Add Hugging Face Rerank support

4d9ec59

elasticsearchmachine added needs:triage Requires assignment of a team area label v9.1.0 external-contributor Pull request authored by a developer outside the Elasticsearch team labels May 9, 2025

Jan-Kazlouski-elastic reviewed May 9, 2025

View reviewed changes

jonathan-buttner requested changes May 9, 2025

View reviewed changes

Evgenii-Kazannik added 2 commits May 13, 2025 13:31

Address comments

b58aab4

Merge branch 'main' into Add-Hugging-Face-Rerank-support

c567137

# Conflicts: # server/src/main/java/org/elasticsearch/TransportVersions.java

PeteGillinElastic added :ml Machine learning and removed needs:triage Requires assignment of a team area label labels May 13, 2025

elasticsearchmachine added the Team:ML Meta label for the ML team label May 13, 2025

Evgenii-Kazannik added 5 commits May 13, 2025 14:36

Add transport version

a4ebb87

Merge branch 'main' into Add-Hugging-Face-Rerank-support

5d316c1

Add transport version

f74e9e5

Merge branch 'main' into Add-Hugging-Face-Rerank-support

2ea07f0

Add to inference service and crud IT rerank tests

0054891

Jan-Kazlouski-elastic reviewed May 14, 2025

View reviewed changes

Add Hugging Face Rerank support #127966

Are you sure you want to change the base?

Add Hugging Face Rerank support #127966

Conversation

Evgenii-Kazannik commented May 9, 2025 • edited Loading

Choose a reason for hiding this comment

Evgenii-Kazannik May 13, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Jan-Kazlouski-elastic left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jonathan-buttner left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Evgenii-Kazannik May 13, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

elasticsearchmachine commented May 13, 2025

Choose a reason for hiding this comment

Evgenii-Kazannik commented May 9, 2025 •

edited

Loading

Evgenii-Kazannik May 13, 2025 •

edited

Loading

Evgenii-Kazannik May 13, 2025 •

edited

Loading