Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 11 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 50 additions & 5 deletions output/openapi/elasticsearch-openapi.json

Large diffs are not rendered by default.

55 changes: 50 additions & 5 deletions output/openapi/elasticsearch-serverless-openapi.json

Large diffs are not rendered by default.

715 changes: 65 additions & 650 deletions output/schema/schema-serverless.json

Large diffs are not rendered by default.

153 changes: 118 additions & 35 deletions output/schema/schema.json

Large diffs are not rendered by default.

10 changes: 8 additions & 2 deletions output/typescript/types.ts

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1,727 changes: 527 additions & 1,200 deletions package-lock.json

Large diffs are not rendered by default.

14 changes: 14 additions & 0 deletions specification/inference/_types/CommonTypes.ts
Original file line number Diff line number Diff line change
Expand Up @@ -908,9 +908,23 @@ 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,
sparse_embedding,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jonathan-buttner same thing here

text_embedding
}

Expand Down
6 changes: 4 additions & 2 deletions specification/inference/_types/TaskType.ts
Original file line number Diff line number Diff line change
Expand Up @@ -86,9 +86,11 @@ export enum TaskTypeGoogleVertexAI {
}

export enum TaskTypeHuggingFace {
text_embedding,
chat_completion,
completion
completion,
rerank,
sparse_embedding,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jonathan-buttner Please confirm if we want to include sparse_embedding into a list of supported HF tasks. Because HF models don't provide sparse embeddings support out of the box.

text_embedding
}

export enum TaskTypeMistral {
Expand Down
2 changes: 1 addition & 1 deletion specification/inference/put/PutRequest.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 (`text_embedding`)
* * Hugging Face (`chat_completion`, `completion`, `rerank`, `sparse_embedding`, `text_embedding`)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jonathan-buttner could you please confirm that we want to add sparse_embedding task here?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@davidkyle do we want to document hugging face support for sparse embedding aka elser? I don't think we ever released elser on hugging face for customers right?

* * Mistral (`text_embedding`)
* * OpenAI (`chat_completion`, `completion`, `text_embedding`)
* * VoyageAI (`text_embedding`, `rerank`)
Expand Down
16 changes: 16 additions & 0 deletions specification/inference/put_hugging_face/PutHuggingFaceRequest.ts
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you, Jonathan. Done

Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import { Id } from '@_types/common'
import {
HuggingFaceServiceSettings,
HuggingFaceServiceType,
HuggingFaceTaskSettings,
HuggingFaceTaskType
} from '@inference/_types/CommonTypes'
import { InferenceChunkingSettings } from '@inference/_types/Services'
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
}
}
Original file line number Diff line number Diff line change
@@ -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
}
}
Original file line number Diff line number Diff line change
@@ -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
}
Original file line number Diff line number Diff line change
@@ -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
}
Original file line number Diff line number Diff line change
@@ -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
}
]
}
Original file line number Diff line number Diff line change
@@ -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"
}
]
}