class PoolingParams(
msgspec.Struct,
omit_defaults=True, # type: ignore[call-arg]
array_like=True,
): # type: ignore[call-arg]
"""API parameters for pooling models.
Attributes:
use_activation: Whether to apply activation function to the pooler outputs.
`None` uses the pooler's default, which is `True` in most cases.
dimensions: Reduce the dimensions of embeddings
if model support matryoshka representation.
late_chunking_params: Mean-pool contextual token states into fixed-size
chunks for the token embedding task.
"""
# --8<-- [start:common-pooling-params]
use_activation: bool | None = None
# --8<-- [end:common-pooling-params]
## for embeddings models
# --8<-- [start:embed-pooling-params]
dimensions: int | None = None
# --8<-- [end:embed-pooling-params]
## for step pooling models
step_tag_id: int | None = None
returned_token_ids: list[int] | None = None
## Internal use only
task: PoolingTask | None = None
requires_token_ids: bool = False
skip_reading_prefix_cache: bool | None = None
late_interaction_params: LateInteractionParams | None = None
extra_kwargs: dict[str, Any] | None = None
output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY
## for token embedding models
# --8<-- [start:token-embed-pooling-params]
late_chunking_params: LateChunkingParams | None = None
"""Mean-pool contextual tokens into fixed-size chunks for `token_embed`."""
# --8<-- [end:token-embed-pooling-params]
@property
def all_parameters(self) -> list[str]:
return ["dimensions", "use_activation", "late_chunking_params"]
@property
def valid_parameters(self):
return {
"embed": ["dimensions", "use_activation"],
"classify": ["use_activation"],
"token_embed": ["dimensions", "use_activation", "late_chunking_params"],
"token_classify": ["use_activation"],
}
def clone(self) -> "PoolingParams":
"""Returns a deep copy of the PoolingParams instance."""
return deepcopy(self)
def verify(self, model_config: ModelConfig) -> None:
if self.late_chunking_params is not None:
self.late_chunking_params.verify()
if self.task != "token_embed":
raise VLLMValidationError("Late chunking requires token_embed")
if self.late_interaction_params is not None:
raise VLLMValidationError(
"Late chunking cannot be combined with late-interaction scoring"
)
# plugin task uses io_processor.parse_data to verify inputs,
# skipping PoolingParams verify
if self.task == "plugin":
if self.skip_reading_prefix_cache is None:
self.skip_reading_prefix_cache = True
return
# skipping verify, let plugins configure and validate pooling params
if self.task not in self.valid_parameters:
return
# NOTE: Task validation needs to done against the model instance,
# which is not available in model config. So, it's not included
# in this method
self._merge_default_parameters(model_config)
self._set_default_parameters(model_config)
self._verify_valid_parameters()
if self.late_chunking_params is not None:
if self.step_tag_id is not None:
raise VLLMValidationError("Late chunking requires all token rows")
# KV cache hits omit hidden states needed for complete chunk means.
self.skip_reading_prefix_cache = True
def _merge_default_parameters(self, model_config: ModelConfig) -> None:
pooler_config = model_config.pooler_config
if pooler_config is None:
return
if self.task is None:
raise ValueError("task must be set before merging parameters")
valid_parameters = self.valid_parameters[self.task]
for k in valid_parameters:
if getattr(pooler_config, k, None) is None:
continue
if getattr(self, k, None) is None:
setattr(self, k, getattr(pooler_config, k))
if self.skip_reading_prefix_cache is None:
# If prefix caching is enabled,
# the output of all pooling may less than n_prompt_tokens,
# we need to skip reading cache at this request.
if self.task in ["token_embed", "token_classify"]:
self.skip_reading_prefix_cache = True
else:
self.skip_reading_prefix_cache = False
self._verify_step_pooling(pooler_config, valid_parameters)
def _verify_step_pooling(
self,
pooler_config: PoolerConfig,
valid_parameters: list[str],
):
step_pooling_parameters = ["step_tag_id", "returned_token_ids"]
if pooler_config.tok_pooling_type != "STEP":
invalid_parameters = []
for k in step_pooling_parameters:
if getattr(self, k, None) is not None:
invalid_parameters.append(k)
if invalid_parameters:
raise VLLMValidationError(
f"Task {self.task} only supports {valid_parameters} "
f"parameters, does not support "
f"{invalid_parameters} parameters"
)
else:
for k in step_pooling_parameters:
if getattr(pooler_config, k, None) is None:
continue
if getattr(self, k, None) is None:
setattr(self, k, getattr(pooler_config, k))
def _set_default_parameters(self, model_config: ModelConfig):
if self.task in ["embed", "token_embed"]:
if self.use_activation is None:
self.use_activation = True
if self.dimensions is not None:
dimensions = self.dimensions
model_name = model_config.served_model_name
embedding_size = model_config.embedding_size
valid_range = f"[1, {embedding_size}]"
dimensions_in_range = 1 <= dimensions <= embedding_size
if not model_config.is_matryoshka:
raise VLLMValidationError(
f"Model {model_name!r} does not support Matryoshka "
f"embeddings; dimensions must be unset "
f"(received dimensions={dimensions})."
)
if not dimensions_in_range:
raise VLLMValidationError(
f"Model {model_name!r} only supports dimensions in "
f"range {valid_range}, got {dimensions}."
)
mds = model_config.matryoshka_dimensions
if mds is not None and dimensions not in mds:
raise VLLMValidationError(
f"Model {model_name!r} only supports Matryoshka "
f"dimensions {str(mds)}, got {dimensions}."
)
elif self.task in ["classify", "token_classify"]:
if self.use_activation is None:
self.use_activation = True
else:
raise ValueError(f"Unknown pooling task: {self.task!r}")
def _verify_valid_parameters(self):
if self.task is None:
raise ValueError("task must be set before verifying parameters")
valid_parameters = self.valid_parameters[self.task]
invalid_parameters = []
for k in self.all_parameters:
if k in valid_parameters:
continue
if getattr(self, k, None) is not None:
invalid_parameters.append(k)
if invalid_parameters:
raise VLLMValidationError(
f"Task {self.task!r} only supports {valid_parameters} "
f"parameters, does not support "
f"{invalid_parameters} parameters"
)
def __repr__(self) -> str:
return (
f"PoolingParams("
f"task={self.task}, "
f"dimensions={self.dimensions}, "
f"use_activation={self.use_activation}, "
f"step_tag_id={self.step_tag_id}, "
f"returned_token_ids={self.returned_token_ids}, "
f"requires_token_ids={self.requires_token_ids}, "
f"skip_reading_prefix_cache={self.skip_reading_prefix_cache}, "
f"late_interaction_params={self.late_interaction_params}, "
f"late_chunking_params={self.late_chunking_params}, "
f"extra_kwargs={self.extra_kwargs})"
)
def __post_init__(self) -> None:
check_removed_pooling_task(self.task)
if self.output_kind != RequestOutputKind.FINAL_ONLY:
raise VLLMValidationError(
"For pooling output_kind has to be FINAL_ONLY, "
f"got {self.output_kind!r}"
)