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.

Sign up for GitHub

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

Jump to bottom

[Frontend] Rerank API (Jina- and Cohere-compatible API) #12376

Open

K-Mistele wants to merge 19 commits into vllm-project:main from K-Mistele:main

+583 −5

Contributor

K-Mistele commented Jan 24, 2025 •

edited by github-actions bot

Loading

This PR introduces an explicit re-rank API at /rerank and /v1/rerank. The same fundamental relevant-scoring functionality with cross-encoder models has been previously implemented in /score, but this API is unique and does not follow any standardized API, and therefore requires custom code to consume.

The endpoints introduced in this PR use the same cross-encoder scoring functionality under the hood, but they reproduce the API behavior of Jina AI's re-ranker and therefore also of Cohere's re-ranker (which is a subset of Jina's API) in terms of API request and response schemas and behavior.

Why Jina and Cohere? OpenAI does not have a re-ranking API. Jina and Cohere's API standards are very common for re-rank APIs, simliar to how OpenAI's chat completions API has become near-universally adopted. The Jina/Cohere re-rank API format has notably been followed by other labs and projects, including:

Qwen
Voyage
Baidu
Siliconflow
localai.io
ragflow

The code in serving_rerank.py is 90% the same as serving_score.py, but implements the rerank API as described above. This enhances vLLM's compatibility with Open-source tools & frameworks that use existing well-defined APIs, allowing vLLM to be used for self-hosted RAG setup with off-the-shelf tooling

FIX infiniflow/ragflow#4316

K-Mistele requested review from DarkLight1337, robertgshaw2-redhat and simon-mo as code owners

January 24, 2025 01:45

github-actions bot commented Jan 24, 2025

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

Add ready label to the PR
Enable auto-merge.

🚀

mergify bot added documentation frontend labels

K-Mistele added 4 commits

January 23, 2025 19:46


          feat: serving_rerank implementation

b6610fb

Signed-off-by: Kyle Mistele <[email protected]>


          fix: imports

a82b4bb

Signed-off-by: Kyle Mistele <[email protected]>


          doc: add example requests and scripts

99acff6

Signed-off-by: Kyle Mistele <[email protected]>


          test: rerank

31b5137

also, add documentation and update client with instructions

Signed-off-by: Kyle Mistele <[email protected]>

K-Mistele force-pushed the main branch from 2c2310f to 31b5137 Compare

January 24, 2025 01:47

DarkLight1337 self-assigned this

DarkLight1337 reviewed

View reviewed changes

examples/online_serving/jinjaai_rerank_client.py Show resolved Hide resolved

DarkLight1337 reviewed

View reviewed changes

docs/source/serving/openai_compatible_server.md Outdated Show resolved Hide resolved

DarkLight1337 reviewed

View reviewed changes

vllm/entrypoints/openai/api_server.py Outdated Show resolved Hide resolved

K-Mistele added 9 commits

January 23, 2025 23:05


          feat: serving_rerank implementation

485e328

Signed-off-by: Kyle Mistele <[email protected]>


          fix: imports

8922f81

Signed-off-by: Kyle Mistele <[email protected]>


          doc: add example requests and scripts

dc0d158

Signed-off-by: Kyle Mistele <[email protected]>


          test: rerank

4ed459b

also, add documentation and update client with instructions

Signed-off-by: Kyle Mistele <[email protected]>


          added /v2/rerank route

676eea0

Signed-off-by: Kyle Mistele <[email protected]>


          fix(docs): extra spaces

b66bcc2

Signed-off-by: Kyle Mistele <[email protected]>


          fix(docs): cross-reference target for rerank API

c44dee4

Signed-off-by: Kyle Mistele <[email protected]>


          fix(tests): needed to break up model quotes

cce2873

Signed-off-by: Kyle Mistele <[email protected]>


          doc(example): update jina example to reflect lack of SDK, add cohere …

a38060f

…example

Signed-off-by: Kyle Mistele <[email protected]>

K-Mistele force-pushed the main branch from 2de6088 to a38060f Compare

January 24, 2025 05:06

DarkLight1337 reviewed

View reviewed changes

vllm/entrypoints/openai/api_server.py Show resolved Hide resolved


          fix: remove logger warnings and make the linter happy

901021f

Signed-off-by: Kyle Mistele <[email protected]>

DarkLight1337 reviewed

View reviewed changes

tests/entrypoints/openai/test_rerank.py Outdated Show resolved Hide resolved

vllm/entrypoints/openai/api_server.py Outdated Show resolved Hide resolved

K-Mistele added 2 commits

January 24, 2025 00:20


          fix: file name

Signed-off-by: Kyle Mistele <[email protected]>


          fix(nit): ordering on assertions

36e85a5

Signed-off-by: Kyle Mistele <[email protected]>

K-Mistele added 3 commits

January 24, 2025 12:59


          fix(tests): was using score instead of rerank

4adb94b

Signed-off-by: Kyle Mistele <[email protected]>


          fix(api): use rereank as the default API for scoring

dc92240

Signed-off-by: Kyle Mistele <[email protected]>


          fix(merge)

330aa22

Signed-off-by: Kyle Mistele <[email protected]>

Contributor Author

K-Mistele commented Jan 24, 2025

Is there an issue with docker or dependencies being downloaded? CI is failing for reasons that seem unrelated to my code. Something about invalid flashinfer wheels?

mgoin reviewed

View reviewed changes

vllm/entrypoints/openai/api_server.py

Comment on lines -515 to +555

    
                      "default": (ScoreRequest, create_score),

                      "default": (RerankRequest, do_rerank)

Member

mgoin Jan 24, 2025

Why can we just remove (ScoreRequest, create_score)? I imagine we want to keep that interface around and just offer rerank as a new option

Contributor Author

K-Mistele Jan 24, 2025

This was suggested by @DarkLight1337

vllm/entrypoints/openai/api_server.py

Comment on lines +805 to +809

    
                  state.jinaai_serving_reranking = JinaAIServingRerank(

                      engine_client,

                      model_config,

                      state.openai_serving_models,

                      request_logger=request_logger)

Member

mgoin Jan 24, 2025

I think we should only make it in the rerank task case, like above

Suggested change

      
                state.jinaai_serving_reranking = JinaAIServingRerank(
          
                    engine_client,
          
                    model_config,
          
                    state.openai_serving_models,
          
                    request_logger=request_logger)
          
                state.jinaai_serving_reranking = JinaAIServingRerank(
          
                    engine_client,
          
                    model_config,
          
                    state.openai_serving_models,
          
                    request_logger=request_logger
          
                ) if model_config.task == "rerank" else None

examples/online_serving/jinjaai_rerank_client.py

Member

mgoin Jan 24, 2025

This seems duplicated from examples/online_serving/jinaai_rerank_client.py?

examples/online_serving/jinaai_rerank_client.py

    
              Example of using the OpenAI entrypoint's rerank API which is compatible with

              Jina and Cohere https://jina.ai/reranker

              run: vllm serve --model BAAI/bge-reranker-base

Member

mgoin Jan 24, 2025

Suggested change

      
            run: vllm serve --model BAAI/bge-reranker-base
          
            run: vllm serve BAAI/bge-reranker-base

docs/source/serving/openai_compatible_server.md

    
            @@ -50,6 +50,11 @@ In addition, we have the following custom APIs:
          
                - Applicable to all [pooling models](../models/pooling_models.md).

              - [Score API](#score-api) (`/score`)

                - Only applicable to [cross-encoder models](../models/pooling_models.md) (`--task score`).

              - [Re-rank API](#rerank-api) (`/rerank`, `/v1/rerank`)

Member

mgoin Jan 24, 2025

Suggested change

      
            - [Re-rank API](#rerank-api) (`/rerank`, `/v1/rerank`)
          
            - [Re-rank API](#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

DarkLight1337 DarkLight1337 left review comments

mgoin mgoin left review comments

robertgshaw2-redhat Awaiting requested review from robertgshaw2-redhat robertgshaw2-redhat is a code owner

simon-mo Awaiting requested review from simon-mo simon-mo is a code owner

At least 1 approving review is required to merge this pull request.

Labels

documentation frontend