@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1c3c975766
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 22:15 UTC

[FEATURE] Enables /score endpoint for embedding models (#12846)

Gabriel Marinho committed Feb 21, 2025 at 03:09 UTC 1c3c9757668238485e880afff2c42611ac40763b
11 files changed +599 -522
docs/source/models/pooling_models.md
+1 -2
@@ -108,8 +108,7 @@ A code example can be found here: <gh-file:examples/offline_inference/basic/clas
108 ### `LLM.score`
109
110 The {class}`~vllm.LLM.score` method outputs similarity scores between sentence pairs.
111 -It is primarily designed for [cross-encoder models](https://www.sbert.net/examples/applications/cross-encoder/README.html).
112 -These types of models serve as rerankers between candidate query-document pairs in RAG systems.
111 +It is designed for embedding models and cross encoder models. Embedding models use cosine similarity, and [cross-encoder models](https://www.sbert.net/examples/applications/cross-encoder/README.html) serve as rerankers between candidate query-document pairs in RAG systems.
112
113 :::{note}
114 vLLM can only perform the model inference component (e.g. embedding, reranking) of RAG.
docs/source/serving/openai_compatible_server.md
+5 -5
@@ -51,7 +51,7 @@ In addition, we have the following custom APIs:
51 - [Pooling API](#pooling-api) (`/pooling`)
52 - Applicable to all [pooling models](../models/pooling_models.md).
53 - [Score API](#score-api) (`/score`)
54 - - Only applicable to [cross-encoder models](../models/pooling_models.md) (`--task score`).
54 + - Applicable to embedding models and [cross-encoder models](../models/pooling_models.md) (`--task score`).
55 - [Re-rank API](#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
56 - Implements [Jina AI's v1 re-rank API](https://jina.ai/reranker/)
57 - Also compatible with [Cohere's v1 & v2 re-rank APIs](https://docs.cohere.com/v2/reference/rerank)
@@ -333,10 +333,10 @@ Code example: <gh-file:examples/online_serving/openai_pooling_client.py>
333
334 ### Score API
335
336 -Our Score API applies a cross-encoder model to predict scores for sentence pairs.
336 +Our Score API can apply a cross-encoder model or an embedding model to predict scores for sentence pairs. When using an embedding model the score corresponds to the cosine similarity between each embedding pair.
337 Usually, the score for a sentence pair refers to the similarity between two sentences, on a scale of 0 to 1.
338
339 -You can find the documentation for these kind of models at [sbert.net](https://www.sbert.net/docs/package_reference/cross_encoder/cross_encoder.html).
339 +You can find the documentation for cross encoder models at [sbert.net](https://www.sbert.net/docs/package_reference/cross_encoder/cross_encoder.html).
340
341 Code example: <gh-file:examples/online_serving/openai_cross_encoder_score.py>
342
@@ -496,11 +496,11 @@ The following extra parameters are supported:
496
497 ### Re-rank API
498
499 -Our Re-rank API applies a cross-encoder model to predict relevant scores between a single query, and
499 +Our Re-rank API can apply an embedding model or a cross-encoder model to predict relevant scores between a single query, and
500 each of a list of documents. Usually, the score for a sentence pair refers to the similarity between two sentences, on
501 a scale of 0 to 1.
502
503 -You can find the documentation for these kind of models at [sbert.net](https://www.sbert.net/docs/package_reference/cross_encoder/cross_encoder.html).
503 +You can find the documentation for cross encoder models at [sbert.net](https://www.sbert.net/docs/package_reference/cross_encoder/cross_encoder.html).
504
505 The rerank endpoints support popular re-rank models such as `BAAI/bge-reranker-base` and other models supporting the
506 `score` task. Additionally, `/rerank`, `/v1/rerank`, and `/v2/rerank`
tests/entrypoints/openai/test_rerank.py
+2 -4
@@ -8,17 +8,17 @@ from vllm.entrypoints.openai.protocol import RerankResponse
8 from ...utils import RemoteOpenAIServer
9
10 MODEL_NAME = "BAAI/bge-reranker-base"
11 +DTYPE = "bfloat16"
12
13
14 @pytest.fixture(scope="module")
15 def server():
15 - args = ["--enforce-eager", "--max-model-len", "100"]
16 + args = ["--enforce-eager", "--max-model-len", "100", "--dtype", DTYPE]
17
18 with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
19 yield remote_server
20
21
21 -@pytest.mark.asyncio
22 @pytest.mark.parametrize("model_name", [MODEL_NAME])
23 def test_rerank_texts(server: RemoteOpenAIServer, model_name: str):
24 query = "What is the capital of France?"
@@ -42,7 +42,6 @@ def test_rerank_texts(server: RemoteOpenAIServer, model_name: str):
42 assert rerank.results[1].relevance_score <= 0.01
43
44
45 -@pytest.mark.asyncio
45 @pytest.mark.parametrize("model_name", [MODEL_NAME])
46 def test_top_n(server: RemoteOpenAIServer, model_name: str):
47 query = "What is the capital of France?"
@@ -68,7 +67,6 @@ def test_top_n(server: RemoteOpenAIServer, model_name: str):
67 assert rerank.results[1].relevance_score <= 0.01
68
69
71 -@pytest.mark.asyncio
70 @pytest.mark.parametrize("model_name", [MODEL_NAME])
71 def test_rerank_max_model_len(server: RemoteOpenAIServer, model_name: str):
72
tests/entrypoints/openai/test_score.py
+173 -111
@@ -1,123 +1,185 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 +import math
4 +from typing import Any
5 +
6 import pytest
7 import requests
8 +import torch.nn.functional as F
9 +from torch import tensor
10
11 from vllm.entrypoints.openai.protocol import ScoreResponse
12
13 from ...utils import RemoteOpenAIServer
14
10 -MODEL_NAME = "BAAI/bge-reranker-v2-m3"
11 -
12 -
13 -@pytest.fixture(scope="module")
14 -def server():
15 - args = ["--enforce-eager", "--max-model-len", "100"]
16 -
17 - with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
15 +MODELS = [
16 + {
17 + "name": "BAAI/bge-reranker-v2-m3",
18 + "is_cross_encoder": True
19 + },
20 + {
21 + "name": "BAAI/bge-base-en-v1.5",
22 + "is_cross_encoder": False
23 + },
24 +]
25 +DTYPE = "half"
26 +
27 +
28 +def run_transformers(hf_model, model, text_pairs):
29 + if model["is_cross_encoder"]:
30 + return hf_model.predict(text_pairs).tolist()
31 + else:
32 + hf_embeddings = [
33 + hf_model.encode(text_pair) for text_pair in text_pairs
34 + ]
35 + return [
36 + F.cosine_similarity(tensor(pair[0]), tensor(pair[1]), dim=0)
37 + for pair in hf_embeddings
38 + ]
39 +
40 +
41 +@pytest.fixture(scope="class", params=MODELS)
42 +def model(request):
43 + yield request.param
44 +
45 +
46 +@pytest.fixture(scope="class")
47 +def server(model: dict[str, Any]):
48 + args = ["--enforce-eager", "--max-model-len", "100", "--dtype", DTYPE]
49 +
50 + with RemoteOpenAIServer(model["name"], args) as remote_server:
51 yield remote_server
52
53
21 -@pytest.mark.asyncio
22 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
23 -def test_text_1_str_text_2_list(server: RemoteOpenAIServer, model_name: str):
24 - text_1 = "What is the capital of France?"
25 - text_2 = [
26 - "The capital of Brazil is Brasilia.", "The capital of France is Paris."
27 - ]
28 -
29 - score_response = requests.post(server.url_for("score"),
30 - json={
31 - "model": model_name,
32 - "text_1": text_1,
33 - "text_2": text_2,
34 - })
35 - score_response.raise_for_status()
36 - score = ScoreResponse.model_validate(score_response.json())
37 -
38 - assert score.id is not None
39 - assert score.data is not None
40 - assert len(score.data) == 2
41 - assert score.data[0].score <= 0.01
42 - assert score.data[1].score >= 0.9
43 -
44 -
45 -@pytest.mark.asyncio
46 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
47 -def test_text_1_list_text_2_list(server: RemoteOpenAIServer, model_name: str):
48 - text_1 = [
49 - "What is the capital of the United States?",
50 - "What is the capital of France?"
51 - ]
52 - text_2 = [
53 - "The capital of Brazil is Brasilia.", "The capital of France is Paris."
54 - ]
55 -
56 - score_response = requests.post(server.url_for("score"),
57 - json={
58 - "model": model_name,
59 - "text_1": text_1,
60 - "text_2": text_2,
61 - })
62 - score_response.raise_for_status()
63 - score = ScoreResponse.model_validate(score_response.json())
64 -
65 - assert score.id is not None
66 - assert score.data is not None
67 - assert len(score.data) == 2
68 - assert score.data[0].score <= 0.01
69 - assert score.data[1].score >= 0.9
70 -
71 -
72 -@pytest.mark.asyncio
73 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
74 -def test_text_1_str_text_2_str(server: RemoteOpenAIServer, model_name: str):
75 - text_1 = "What is the capital of France?"
76 - text_2 = "The capital of France is Paris."
77 -
78 - score_response = requests.post(server.url_for("score"),
79 - json={
80 - "model": model_name,
81 - "text_1": text_1,
82 - "text_2": text_2,
83 - })
84 - score_response.raise_for_status()
85 - score = ScoreResponse.model_validate(score_response.json())
86 -
87 - assert score.id is not None
88 - assert score.data is not None
89 - assert len(score.data) == 1
90 - assert score.data[0].score >= 0.9
91 -
92 -
93 -@pytest.mark.asyncio
94 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
95 -def test_score_max_model_len(server: RemoteOpenAIServer, model_name: str):
96 -
97 - text_1 = "What is the capital of France?" * 20
98 - text_2 = [
99 - "The capital of Brazil is Brasilia.", "The capital of France is Paris."
100 - ]
101 -
102 - score_response = requests.post(server.url_for("score"),
103 - json={
104 - "model": model_name,
105 - "text_1": text_1,
106 - "text_2": text_2,
107 - })
108 - assert score_response.status_code == 400
109 - # Assert just a small fragments of the response
110 - assert "Please reduce the length of the input." in \
111 - score_response.text
112 -
113 - # Test truncation
114 - score_response = requests.post(server.url_for("score"),
115 - json={
116 - "model": model_name,
117 - "text_1": text_1,
118 - "text_2": text_2,
119 - "truncate_prompt_tokens": 101
120 - })
121 - assert score_response.status_code == 400
122 - assert "Please, select a smaller truncation size." in \
123 - score_response.text
54 +@pytest.fixture(scope="class")
55 +def runner(model: dict[str, Any], hf_runner):
56 + kwargs = {
57 + "dtype": DTYPE,
58 + "is_cross_encoder" if model["is_cross_encoder"]\
59 + else "is_sentence_transformer": True
60 + }
61 +
62 + with hf_runner(model["name"], **kwargs) as hf_model:
63 + yield hf_model
64 +
65 +
66 +class TestModel:
67 +
68 + def test_text_1_str_text_2_list(self, server: RemoteOpenAIServer,
69 + model: dict[str, Any], runner):
70 + text_1 = "What is the capital of France?"
71 + text_2 = [
72 + "The capital of Brazil is Brasilia.",
73 + "The capital of France is Paris."
74 + ]
75 +
76 + score_response = requests.post(server.url_for("score"),
77 + json={
78 + "model": model["name"],
79 + "text_1": text_1,
80 + "text_2": text_2,
81 + })
82 + score_response.raise_for_status()
83 + score = ScoreResponse.model_validate(score_response.json())
84 +
85 + assert score.id is not None
86 + assert score.data is not None
87 + assert len(score.data) == 2
88 +
89 + vllm_outputs = [d.score for d in score.data]
90 +
91 + text_pairs = [[text_1, text_2[0]], [text_1, text_2[1]]]
92 + hf_outputs = run_transformers(runner, model, text_pairs)
93 +
94 + for i in range(len(vllm_outputs)):
95 + assert math.isclose(hf_outputs[i], vllm_outputs[i], rel_tol=0.01)
96 +
97 + def test_text_1_list_text_2_list(self, server: RemoteOpenAIServer,
98 + model: dict[str, Any], runner):
99 + text_1 = [
100 + "What is the capital of the United States?",
101 + "What is the capital of France?"
102 + ]
103 + text_2 = [
104 + "The capital of Brazil is Brasilia.",
105 + "The capital of France is Paris."
106 + ]
107 +
108 + score_response = requests.post(server.url_for("score"),
109 + json={
110 + "model": model["name"],
111 + "text_1": text_1,
112 + "text_2": text_2,
113 + })
114 + score_response.raise_for_status()
115 + score = ScoreResponse.model_validate(score_response.json())
116 +
117 + assert score.id is not None
118 + assert score.data is not None
119 + assert len(score.data) == 2
120 +
121 + vllm_outputs = [d.score for d in score.data]
122 +
123 + text_pairs = [[text_1[0], text_2[0]], [text_1[1], text_2[1]]]
124 + hf_outputs = run_transformers(runner, model, text_pairs)
125 +
126 + for i in range(len(vllm_outputs)):
127 + assert math.isclose(hf_outputs[i], vllm_outputs[i], rel_tol=0.01)
128 +
129 + def test_text_1_str_text_2_str(self, server: RemoteOpenAIServer,
130 + model: dict[str, Any], runner):
131 + text_1 = "What is the capital of France?"
132 + text_2 = "The capital of France is Paris."
133 +
134 + score_response = requests.post(server.url_for("score"),
135 + json={
136 + "model": model["name"],
137 + "text_1": text_1,
138 + "text_2": text_2,
139 + })
140 + score_response.raise_for_status()
141 + score = ScoreResponse.model_validate(score_response.json())
142 +
143 + assert score.id is not None
144 + assert score.data is not None
145 + assert len(score.data) == 1
146 +
147 + vllm_outputs = [d.score for d in score.data]
148 +
149 + text_pairs = [[text_1, text_2]]
150 + hf_outputs = run_transformers(runner, model, text_pairs)
151 +
152 + for i in range(len(vllm_outputs)):
153 + assert math.isclose(hf_outputs[i], vllm_outputs[i], rel_tol=0.01)
154 +
155 + def test_score_max_model_len(self, server: RemoteOpenAIServer,
156 + model: dict[str, Any]):
157 +
158 + text_1 = "What is the capital of France?" * 20
159 + text_2 = [
160 + "The capital of Brazil is Brasilia.",
161 + "The capital of France is Paris."
162 + ]
163 +
164 + score_response = requests.post(server.url_for("score"),
165 + json={
166 + "model": model["name"],
167 + "text_1": text_1,
168 + "text_2": text_2,
169 + })
170 + assert score_response.status_code == 400
171 + # Assert just a small fragments of the response
172 + assert "Please reduce the length of the input." in \
173 + score_response.text
174 +
175 + # Test truncation
176 + score_response = requests.post(server.url_for("score"),
177 + json={
178 + "model": model["name"],
179 + "text_1": text_1,
180 + "text_2": text_2,
181 + "truncate_prompt_tokens": 101
182 + })
183 + assert score_response.status_code == 400
184 + assert "Please, select a smaller truncation size." in \
185 + score_response.text
vllm/entrypoints/llm.py
+13 -33
@@ -7,7 +7,6 @@ from typing import (Any, Callable, ClassVar, Dict, List, Optional, Sequence,
7 Tuple, Type, Union, cast, overload)
8
9 import cloudpickle
10 -import torch
10 import torch.nn as nn
11 from tqdm import tqdm
12 from typing_extensions import TypeVar, deprecated
@@ -25,6 +24,8 @@ from vllm.entrypoints.chat_utils import (ChatCompletionMessageParam,
24 apply_mistral_chat_template,
25 parse_chat_messages,
26 resolve_chat_template_content_format)
27 +from vllm.entrypoints.score_utils import (_cosine_similarity,
28 + _validate_score_input_lens)
29 from vllm.inputs import PromptType, SingletonPrompt, TextPrompt, TokensPrompt
30 from vllm.inputs.parse import is_token_prompt, parse_and_batch_prompt
31 from vllm.logger import init_logger
@@ -1010,40 +1011,25 @@ class LLM:
1011 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
1012 ) -> List[ScoringRequestOutput]:
1013
1013 - encoded_output = self.encode(
1014 + encoded_output: List[PoolingRequestOutput] = self.encode(
1015 text_1 + text_2,
1016 use_tqdm=use_tqdm,
1017 lora_request=lora_request,
1018 prompt_adapter_request=prompt_adapter_request)
1018 - encoded_output_1 = encoded_output[0:len(text_1)]
1019 - encoded_output_2 = encoded_output[len(text_1):]
1019 +
1020 + encoded_output_1: List[PoolingRequestOutput] = encoded_output[
1021 + 0:len(text_1)]
1022 + encoded_output_2: List[PoolingRequestOutput] = encoded_output[
1023 + len(text_1):]
1024
1025 if len(encoded_output_1) == 1:
1026 encoded_output_1 = encoded_output_1 * len(encoded_output_2)
1027
1024 - output_pairs = [(t1, t2)
1025 - for t1, t2 in zip(encoded_output_1, encoded_output_2)]
1026 -
1027 - scores = []
1028 - scorer = torch.nn.CosineSimilarity(0)
1029 -
1030 - for embed_1, embed_2 in output_pairs:
1031 - pair_score = scorer(embed_1.outputs.data, embed_2.outputs.data)
1028 + scores: List[PoolingRequestOutput] = []
1029
1033 - if (pad_token_id := getattr(tokenizer, "pad_token_id",
1034 - None)) is not None:
1035 - tokens = embed_1.prompt_token_ids + [
1036 - pad_token_id
1037 - ] + embed_2.prompt_token_ids
1038 - else:
1039 - tokens = embed_1.prompt_token_ids + embed_2.prompt_token_ids
1040 -
1041 - scores.append(
1042 - PoolingRequestOutput(
1043 - request_id=f"{embed_1.request_id}_{embed_2.request_id}",
1044 - outputs=pair_score,
1045 - prompt_token_ids=tokens,
1046 - finished=True))
1030 + scores = _cosine_similarity(tokenizer=tokenizer,
1031 + embed_1=encoded_output_1,
1032 + embed_2=encoded_output_2)
1033
1034 items = self.engine_class.validate_outputs(scores,
1035 PoolingRequestOutput)
@@ -1183,12 +1169,7 @@ class LLM:
1169 text_2 = [text_2]
1170 input_text_2: List[str] = [ensure_str(t) for t in text_2]
1171
1186 - if len(input_text_1) > 1 and len(input_text_1) != len(input_text_2):
1187 - raise ValueError("Input lengths must be either 1:1, 1:N or N:N")
1188 - if len(input_text_1) == 0:
1189 - raise ValueError("At least one text element must be given")
1190 - if len(input_text_2) == 0:
1191 - raise ValueError("At least one text_pair element must be given")
1172 + _validate_score_input_lens(input_text_1, input_text_2)
1173
1174 if self.llm_engine.model_config.is_cross_encoder:
1175 return self._cross_encoding_score(tokenizer, input_text_1,
@@ -1197,7 +1178,6 @@ class LLM:
1178 lora_request,
1179 prompt_adapter_request)
1180 else:
1200 -
1181 return self._embedding_score(
1182 tokenizer,
1183 input_text_1, # type: ignore[arg-type]
vllm/entrypoints/openai/api_server.py
+8 -9
@@ -73,8 +73,7 @@ from vllm.entrypoints.openai.serving_engine import OpenAIServing
73 from vllm.entrypoints.openai.serving_models import (BaseModelPath,
74 OpenAIServingModels)
75 from vllm.entrypoints.openai.serving_pooling import OpenAIServingPooling
76 -from vllm.entrypoints.openai.serving_rerank import JinaAIServingRerank
77 -from vllm.entrypoints.openai.serving_score import OpenAIServingScores
76 +from vllm.entrypoints.openai.serving_score import ServingScores
77 from vllm.entrypoints.openai.serving_tokenization import (
78 OpenAIServingTokenization)
79 from vllm.entrypoints.openai.serving_transcription import (
@@ -320,12 +319,12 @@ def embedding(request: Request) -> Optional[OpenAIServingEmbedding]:
319 return request.app.state.openai_serving_embedding
320
321
323 -def score(request: Request) -> Optional[OpenAIServingScores]:
322 +def score(request: Request) -> Optional[ServingScores]:
323 return request.app.state.openai_serving_scores
324
325
327 -def rerank(request: Request) -> Optional[JinaAIServingRerank]:
328 - return request.app.state.jinaai_serving_reranking
326 +def rerank(request: Request) -> Optional[ServingScores]:
327 + return request.app.state.openai_serving_scores
328
329
330 def tokenization(request: Request) -> OpenAIServingTokenization:
@@ -866,13 +865,13 @@ async def init_app_state(
865 chat_template=resolved_chat_template,
866 chat_template_content_format=args.chat_template_content_format,
867 ) if model_config.task == "embed" else None
869 - state.openai_serving_scores = OpenAIServingScores(
868 + state.openai_serving_scores = ServingScores(
869 engine_client,
870 model_config,
871 state.openai_serving_models,
873 - request_logger=request_logger
874 - ) if model_config.task == "score" else None
875 - state.jinaai_serving_reranking = JinaAIServingRerank(
872 + request_logger=request_logger) if model_config.task in (
873 + "score", "embed", "pooling") else None
874 + state.jinaai_serving_reranking = ServingScores(
875 engine_client,
876 model_config,
877 state.openai_serving_models,
vllm/entrypoints/openai/run_batch.py
+17 -14
@@ -26,7 +26,7 @@ from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
26 from vllm.entrypoints.openai.serving_embedding import OpenAIServingEmbedding
27 from vllm.entrypoints.openai.serving_models import (BaseModelPath,
28 OpenAIServingModels)
29 -from vllm.entrypoints.openai.serving_score import OpenAIServingScores
29 +from vllm.entrypoints.openai.serving_score import ServingScores
30 from vllm.usage.usage_lib import UsageContext
31 from vllm.utils import FlexibleArgumentParser, random_uuid
32 from vllm.version import __version__ as VLLM_VERSION
@@ -342,7 +342,7 @@ async def main(args):
342 chat_template=None,
343 chat_template_content_format="auto",
344 ) if model_config.task == "embed" else None
345 - openai_serving_scores = (OpenAIServingScores(
345 + openai_serving_scores = (ServingScores(
346 engine,
347 model_config,
348 openai_serving_models,
@@ -364,9 +364,9 @@ async def main(args):
364
365 # Determine the type of request and run it.
366 if request.url == "/v1/chat/completions":
367 - handler_fn = (None if openai_serving_chat is None else
368 - openai_serving_chat.create_chat_completion)
369 - if handler_fn is None:
367 + chat_handler_fn = (None if openai_serving_chat is None else
368 + openai_serving_chat.create_chat_completion)
369 + if chat_handler_fn is None:
370 response_futures.append(
371 make_async_error_request_output(
372 request,
@@ -375,12 +375,13 @@ async def main(args):
375 ))
376 continue
377
378 - response_futures.append(run_request(handler_fn, request, tracker))
378 + response_futures.append(
379 + run_request(chat_handler_fn, request, tracker))
380 tracker.submitted()
381 elif request.url == "/v1/embeddings":
381 - handler_fn = (None if openai_serving_embedding is None else
382 - openai_serving_embedding.create_embedding)
383 - if handler_fn is None:
382 + embed_handler_fn = (None if openai_serving_embedding is None else
383 + openai_serving_embedding.create_embedding)
384 + if embed_handler_fn is None:
385 response_futures.append(
386 make_async_error_request_output(
387 request,
@@ -388,12 +389,13 @@ async def main(args):
389 ))
390 continue
391
391 - response_futures.append(run_request(handler_fn, request, tracker))
392 + response_futures.append(
393 + run_request(embed_handler_fn, request, tracker))
394 tracker.submitted()
395 elif request.url == "/v1/score":
394 - handler_fn = (None if openai_serving_scores is None else
395 - openai_serving_scores.create_score)
396 - if handler_fn is None:
396 + score_handler_fn = (None if openai_serving_scores is None else
397 + openai_serving_scores.create_score)
398 + if score_handler_fn is None:
399 response_futures.append(
400 make_async_error_request_output(
401 request,
@@ -401,7 +403,8 @@ async def main(args):
403 ))
404 continue
405
404 - response_futures.append(run_request(handler_fn, request, tracker))
406 + response_futures.append(
407 + run_request(score_handler_fn, request, tracker))
408 tracker.submitted()
409 else:
410 response_futures.append(
vllm/entrypoints/openai/serving_engine.py
+2 -2
@@ -52,8 +52,8 @@ from vllm.utils import is_list_of, make_async, random_uuid
52 logger = init_logger(__name__)
53
54 CompletionLikeRequest = Union[CompletionRequest, DetokenizeRequest,
55 - EmbeddingCompletionRequest, ScoreRequest,
56 - TokenizeCompletionRequest]
55 + EmbeddingCompletionRequest, RerankRequest,
56 + ScoreRequest, TokenizeCompletionRequest]
57
58 ChatLikeRequest = Union[ChatCompletionRequest, EmbeddingChatRequest,
59 TokenizeChatRequest]
vllm/entrypoints/openai/serving_rerank.py deleted
-208
@@ -1,208 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -
3 -import asyncio
4 -from typing import Any, AsyncGenerator, Dict, List, Optional, Union, cast
5 -
6 -from fastapi import Request
7 -
8 -from vllm.config import ModelConfig
9 -from vllm.engine.protocol import EngineClient
10 -from vllm.entrypoints.logger import RequestLogger
11 -from vllm.entrypoints.openai.protocol import (ErrorResponse, RerankDocument,
12 - RerankRequest, RerankResponse,
13 - RerankResult, RerankUsage)
14 -from vllm.entrypoints.openai.serving_engine import OpenAIServing
15 -from vllm.entrypoints.openai.serving_models import OpenAIServingModels
16 -from vllm.inputs.data import TokensPrompt
17 -from vllm.logger import init_logger
18 -from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput
19 -from vllm.transformers_utils.tokenizers.mistral import MistralTokenizer
20 -from vllm.utils import make_async, merge_async_iterators
21 -
22 -logger = init_logger(__name__)
23 -
24 -
25 -class JinaAIServingRerank(OpenAIServing):
26 -
27 - def __init__(
28 - self,
29 - engine_client: EngineClient,
30 - model_config: ModelConfig,
31 - models: OpenAIServingModels,
32 - *,
33 - request_logger: Optional[RequestLogger],
34 - ) -> None:
35 - super().__init__(engine_client=engine_client,
36 - model_config=model_config,
37 - models=models,
38 - request_logger=request_logger)
39 -
40 - async def do_rerank(
41 - self,
42 - request: RerankRequest,
43 - raw_request: Optional[Request] = None
44 - ) -> Union[RerankResponse, ErrorResponse]:
45 - """
46 - Rerank API based on JinaAI's rerank API; implements the same
47 - API interface. Designed for compatibility with off-the-shelf
48 - tooling, since this is a common standard for reranking APIs
49 -
50 - See example client implementations at
51 - https://github.com/infiniflow/ragflow/blob/main/rag/llm/rerank_model.py
52 - numerous clients use this standard.
53 - """
54 - error_check_ret = await self._check_model(request)
55 - if error_check_ret is not None:
56 - return error_check_ret
57 -
58 - model_name = request.model
59 - request_id = f"rerank-{self._base_request_id(raw_request)}"
60 - truncate_prompt_tokens = request.truncate_prompt_tokens
61 - query = request.query
62 - documents = request.documents
63 - request_prompts = []
64 - engine_prompts = []
65 - top_n = request.top_n if request.top_n > 0 else len(documents)
66 -
67 - try:
68 - (
69 - lora_request,
70 - prompt_adapter_request,
71 - ) = self._maybe_get_adapters(request)
72 -
73 - tokenizer = await self.engine_client.get_tokenizer(lora_request)
74 -
75 - if prompt_adapter_request is not None:
76 - raise NotImplementedError("Prompt adapter is not supported "
77 - "for scoring models")
78 -
79 - if isinstance(tokenizer, MistralTokenizer):
80 - raise ValueError(
81 - "MistralTokenizer not supported for cross-encoding")
82 -
83 - if not self.model_config.is_cross_encoder:
84 - raise ValueError("Model is not cross encoder.")
85 -
86 - if truncate_prompt_tokens is not None and \
87 - truncate_prompt_tokens > self.max_model_len:
88 - raise ValueError(
89 - f"truncate_prompt_tokens value ({truncate_prompt_tokens}) "
90 - f"is greater than max_model_len ({self.max_model_len})."
91 - f" Please, select a smaller truncation size.")
92 - for doc in documents:
93 - request_prompt = f"{query}{tokenizer.sep_token}{doc}"
94 - tokenization_kwargs: Dict[str, Any] = {}
95 - if truncate_prompt_tokens is not None:
96 - tokenization_kwargs["truncation"] = True
97 - tokenization_kwargs["max_length"] = truncate_prompt_tokens
98 -
99 - tokenize_async = make_async(tokenizer.__call__,
100 - executor=self._tokenizer_executor)
101 - prompt_inputs = await tokenize_async(text=query,
102 - text_pair=doc,
103 - **tokenization_kwargs)
104 -
105 - input_ids = prompt_inputs["input_ids"]
106 - text_token_prompt = \
107 - self._validate_input(request, input_ids, request_prompt)
108 - engine_prompt = TokensPrompt(
109 - prompt_token_ids=text_token_prompt["prompt_token_ids"],
110 - token_type_ids=prompt_inputs.get("token_type_ids"))
111 -
112 - request_prompts.append(request_prompt)
113 - engine_prompts.append(engine_prompt)
114 -
115 - except ValueError as e:
116 - logger.exception("Error in preprocessing prompt inputs")
117 - return self.create_error_response(str(e))
118 -
119 - # Schedule the request and get the result generator.
120 - generators: List[AsyncGenerator[PoolingRequestOutput, None]] = []
121 -
122 - try:
123 - pooling_params = request.to_pooling_params()
124 -
125 - for i, engine_prompt in enumerate(engine_prompts):
126 - request_id_item = f"{request_id}-{i}"
127 -
128 - self._log_inputs(request_id_item,
129 - request_prompts[i],
130 - params=pooling_params,
131 - lora_request=lora_request,
132 - prompt_adapter_request=prompt_adapter_request)
133 -
134 - trace_headers = (None if raw_request is None else await
135 - self._get_trace_headers(raw_request.headers))
136 -
137 - generator = self.engine_client.encode(
138 - engine_prompt,
139 - pooling_params,
140 - request_id_item,
141 - lora_request=lora_request,
142 - trace_headers=trace_headers,
143 - priority=request.priority,
144 - )
145 -
146 - generators.append(generator)
147 - except ValueError as e:
148 - # TODO: Use a vllm-specific Validation Error
149 - return self.create_error_response(str(e))
150 - result_generator = merge_async_iterators(*generators)
151 -
152 - num_prompts = len(engine_prompts)
153 -
154 - # Non-streaming response
155 - final_res_batch: List[Optional[PoolingRequestOutput]]
156 - final_res_batch = [None] * num_prompts
157 -
158 - try:
159 - async for i, res in result_generator:
160 - final_res_batch[i] = res
161 -
162 - assert all(final_res is not None for final_res in final_res_batch)
163 -
164 - final_res_batch_checked = cast(List[PoolingRequestOutput],
165 - final_res_batch)
166 -
167 - response = self.request_output_to_rerank_response(
168 - final_res_batch_checked, request_id, model_name, documents,
169 - top_n)
170 - except asyncio.CancelledError:
171 - return self.create_error_response("Client disconnected")
172 - except ValueError as e:
173 - # TODO: Use a vllm-specific Validation Error
174 - return self.create_error_response(str(e))
175 -
176 - return response
177 -
178 - def request_output_to_rerank_response(
179 - self, final_res_batch: List[PoolingRequestOutput], request_id: str,
180 - model_name: str, documents: List[str],
181 - top_n: int) -> RerankResponse:
182 - """
183 - Convert the output of do_rank to a RerankResponse
184 - """
185 - results: List[RerankResult] = []
186 - num_prompt_tokens = 0
187 - for idx, final_res in enumerate(final_res_batch):
188 - classify_res = ScoringRequestOutput.from_base(final_res)
189 -
190 - result = RerankResult(
191 - index=idx,
192 - document=RerankDocument(text=documents[idx]),
193 - relevance_score=classify_res.outputs.score,
194 - )
195 - results.append(result)
196 - prompt_token_ids = final_res.prompt_token_ids
197 - num_prompt_tokens += len(prompt_token_ids)
198 -
199 - # sort by relevance, then return the top n if set
200 - results.sort(key=lambda x: x.relevance_score, reverse=True)
201 - if top_n < len(documents):
202 - results = results[:top_n]
203 -
204 - return RerankResponse(
205 - id=request_id,
206 - model=model_name,
207 - results=results,
208 - usage=RerankUsage(total_tokens=num_prompt_tokens))
vllm/entrypoints/openai/serving_score.py
+329 -134
@@ -1,53 +1,36 @@
1 # SPDX-License-Identifier: Apache-2.0
2 -
2 import asyncio
3 import time
5 -from typing import Any, AsyncGenerator, Dict, List, Optional, Union, cast
4 +from typing import Any, AsyncGenerator, Dict, List, Mapping, Optional, Union
5
6 from fastapi import Request
7
8 from vllm.config import ModelConfig
9 from vllm.engine.protocol import EngineClient
10 from vllm.entrypoints.logger import RequestLogger
12 -from vllm.entrypoints.openai.protocol import (ErrorResponse, ScoreRequest,
13 - ScoreResponse, ScoreResponseData,
14 - UsageInfo)
11 +from vllm.entrypoints.openai.protocol import (ErrorResponse, RerankDocument,
12 + RerankRequest, RerankResponse,
13 + RerankResult, RerankUsage,
14 + ScoreRequest, ScoreResponse,
15 + ScoreResponseData, UsageInfo)
16 from vllm.entrypoints.openai.serving_engine import OpenAIServing
17 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
18 +from vllm.entrypoints.score_utils import (_cosine_similarity,
19 + _validate_score_input_lens)
20 from vllm.inputs.data import TokensPrompt
21 from vllm.logger import init_logger
22 +from vllm.lora.request import LoRARequest
23 from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput
20 -from vllm.transformers_utils.tokenizers.mistral import MistralTokenizer
24 +from vllm.prompt_adapter.request import PromptAdapterRequest
25 +from vllm.transformers_utils.tokenizer import (AnyTokenizer, MistralTokenizer,
26 + PreTrainedTokenizer,
27 + PreTrainedTokenizerFast)
28 from vllm.utils import make_async, merge_async_iterators
29
30 logger = init_logger(__name__)
31
32
26 -def make_pairs(text_1: Union[List[str], str], text_2: Union[List[str],
27 - str]) -> List:
28 - if isinstance(text_1, (str, dict)):
29 - # Convert a single prompt to a list.
30 - text_1 = [text_1]
31 - text_1 = [t for t in text_1]
32 -
33 - if isinstance(text_2, (str, dict)):
34 - # Convert a single prompt to a list.
35 - text_2 = [text_2]
36 - text_2 = [t for t in text_2]
37 - if len(text_1) > 1 and len(text_1) != len(text_2):
38 - raise ValueError("Input lengths must be either 1:1, 1:N or N:N")
39 - if len(text_1) == 0:
40 - raise ValueError("At least one text element must be given")
41 - if len(text_2) == 0:
42 - raise ValueError("At least one text_pair element must be given")
43 -
44 - if len(text_1) == 1:
45 - text_1 = text_1 * len(text_2)
46 -
47 - return [(t1, t2) for t1, t2 in zip(text_1, text_2)]
48 -
49 -
50 -class OpenAIServingScores(OpenAIServing):
33 +class ServingScores(OpenAIServing):
34
35 def __init__(
36 self,
@@ -62,137 +45,280 @@ class OpenAIServingScores(OpenAIServing):
45 models=models,
46 request_logger=request_logger)
47
65 - async def create_score(
48 + async def _embedding_score(
49 self,
67 - request: ScoreRequest,
68 - raw_request: Optional[Request] = None,
69 - ) -> Union[ScoreResponse, ErrorResponse]:
70 - """
71 - Score API similar to Sentence Transformers cross encoder
50 + tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
51 + texts_1: List[str],
52 + texts_2: List[str],
53 + request: Union[RerankRequest, ScoreRequest],
54 + request_id=str,
55 + tokenization_kwargs: Optional[dict[str, Any]] = None,
56 + lora_request: Optional[Union[LoRARequest, None]] = None,
57 + prompt_adapter_request: Optional[Union[PromptAdapterRequest,
58 + None]] = None,
59 + trace_headers: Optional[Mapping[str, str]] = None,
60 + ) -> List[PoolingRequestOutput]:
61 +
62 + input_texts = texts_1 + texts_2
63 +
64 + engine_prompts: List[TokensPrompt] = []
65 + tokenize_async = make_async(tokenizer.__call__,
66 + executor=self._tokenizer_executor)
67 +
68 + tokenization_kwargs = tokenization_kwargs or {}
69 + tokenized_prompts = await asyncio.gather(
70 + *(tokenize_async(t, **tokenization_kwargs) for t in input_texts))
71 +
72 + for tok_result, input_text in zip(tokenized_prompts, input_texts):
73 +
74 + text_token_prompt = \
75 + self._validate_input(
76 + request,
77 + tok_result["input_ids"],
78 + input_text)
79 +
80 + engine_prompts.append(
81 + TokensPrompt(
82 + prompt_token_ids=text_token_prompt["prompt_token_ids"]))
83
73 - See https://sbert.net/docs/package_reference/cross_encoder
74 - """
75 - error_check_ret = await self._check_model(request)
76 - if error_check_ret is not None:
77 - return error_check_ret
84 + # Schedule the request and get the result generator.
85 + generators: List[AsyncGenerator[PoolingRequestOutput, None]] = []
86 + pooling_params = request.to_pooling_params()
87
79 - model_name = request.model
80 - request_id = f"score-{self._base_request_id(raw_request)}"
81 - created_time = int(time.time())
82 - truncate_prompt_tokens = request.truncate_prompt_tokens
88 + for i, engine_prompt in enumerate(engine_prompts):
89
84 - request_prompts = []
85 - engine_prompts = []
90 + request_id_item = f"{request_id}-{i}"
91
87 - try:
88 - (
89 - lora_request,
90 - prompt_adapter_request,
91 - ) = self._maybe_get_adapters(request)
92 + self._log_inputs(request_id_item,
93 + input_texts[i],
94 + params=pooling_params,
95 + lora_request=lora_request,
96 + prompt_adapter_request=prompt_adapter_request)
97
93 - tokenizer = await self.engine_client.get_tokenizer(lora_request)
98 + generators.append(
99 + self.engine_client.encode(
100 + engine_prompt,
101 + pooling_params,
102 + request_id_item,
103 + lora_request=lora_request,
104 + trace_headers=trace_headers,
105 + priority=request.priority,
106 + ))
107
95 - if prompt_adapter_request is not None:
96 - raise NotImplementedError("Prompt adapter is not supported "
97 - "for scoring models")
108 + result_generator = merge_async_iterators(*generators)
109
99 - if isinstance(tokenizer, MistralTokenizer):
100 - raise ValueError(
101 - "MistralTokenizer not supported for cross-encoding")
110 + # Non-streaming response
111 + final_res_batch: List[PoolingRequestOutput] = []
112
103 - if not self.model_config.is_cross_encoder:
104 - raise ValueError("Model is not cross encoder.")
113 + embeddings: List[Optional[PoolingRequestOutput]] =\
114 + [None] * len(engine_prompts)
115
106 - if truncate_prompt_tokens is not None and \
107 - truncate_prompt_tokens > self.max_model_len:
108 - raise ValueError(
109 - f"truncate_prompt_tokens value ({truncate_prompt_tokens}) "
110 - f"is greater than max_model_len ({self.max_model_len})."
111 - f" Please, select a smaller truncation size.")
112 -
113 - input_pairs = make_pairs(request.text_1, request.text_2)
114 - for q, t in input_pairs:
115 - request_prompt = f"{q}{tokenizer.sep_token}{t}"
116 -
117 - tokenization_kwargs: Dict[str, Any] = {}
118 - if truncate_prompt_tokens is not None:
119 - tokenization_kwargs["truncation"] = True
120 - tokenization_kwargs["max_length"] = truncate_prompt_tokens
121 -
122 - tokenize_async = make_async(tokenizer.__call__,
123 - executor=self._tokenizer_executor)
124 - prompt_inputs = await tokenize_async(q,
125 - text_pair=t,
126 - **tokenization_kwargs)
127 -
128 - input_ids = prompt_inputs["input_ids"]
129 - text_token_prompt = \
130 - self._validate_input(request, input_ids, request_prompt)
131 - engine_prompt = TokensPrompt(
132 - prompt_token_ids=text_token_prompt["prompt_token_ids"],
133 - token_type_ids=prompt_inputs.get("token_type_ids"))
134 -
135 - request_prompts.append(request_prompt)
136 - engine_prompts.append(engine_prompt)
116 + async for i, res in result_generator:
117 + embeddings[i] = res
118
138 - except ValueError as e:
139 - logger.exception("Error in preprocessing prompt inputs")
140 - return self.create_error_response(str(e))
119 + emb_texts_1: List[PoolingRequestOutput] = []
120 + emb_texts_2: List[PoolingRequestOutput] = []
121 +
122 + for i in range(0, len(texts_1)):
123 + assert (emb := embeddings[i]) is not None
124 + emb_texts_1.append(emb)
125 +
126 + for i in range(len(texts_1), len(embeddings)):
127 + assert (emb := embeddings[i]) is not None
128 + emb_texts_2.append(emb)
129 +
130 + if len(emb_texts_1) == 1:
131 + emb_texts_1 = emb_texts_1 * len(emb_texts_2)
132 +
133 + final_res_batch = _cosine_similarity(tokenizer=tokenizer,
134 + embed_1=emb_texts_1,
135 + embed_2=emb_texts_2)
136 +
137 + return final_res_batch
138 +
139 + async def _cross_encoding_score(
140 + self,
141 + tokenizer: Union[AnyTokenizer],
142 + texts_1: List[str],
143 + texts_2: List[str],
144 + request: Union[RerankRequest, ScoreRequest],
145 + request_id=str,
146 + tokenization_kwargs: Optional[dict[str, Any]] = None,
147 + lora_request: Optional[Union[LoRARequest, None]] = None,
148 + prompt_adapter_request: Optional[Union[PromptAdapterRequest,
149 + None]] = None,
150 + trace_headers: Optional[Mapping[str, str]] = None,
151 + ) -> List[PoolingRequestOutput]:
152 +
153 + request_prompts: List[str] = []
154 + engine_prompts: List[TokensPrompt] = []
155 +
156 + if len(texts_1) == 1:
157 + texts_1 = texts_1 * len(texts_2)
158 +
159 + input_pairs = [(t1, t2) for t1, t2 in zip(texts_1, texts_2)]
160 +
161 + if isinstance(tokenizer, MistralTokenizer):
162 + raise ValueError(
163 + "MistralTokenizer not supported for cross-encoding")
164 +
165 + tokenize_async = make_async(tokenizer.__call__,
166 + executor=self._tokenizer_executor)
167 +
168 + tokenization_kwargs = tokenization_kwargs or {}
169 + tokenized_prompts = await asyncio.gather(
170 + *(tokenize_async(text=t1, text_pair=t2, **tokenization_kwargs)
171 + for t1, t2 in input_pairs))
172 +
173 + for prompt_inputs, (t1, t2) in zip(tokenized_prompts, input_pairs):
174 +
175 + request_prompt = f"{t1}{tokenizer.sep_token}{t2}"
176 +
177 + input_ids = prompt_inputs["input_ids"]
178 + text_token_prompt = \
179 + self._validate_input(request, input_ids, request_prompt)
180 + engine_prompt = TokensPrompt(
181 + prompt_token_ids=text_token_prompt["prompt_token_ids"],
182 + token_type_ids=prompt_inputs.get("token_type_ids"))
183 +
184 + request_prompts.append(request_prompt)
185 + engine_prompts.append(engine_prompt)
186
187 # Schedule the request and get the result generator.
188 generators: List[AsyncGenerator[PoolingRequestOutput, None]] = []
189
145 - try:
146 - pooling_params = request.to_pooling_params()
147 -
148 - for i, engine_prompt in enumerate(engine_prompts):
149 - request_id_item = f"{request_id}-{i}"
190 + pooling_params = request.to_pooling_params()
191
151 - self._log_inputs(request_id_item,
152 - request_prompts[i],
153 - params=pooling_params,
154 - lora_request=lora_request,
155 - prompt_adapter_request=prompt_adapter_request)
192 + for i, engine_prompt in enumerate(engine_prompts):
193 + request_id_item = f"{request_id}-{i}"
194
157 - trace_headers = (None if raw_request is None else await
158 - self._get_trace_headers(raw_request.headers))
195 + self._log_inputs(request_id_item,
196 + request_prompts[i],
197 + params=pooling_params,
198 + lora_request=lora_request,
199 + prompt_adapter_request=prompt_adapter_request)
200
160 - generator = self.engine_client.encode(
161 - engine_prompt,
162 - pooling_params,
163 - request_id_item,
164 - lora_request=lora_request,
165 - trace_headers=trace_headers,
166 - priority=request.priority,
167 - )
201 + generator = self.engine_client.encode(
202 + engine_prompt,
203 + pooling_params,
204 + request_id_item,
205 + lora_request=lora_request,
206 + trace_headers=trace_headers,
207 + priority=request.priority,
208 + )
209
169 - generators.append(generator)
170 - except ValueError as e:
171 - # TODO: Use a vllm-specific Validation Error
172 - return self.create_error_response(str(e))
210 + generators.append(generator)
211
212 result_generator = merge_async_iterators(*generators)
213
176 - num_prompts = len(engine_prompts)
177 -
214 # Non-streaming response
179 - final_res_batch: List[Optional[PoolingRequestOutput]]
180 - final_res_batch = [None] * num_prompts
215 + final_res_batch: List[
216 + Optional[PoolingRequestOutput]] = [None] * len(engine_prompts)
217
182 - try:
183 - async for i, res in result_generator:
184 - final_res_batch[i] = res
218 + async for i, res in result_generator:
219 + final_res_batch[i] = res
220 +
221 + return [out for out in final_res_batch if out is not None]
222 +
223 + async def _run_scoring(
224 + self,
225 + texts_1: Union[str, list[str]],
226 + texts_2: Union[str, list[str]],
227 + request: Union[ScoreRequest, RerankRequest],
228 + request_id: str,
229 + raw_request: Optional[Request] = None,
230 + truncate_prompt_tokens: Optional[int] = None,
231 + ) -> List[PoolingRequestOutput]:
232 +
233 + tokenization_kwargs: Dict[str, Any] = {}
234 + if truncate_prompt_tokens is not None:
235 + tokenization_kwargs["truncation"] = True
236 + tokenization_kwargs["max_length"] = truncate_prompt_tokens
237 +
238 + (
239 + lora_request,
240 + prompt_adapter_request,
241 + ) = self._maybe_get_adapters(request)
242 +
243 + if prompt_adapter_request is not None:
244 + raise NotImplementedError("Prompt adapter is not supported "
245 + "for scoring models")
246 +
247 + tokenizer = await self.engine_client.get_tokenizer(lora_request)
248 +
249 + if truncate_prompt_tokens is not None and \
250 + truncate_prompt_tokens > self.max_model_len:
251 + raise ValueError(
252 + f"truncate_prompt_tokens value ({truncate_prompt_tokens}) "
253 + f"is greater than max_model_len ({self.max_model_len})."
254 + f" Please, select a smaller truncation size.")
255 +
256 + trace_headers = (None if raw_request is None else await
257 + self._get_trace_headers(raw_request.headers))
258 +
259 + if isinstance(texts_1, str):
260 + texts_1 = [texts_1]
261 + if isinstance(texts_2, str):
262 + texts_2 = [texts_2]
263 +
264 + _validate_score_input_lens(texts_1, texts_2)
265 +
266 + if self.model_config.is_cross_encoder:
267 + return await self._cross_encoding_score(
268 + tokenizer=tokenizer,
269 + texts_1=texts_1,
270 + texts_2=texts_2,
271 + request=request,
272 + request_id=request_id,
273 + tokenization_kwargs=tokenization_kwargs,
274 + lora_request=lora_request,
275 + prompt_adapter_request=prompt_adapter_request,
276 + trace_headers=trace_headers)
277 +
278 + else:
279 + return await self._embedding_score(
280 + tokenizer=tokenizer,
281 + texts_1=texts_1,
282 + texts_2=texts_2,
283 + request=request,
284 + request_id=request_id,
285 + tokenization_kwargs=tokenization_kwargs,
286 + lora_request=lora_request,
287 + prompt_adapter_request=prompt_adapter_request,
288 + trace_headers=trace_headers)
289 +
290 + async def create_score(
291 + self,
292 + request: ScoreRequest,
293 + raw_request: Optional[Request] = None,
294 + ) -> Union[ScoreResponse, ErrorResponse]:
295 + """
296 + Score API similar to Sentence Transformers cross encoder
297 +
298 + See https://sbert.net/docs/package_reference/cross_encoder
299 + """
300 + error_check_ret = await self._check_model(request)
301 + if error_check_ret is not None:
302 + return error_check_ret
303
186 - assert all(final_res is not None for final_res in final_res_batch)
304 + request_id = f"score-{self._base_request_id(raw_request)}"
305 + created_time = int(time.time())
306
188 - final_res_batch_checked = cast(List[PoolingRequestOutput],
189 - final_res_batch)
307 + try:
308 + final_res_batch = await self._run_scoring(
309 + request.text_1,
310 + request.text_2,
311 + request,
312 + request_id,
313 + raw_request,
314 + request.truncate_prompt_tokens,
315 + )
316
191 - response = self.request_output_to_score_response(
192 - final_res_batch_checked,
317 + return self.request_output_to_score_response(
318 + final_res_batch,
319 request_id,
320 created_time,
195 - model_name,
321 + request.model,
322 )
323 except asyncio.CancelledError:
324 return self.create_error_response("Client disconnected")
@@ -200,7 +326,44 @@ class OpenAIServingScores(OpenAIServing):
326 # TODO: Use a vllm-specific Validation Error
327 return self.create_error_response(str(e))
328
203 - return response
329 + async def do_rerank(
330 + self,
331 + request: RerankRequest,
332 + raw_request: Optional[Request] = None
333 + ) -> Union[RerankResponse, ErrorResponse]:
334 + """
335 + Rerank API based on JinaAI's rerank API; implements the same
336 + API interface. Designed for compatibility with off-the-shelf
337 + tooling, since this is a common standard for reranking APIs
338 +
339 + See example client implementations at
340 + https://github.com/infiniflow/ragflow/blob/main/rag/llm/rerank_model.py
341 + numerous clients use this standard.
342 + """
343 + error_check_ret = await self._check_model(request)
344 + if error_check_ret is not None:
345 + return error_check_ret
346 +
347 + request_id = f"rerank-{self._base_request_id(raw_request)}"
348 + documents = request.documents
349 + top_n = request.top_n if request.top_n > 0 else len(documents)
350 +
351 + try:
352 + final_res_batch = await self._run_scoring(
353 + request.query,
354 + documents,
355 + request,
356 + request_id,
357 + raw_request,
358 + request.truncate_prompt_tokens,
359 + )
360 + return self.request_output_to_rerank_response(
361 + final_res_batch, request_id, request.model, documents, top_n)
362 + except asyncio.CancelledError:
363 + return self.create_error_response("Client disconnected")
364 + except ValueError as e:
365 + # TODO: Use a vllm-specific Validation Error
366 + return self.create_error_response(str(e))
367
368 def request_output_to_score_response(
369 self,
@@ -236,3 +399,35 @@ class OpenAIServingScores(OpenAIServing):
399 data=items,
400 usage=usage,
401 )
402 +
403 + def request_output_to_rerank_response(
404 + self, final_res_batch: List[PoolingRequestOutput], request_id: str,
405 + model_name: str, documents: List[str],
406 + top_n: int) -> RerankResponse:
407 + """
408 + Convert the output of do_rank to a RerankResponse
409 + """
410 + results: List[RerankResult] = []
411 + num_prompt_tokens = 0
412 + for idx, final_res in enumerate(final_res_batch):
413 + classify_res = ScoringRequestOutput.from_base(final_res)
414 +
415 + result = RerankResult(
416 + index=idx,
417 + document=RerankDocument(text=documents[idx]),
418 + relevance_score=classify_res.outputs.score,
419 + )
420 + results.append(result)
421 + prompt_token_ids = final_res.prompt_token_ids
422 + num_prompt_tokens += len(prompt_token_ids)
423 +
424 + # sort by relevance, then return the top n if set
425 + results.sort(key=lambda x: x.relevance_score, reverse=True)
426 + if top_n < len(documents):
427 + results = results[:top_n]
428 +
429 + return RerankResponse(
430 + id=request_id,
431 + model=model_name,
432 + results=results,
433 + usage=RerankUsage(total_tokens=num_prompt_tokens))
vllm/entrypoints/score_utils.py new
+49
@@ -0,0 +1,49 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import List, Union
3 +
4 +from torch.nn import CosineSimilarity
5 +
6 +from vllm.outputs import PoolingRequestOutput
7 +from vllm.transformers_utils.tokenizer import (PreTrainedTokenizer,
8 + PreTrainedTokenizerFast)
9 +
10 +
11 +def _cosine_similarity(
12 + tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
13 + embed_1: List[PoolingRequestOutput],
14 + embed_2: List[PoolingRequestOutput],
15 +) -> List[PoolingRequestOutput]:
16 +
17 + scorer = CosineSimilarity(0)
18 + scores: Union[List[PoolingRequestOutput]] = []
19 +
20 + for emb_1, emb_2 in zip(embed_1, embed_2):
21 + pair_score = scorer(emb_1.outputs.data, emb_2.outputs.data)
22 +
23 + padding = []
24 + if (pad_token_id := getattr(tokenizer, "pad_token_id",
25 + None)) is not None:
26 + padding = [pad_token_id]
27 +
28 + tokens = emb_1.prompt_token_ids + padding + emb_2.prompt_token_ids
29 +
30 + scores.append(
31 + PoolingRequestOutput(
32 + request_id=f"{emb_1.request_id}_{emb_2.request_id}",
33 + outputs=pair_score,
34 + prompt_token_ids=tokens,
35 + finished=True))
36 +
37 + return scores
38 +
39 +
40 +def _validate_score_input_lens(
41 + texts_1: Union[List[str], List[dict]],
42 + texts_2: Union[List[str], List[dict]],
43 +):
44 + if len(texts_1) > 1 and len(texts_1) != len(texts_2):
45 + raise ValueError("Input lengths must be either 1:1, 1:N or N:N")
46 + if len(texts_1) == 0:
47 + raise ValueError("At least one text element must be given")
48 + if len(texts_2) == 0:
49 + raise ValueError("At least one text_pair element must be given")