wehub-resource-sync
288 lines
9.8 KiB
Python
288 lines
9.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from fastapi.responses import JSONResponse, Response
|
|
|
|
from vllm import PoolingParams
|
|
from vllm.engine.protocol import EngineClient
|
|
from vllm.entrypoints.openai.engine.protocol import UsageInfo
|
|
from vllm.logger import init_logger
|
|
from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput
|
|
from vllm.tasks import SCORE_TYPE_MAP, SupportedTask
|
|
from vllm.v1.pool.late_interaction import (
|
|
build_late_interaction_doc_params,
|
|
build_late_interaction_query_params,
|
|
)
|
|
|
|
from ..base.io_processor import PoolingIOProcessor
|
|
from ..base.serving import PoolingServing
|
|
from .io_processor import ScoringIOProcessors, ScoringServeContext
|
|
from .protocol import (
|
|
RerankDocument,
|
|
RerankRequest,
|
|
RerankResponse,
|
|
RerankResult,
|
|
RerankUsage,
|
|
ScoreRequest,
|
|
ScoreResponse,
|
|
ScoreResponseData,
|
|
)
|
|
from .typing import ScoreInput
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class ServingScores(PoolingServing):
|
|
request_id_prefix = "score"
|
|
|
|
def __init__(
|
|
self,
|
|
engine_client: EngineClient,
|
|
*args,
|
|
supported_tasks: tuple[SupportedTask, ...],
|
|
enable_flash_late_interaction: bool = True,
|
|
**kwargs,
|
|
):
|
|
pooling_task = engine_client.model_config.get_pooling_task(supported_tasks)
|
|
score_type = SCORE_TYPE_MAP.get(pooling_task, None) # type: ignore[arg-type]
|
|
assert score_type is not None
|
|
|
|
self.io_processor_name: str = score_type
|
|
self.enable_flash_late_interaction = (
|
|
self.io_processor_name == "late-interaction"
|
|
and enable_flash_late_interaction
|
|
)
|
|
|
|
if self.enable_flash_late_interaction:
|
|
self.io_processor_name = "flash-late-interaction"
|
|
|
|
if engine_client.model_config.architecture == "JinaForRanking":
|
|
self.io_processor_name = "jina-reranking-scoring"
|
|
self.enable_flash_late_interaction = False
|
|
|
|
super().__init__(engine_client, *args, **kwargs)
|
|
|
|
def init_io_processor(self, *args, **kwargs) -> PoolingIOProcessor:
|
|
return ScoringIOProcessors[self.io_processor_name](*args, **kwargs)
|
|
|
|
async def __call__(self, *args, **kwargs) -> Response:
|
|
if not self.enable_flash_late_interaction:
|
|
return await super().__call__(*args, **kwargs)
|
|
|
|
return await self.flash_late_interaction(*args, **kwargs)
|
|
|
|
def _build_response(
|
|
self,
|
|
ctx: ScoringServeContext,
|
|
) -> JSONResponse:
|
|
final_res_batch = ctx.final_res_batch
|
|
request_id = ctx.request_id
|
|
created_time = ctx.created_time
|
|
model_name = self.models.model_name()
|
|
|
|
if isinstance(ctx.request, ScoreRequest):
|
|
return self._request_output_to_score_response(
|
|
final_res_batch,
|
|
request_id,
|
|
created_time,
|
|
model_name,
|
|
)
|
|
elif isinstance(ctx.request, RerankRequest):
|
|
return self._request_output_to_rerank_response(
|
|
final_res_batch,
|
|
request_id,
|
|
model_name,
|
|
ctx.request.documents,
|
|
ctx.request.top_n if ctx.request.top_n > 0 else len(final_res_batch),
|
|
)
|
|
else:
|
|
raise ValueError(f"Invalid {self.request_id_prefix} request type")
|
|
|
|
def _request_output_to_score_response(
|
|
self,
|
|
final_res_batch: list[PoolingRequestOutput],
|
|
request_id: str,
|
|
created_time: int,
|
|
model_name: str,
|
|
) -> JSONResponse:
|
|
items: list[ScoreResponseData] = []
|
|
num_prompt_tokens = 0
|
|
|
|
for idx, final_res in enumerate(final_res_batch):
|
|
classify_res = ScoringRequestOutput.from_base(final_res)
|
|
|
|
item = ScoreResponseData(
|
|
index=idx,
|
|
score=classify_res.outputs.score,
|
|
)
|
|
prompt_token_ids = final_res.prompt_token_ids
|
|
|
|
items.append(item)
|
|
num_prompt_tokens += len(prompt_token_ids)
|
|
|
|
usage = UsageInfo(
|
|
prompt_tokens=num_prompt_tokens,
|
|
total_tokens=num_prompt_tokens,
|
|
)
|
|
|
|
response = ScoreResponse(
|
|
id=request_id,
|
|
created=created_time,
|
|
model=model_name,
|
|
data=items,
|
|
usage=usage,
|
|
)
|
|
|
|
return JSONResponse(content=response.model_dump())
|
|
|
|
def _request_output_to_rerank_response(
|
|
self,
|
|
final_res_batch: list[PoolingRequestOutput],
|
|
request_id: str,
|
|
model_name: str,
|
|
documents: ScoreInput | list[ScoreInput],
|
|
top_n: int,
|
|
) -> JSONResponse:
|
|
if not isinstance(documents, list):
|
|
documents = [documents]
|
|
|
|
results: list[RerankResult] = []
|
|
num_prompt_tokens = 0
|
|
for idx, final_res in enumerate(final_res_batch):
|
|
classify_res = ScoringRequestOutput.from_base(final_res)
|
|
|
|
document = documents[idx]
|
|
if isinstance(document, str):
|
|
rerank_document = RerankDocument(text=document)
|
|
else:
|
|
rerank_document = RerankDocument(
|
|
multi_modal=document.get("content", [])
|
|
)
|
|
|
|
result = RerankResult(
|
|
index=idx,
|
|
document=rerank_document,
|
|
relevance_score=classify_res.outputs.score,
|
|
)
|
|
results.append(result)
|
|
prompt_token_ids = final_res.prompt_token_ids
|
|
num_prompt_tokens += len(prompt_token_ids)
|
|
|
|
# sort by relevance, then return the top n if set
|
|
results.sort(key=lambda x: x.relevance_score, reverse=True)
|
|
if top_n < len(documents):
|
|
results = results[:top_n]
|
|
|
|
response = RerankResponse(
|
|
id=request_id,
|
|
model=model_name,
|
|
results=results,
|
|
usage=RerankUsage(
|
|
total_tokens=num_prompt_tokens, prompt_tokens=num_prompt_tokens
|
|
),
|
|
)
|
|
|
|
return JSONResponse(content=response.model_dump())
|
|
|
|
###################################################################################
|
|
### Run pooling score MaxSim on worker side (GPU) in the API server process
|
|
### Can significantly improve late-interaction scoring performance.
|
|
|
|
async def flash_late_interaction(self, *args, **kwargs) -> Response:
|
|
ctx = await self._init_ctx(self.io_processor, *args, **kwargs)
|
|
await self._preprocessing_async(self.io_processor, ctx)
|
|
|
|
# stage 1: encode queries and cache token embeddings on workers.
|
|
await self._flash_late_interaction_encode_queries(ctx)
|
|
# stage 2: encode docs and return scalar scores from workers.
|
|
await self._flash_late_interaction_encode_docs(ctx)
|
|
|
|
return await self._postprocessing_async(self.io_processor, ctx)
|
|
|
|
async def _flash_late_interaction_encode_queries(self, ctx: ScoringServeContext):
|
|
assert ctx.n_queries is not None
|
|
assert ctx.engine_inputs is not None
|
|
assert isinstance(ctx.pooling_params, PoolingParams)
|
|
|
|
n_queries = ctx.n_queries
|
|
n_docs = len(ctx.engine_inputs) - n_queries
|
|
query_engine_inputs = ctx.engine_inputs[:n_queries]
|
|
|
|
query_keys = [f"{ctx.request_id}-query-{i}" for i in range(n_queries)]
|
|
query_uses = [n_docs if n_queries == 1 else 1] * n_queries
|
|
|
|
query_pooling_params_list = []
|
|
for i in range(n_queries):
|
|
pooling_params = ctx.pooling_params.clone()
|
|
pooling_params.late_interaction_params = (
|
|
build_late_interaction_query_params(
|
|
query_key=query_keys[i],
|
|
query_uses=query_uses[i],
|
|
)
|
|
)
|
|
query_pooling_params_list.append(pooling_params)
|
|
|
|
assert (
|
|
n_queries
|
|
== len(query_pooling_params_list)
|
|
== len(query_engine_inputs)
|
|
== len(query_keys)
|
|
)
|
|
|
|
query_ctx = ScoringServeContext(
|
|
request=ctx.request,
|
|
raw_request=ctx.raw_request,
|
|
model_name=ctx.model_name,
|
|
request_id=ctx.request_id,
|
|
pooling_params=query_pooling_params_list,
|
|
prompt_request_ids=query_keys,
|
|
engine_inputs=query_engine_inputs,
|
|
)
|
|
|
|
await self._prepare_generators(query_ctx)
|
|
await self._collect_batch(query_ctx)
|
|
ctx.query_final_res_batch = query_ctx.final_res_batch
|
|
|
|
async def _flash_late_interaction_encode_docs(self, ctx: ScoringServeContext):
|
|
assert ctx.n_queries is not None
|
|
assert ctx.engine_inputs is not None
|
|
assert isinstance(ctx.pooling_params, PoolingParams)
|
|
|
|
n_queries = ctx.n_queries
|
|
n_docs = len(ctx.engine_inputs) - n_queries
|
|
doc_engine_inputs = ctx.engine_inputs[n_queries:]
|
|
|
|
query_keys = [f"{ctx.request_id}-query-{i}" for i in range(n_queries)]
|
|
doc_keys = [f"{ctx.request_id}-doc-{i}" for i in range(n_docs)]
|
|
|
|
doc_pooling_params_list = []
|
|
for i in range(n_docs):
|
|
query_idx = 0 if n_queries == 1 else i
|
|
pooling_params = ctx.pooling_params.clone()
|
|
pooling_params.late_interaction_params = build_late_interaction_doc_params(
|
|
query_key=query_keys[query_idx]
|
|
)
|
|
doc_pooling_params_list.append(pooling_params)
|
|
|
|
assert (
|
|
n_docs
|
|
== len(doc_pooling_params_list)
|
|
== len(doc_engine_inputs)
|
|
== len(doc_keys)
|
|
)
|
|
|
|
doc_ctx = ScoringServeContext(
|
|
request=ctx.request,
|
|
raw_request=ctx.raw_request,
|
|
model_name=ctx.model_name,
|
|
request_id=ctx.request_id,
|
|
pooling_params=doc_pooling_params_list,
|
|
prompt_request_ids=doc_keys,
|
|
engine_inputs=doc_engine_inputs,
|
|
)
|
|
|
|
await self._prepare_generators(doc_ctx)
|
|
await self._collect_batch(doc_ctx)
|
|
|
|
ctx.final_res_batch = doc_ctx.final_res_batch
|