@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 187 lines 5.24 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import time
4 from typing import Any, TypeAlias
5
6 from pydantic import BaseModel, Field, model_validator
7
8 from vllm import PoolingParams
9 from vllm.config import ModelConfig
10 from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel, UsageInfo
11 from vllm.renderers import TokenizeParams
12 from vllm.tasks import PoolingTask
13 from vllm.utils import random_uuid
14
15 from ..base.protocol import ClassifyRequestMixin, PoolingBasicRequestMixin
16 from .typing import ScoreContentPartParam, ScoreInput
17
18
19 class ScoringRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin):
20 # --8<-- [start:scoring-common-params]
21 max_tokens_per_query: int = Field(
22 default=0,
23 description=(
24 "Maximum number of tokens per query. Queries longer than "
25 "this will be truncated to this length. 0 means no "
26 "query-level truncation is applied."
27 ),
28 )
29 max_tokens_per_doc: int = Field(
30 default=0,
31 description=(
32 "Maximum number of tokens per document. Documents longer than "
33 "this will be truncated to this length. 0 means no "
34 "document-level truncation is applied (only truncate_prompt_tokens "
35 "applies to the combined query+document)."
36 ),
37 )
38 instruction: str | None = Field(
39 default=None,
40 description=(
41 "Task instruction prepended to each scored pair via the chat "
42 "template. Equivalent to passing "
43 "chat_template_kwargs={'instruction': ...}."
44 ),
45 )
46 chat_template_kwargs: dict[str, Any] | None = Field(
47 default=None,
48 description=(
49 "Additional keyword args to pass to the chat template renderer. "
50 "Will be accessible by the score/rerank chat template."
51 ),
52 )
53 # --8<-- [end:scoring-common-params]
54
55 @model_validator(mode="after")
56 def _merge_instruction_into_kwargs(self) -> "ScoringRequestMixin":
57 """Fold the top-level `instruction` field into `chat_template_kwargs`.
58
59 This allows callers to use either the convenience field or the generic
60 dict. Explicit keys inside `chat_template_kwargs` take precedence over
61 the top-level `instruction` field.
62 """
63 if self.instruction is not None:
64 merged = dict(self.chat_template_kwargs or {})
65 merged.setdefault("instruction", self.instruction)
66 self.chat_template_kwargs = merged
67 return self
68
69 def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
70 return self._build_pooling_tok_params(
71 model_config,
72 add_special_tokens=True,
73 max_total_tokens=model_config.max_model_len,
74 max_output_tokens=0,
75 )
76
77 def to_pooling_params(self, task: PoolingTask = "classify"):
78 return PoolingParams(
79 task=task,
80 use_activation=self.use_activation,
81 )
82
83
84 class ScoreDataRequest(ScoringRequestMixin):
85 data_1: ScoreInput | list[ScoreInput]
86 data_2: ScoreInput | list[ScoreInput]
87
88
89 class ScoreQueriesDocumentsRequest(ScoringRequestMixin):
90 # --8<-- [start:score-request-params]
91 queries: ScoreInput | list[ScoreInput]
92 documents: ScoreInput | list[ScoreInput]
93 # --8<-- [end:score-request-params]
94
95 @property
96 def data_1(self):
97 return self.queries
98
99 @property
100 def data_2(self):
101 return self.documents
102
103
104 class ScoreQueriesItemsRequest(ScoringRequestMixin):
105 queries: ScoreInput | list[ScoreInput]
106 items: ScoreInput | list[ScoreInput]
107
108 @property
109 def data_1(self):
110 return self.queries
111
112 @property
113 def data_2(self):
114 return self.items
115
116
117 class ScoreTextRequest(ScoringRequestMixin):
118 text_1: ScoreInput | list[ScoreInput]
119 text_2: ScoreInput | list[ScoreInput]
120
121 @property
122 def data_1(self):
123 return self.text_1
124
125 @property
126 def data_2(self):
127 return self.text_2
128
129
130 ScoreRequest: TypeAlias = (
131 ScoreQueriesDocumentsRequest
132 | ScoreQueriesItemsRequest
133 | ScoreDataRequest
134 | ScoreTextRequest
135 )
136
137
138 class RerankRequest(ScoringRequestMixin):
139 # --8<-- [start:rerank-request-params]
140 query: ScoreInput
141 documents: ScoreInput | list[ScoreInput]
142 top_n: int = Field(default=0, ge=0)
143 # --8<-- [end:rerank-request-params]
144
145
146 ScoringRequest: TypeAlias = ScoreRequest | RerankRequest
147
148
149 class RerankDocument(BaseModel):
150 text: str | None = None
151 multi_modal: list[ScoreContentPartParam] | None = None
152
153
154 class RerankResult(BaseModel):
155 index: int
156 document: RerankDocument
157 relevance_score: float
158
159
160 class RerankUsage(BaseModel):
161 prompt_tokens: int
162 total_tokens: int
163
164
165 class RerankResponse(OpenAIBaseModel):
166 id: str
167 model: str
168 usage: RerankUsage
169 results: list[RerankResult]
170
171
172 class ScoreResponseData(OpenAIBaseModel):
173 index: int
174 object: str = "score"
175 score: float
176
177
178 class ScoreResponse(OpenAIBaseModel):
179 id: str = Field(default_factory=lambda: f"embd-{random_uuid()}")
180 object: str = "list"
181 created: int = Field(default_factory=lambda: int(time.time()))
182 model: str
183 data: list[ScoreResponseData]
184 usage: UsageInfo
185
186
187 ScoringResponse: TypeAlias = RerankResponse | ScoreResponse