| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import TypeAlias |
| 5 | |
| 6 | import numpy as np |
| 7 | from fastapi.responses import JSONResponse |
| 8 | |
| 9 | from vllm.entrypoints.serve.engine.protocol import UsageInfo |
| 10 | from vllm.logger import init_logger |
| 11 | from vllm.outputs import ClassificationOutput |
| 12 | |
| 13 | from ..base.serving import PoolingServing |
| 14 | from ..typing import PoolingServeContext |
| 15 | from .io_processor import ClassifyIOProcessor |
| 16 | from .protocol import ( |
| 17 | ClassificationData, |
| 18 | ClassificationRequest, |
| 19 | ClassificationResponse, |
| 20 | ) |
| 21 | |
| 22 | logger = init_logger(__name__) |
| 23 | |
| 24 | |
| 25 | ClassificationServeContext: TypeAlias = PoolingServeContext[ClassificationRequest] |
| 26 | |
| 27 | |
| 28 | class ServingClassification(PoolingServing): |
| 29 | request_id_prefix = "classify" |
| 30 | |
| 31 | def init_io_processor(self, *args, **kwargs) -> ClassifyIOProcessor: |
| 32 | return ClassifyIOProcessor(*args, **kwargs) |
| 33 | |
| 34 | def _build_response( |
| 35 | self, |
| 36 | ctx: ClassificationServeContext, |
| 37 | ) -> JSONResponse: |
| 38 | id2label = getattr(self.model_config.hf_config, "id2label", {}) |
| 39 | num_prompt_tokens = 0 |
| 40 | items: list[ClassificationData] = [] |
| 41 | for idx, final_res in enumerate(ctx.final_res_batch): |
| 42 | classify_res = ClassificationOutput.from_base(final_res.outputs) |
| 43 | |
| 44 | probs = classify_res.probs |
| 45 | predicted_index = int(np.argmax(probs)) |
| 46 | label = id2label.get(predicted_index) |
| 47 | |
| 48 | item = ClassificationData( |
| 49 | index=idx, |
| 50 | label=label, |
| 51 | probs=probs, |
| 52 | num_classes=len(probs), |
| 53 | ) |
| 54 | |
| 55 | items.append(item) |
| 56 | prompt_token_ids = final_res.prompt_token_ids |
| 57 | num_prompt_tokens += len(prompt_token_ids) |
| 58 | |
| 59 | usage = UsageInfo( |
| 60 | prompt_tokens=num_prompt_tokens, |
| 61 | total_tokens=num_prompt_tokens, |
| 62 | ) |
| 63 | |
| 64 | response = ClassificationResponse( |
| 65 | id=ctx.request_id, |
| 66 | created=ctx.created_time, |
| 67 | model=ctx.model_name, |
| 68 | data=items, |
| 69 | usage=usage, |
| 70 | ) |
| 71 | |
| 72 | return JSONResponse(content=response.model_dump()) |