Update inference specification for Hugging Face's rerank task (#4417) (#4563)

(cherry picked from commit eba5dca)

Author: Evgenii-Kazannik

output/openapi/elasticsearch-openapi.json

@@ -979,9 +979,22 @@ export class HuggingFaceServiceSettings {
   model_id?: string
 }
 
+export class HuggingFaceTaskSettings {
+  /**
+   * For a `rerank` task, return doc text within the results.
+   */
+  return_documents?: boolean
+  /**
+   * For a `rerank` task, the number of most relevant documents to return.
+   * It defaults to the number of the documents.
+   */
+  top_n?: integer
+}
+
 export enum HuggingFaceTaskType {
   chat_completion,
   completion,
+  rerank,
   text_embedding
 }
 

output/openapi/elasticsearch-serverless-openapi.json

@@ -40,7 +40,7 @@ import { TaskType } from '@inference/_types/TaskType'
  * * ELSER (`sparse_embedding`)
  * * Google AI Studio (`completion`, `text_embedding`)
  * * Google Vertex AI (`rerank`, `text_embedding`)
- * * Hugging Face (`chat_completion`, `completion`, `text_embedding`)
+ * * Hugging Face (`chat_completion`, `completion`, `rerank`, `text_embedding`)
  * * Mistral (`chat_completion`, `completion`, `text_embedding`)
  * * OpenAI (`chat_completion`, `completion`, `text_embedding`)
  * * VoyageAI (`text_embedding`, `rerank`)

output/schema/schema.json

@@ -22,6 +22,7 @@ import { Id } from '@_types/common'
 import {
   HuggingFaceServiceSettings,
   HuggingFaceServiceType,
+  HuggingFaceTaskSettings,
   HuggingFaceTaskType
 } from '@inference/_types/CommonTypes'
 import { InferenceChunkingSettings } from '@inference/_types/Services'
@@ -56,6 +57,16 @@ import { InferenceChunkingSettings } from '@inference/_types/Services'
  * * `Mistral-7B-Instruct-v0.2`
  * * `QwQ-32B`
  * * `Phi-3-mini-128k-instruct`
+ *
+ * For Elastic's `rerank` task:
+ * The selected model must support the `sentence-ranking` task and expose OpenAI API.
+ * HuggingFace supports only dedicated (not serverless) endpoints for `Rerank` so far.
+ * After the endpoint is initialized, copy the full endpoint URL for use.
+ * Tested models for `rerank` task:
+ *
+ * * `bge-reranker-base`
+ * * `jina-reranker-v1-turbo-en-GGUF`
+ *
  * @rest_spec_name inference.put_hugging_face
  * @availability stack since=8.12.0 stability=stable visibility=public
  * @availability serverless stability=stable visibility=public
@@ -93,5 +104,10 @@ export interface Request extends RequestBase {
      * Settings used to install the inference model. These settings are specific to the `hugging_face` service.
      */
     service_settings: HuggingFaceServiceSettings
+    /**
+     * Settings to configure the inference task.
+     * These settings are specific to the task type you specified.
+     */
+    task_settings?: HuggingFaceTaskSettings
   }
 }

output/typescript/types.ts

@@ -0,0 +1,16 @@
+summary: A rerank task
+description: Run `PUT _inference/rerank/hugging-face-rerank` to create an inference endpoint that performs a `rerank` task type.
+# method_request: "PUT _inference/rerank/hugging-face-rerank"
+# type: "request"
+value: |-
+  {
+      "service": "hugging_face",
+      "service_settings": {
+          "api_key": "hugging-face-access-token", 
+          "url": "url-endpoint" 
+      },
+      "task_settings": {
+          "return_documents": true,
+          "top_n": 3
+      }
+  }

specification/inference/_types/CommonTypes.ts

@@ -0,0 +1,11 @@
+summary: Rerank task
+description: Run `POST _inference/rerank/bge-reranker-base-mkn` to perform reranking on the example input via Hugging Face
+# method_request: "POST _inference/rerank/bge-reranker-base-mkn"
+# type: "request"
+value: |-
+  {
+    "input": ["luke", "like", "leia", "chewy","r2d2", "star", "wars"],
+    "query": "star wars main character",
+    "return_documents": false,
+    "top_n": 2
+  }

specification/inference/put/PutRequest.ts

@@ -0,0 +1,11 @@
+summary: Rerank task
+description: Run `POST _inference/rerank/bge-reranker-base-mkn` to perform reranking on the example input via Hugging Face
+# method_request: "POST _inference/rerank/bge-reranker-base-mkn"
+# type: "request"
+value: |-
+  {
+    "input": ["luke", "like", "leia", "chewy","r2d2", "star", "wars"],
+    "query": "star wars main character",
+    "return_documents": true,
+    "top_n": 3
+  }

specification/inference/put_hugging_face/PutHuggingFaceRequest.ts

@@ -0,0 +1,18 @@
+summary: Rerank task
+description: >
+  A successful response from `POST _inference/rerank/bge-reranker-base-mkn`.
+# type: "response"
+# response_code:
+value: |-
+  {
+    "rerank": [
+      {
+        "index": 6,
+        "relevance_score": 0.50955844
+      },
+      {
+        "index": 5,
+        "relevance_score": 0.084341794
+      }
+    ]
+  }

specification/inference/put_hugging_face/examples/request/PutHuggingFaceRequestExample2.yaml

@@ -0,0 +1,25 @@
+summary: Rerank task
+description: >
+  A successful response from `POST _inference/rerank/bge-reranker-base-mkn`.
+# type: "response"
+# response_code:
+value: |-
+  {
+    "rerank": [
+      {
+        "index": 6,
+        "relevance_score": 0.50955844,
+        "text": "wars"
+      },
+      {
+        "index": 5,
+        "relevance_score": 0.084341794,
+        "text": "star"
+      },
+      {
+        "index": 3,
+        "relevance_score": 0.004520818,
+        "text": "chewy"
+      }
+    ]
+  }