Skip to content

vllm.pooling_params

PoolingTask module-attribute 

PoolingTask = Literal[
    "encode", "embed", "classify", "score"
]

PoolingParams

Bases: Struct

API parameters for pooling models.

Attributes:

Name Type Description
dimensions Optional[int]

Reduce the dimensions of embeddings if model support matryoshka representation.
Source code in vllm/pooling_params.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
class PoolingParams(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
    """API parameters for pooling models.

    Attributes:
        dimensions: Reduce the dimensions of embeddings
                    if model support matryoshka representation.
    """

    dimensions: Optional[int] = None

    output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY

    task: Optional[PoolingTask] = None
    """Internal use only."""

    requires_token_ids: bool = False
    """Internal use only."""

    def clone(self) -> "PoolingParams":
        """Returns a deep copy of the PoolingParams instance."""
        return PoolingParams(
            dimensions=self.dimensions,
            task=self.task,
            requires_token_ids=self.requires_token_ids,
        )

    def verify(self, task: PoolingTask, model_config: "ModelConfig") -> None:
        if self.task is None:
            self.task = task
        elif self.task != task:
            msg = f"You cannot overwrite {self.task=!r} with {task=!r}!"
            raise ValueError(msg)

        # 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

        if self.dimensions is not None:
            if not model_config.is_matryoshka:
                raise ValueError(
                    f'Model "{model_config.served_model_name}" does not '
                    f'support matryoshka representation, '
                    f'changing output dimensions will lead to poor results.')

            mds = model_config.matryoshka_dimensions
            if mds is not None:
                if self.dimensions not in mds:
                    raise ValueError(
                        f'Model "{model_config.served_model_name}" '
                        f'only supports {str(mds)} matryoshka dimensions, '
                        f'use other output dimensions will '
                        f'lead to poor results.')
            elif self.dimensions < 1:
                raise ValueError("Dimensions must be greater than 0")

    def __repr__(self) -> str:
        return (f"PoolingParams("
                f"dimensions={self.dimensions}, "
                f"task={self.task}, "
                f"requires_token_ids={self.requires_token_ids})")

    def __post_init__(self) -> None:
        assert self.output_kind == RequestOutputKind.FINAL_ONLY,\
            "For pooling output_kind has to be FINAL_ONLY"

dimensions class-attribute instance-attribute 

dimensions: Optional[int] = None

output_kind class-attribute instance-attribute 

output_kind: RequestOutputKind = FINAL_ONLY

requires_token_ids class-attribute instance-attribute 

requires_token_ids: bool = False

Internal use only.

task class-attribute instance-attribute 

task: Optional[PoolingTask] = None

Internal use only.

__post_init__ 

__post_init__() -> None
Source code in vllm/pooling_params.py 
80
81
82
def __post_init__(self) -> None:
    assert self.output_kind == RequestOutputKind.FINAL_ONLY,\
        "For pooling output_kind has to be FINAL_ONLY"

__repr__ 

__repr__() -> str
Source code in vllm/pooling_params.py 
74
75
76
77
78
def __repr__(self) -> str:
    return (f"PoolingParams("
            f"dimensions={self.dimensions}, "
            f"task={self.task}, "
            f"requires_token_ids={self.requires_token_ids})")

clone 

clone() -> PoolingParams

Returns a deep copy of the PoolingParams instance.

Source code in vllm/pooling_params.py 
37
38
39
40
41
42
43
def clone(self) -> "PoolingParams":
    """Returns a deep copy of the PoolingParams instance."""
    return PoolingParams(
        dimensions=self.dimensions,
        task=self.task,
        requires_token_ids=self.requires_token_ids,
    )

verify 

verify(
    task: PoolingTask, model_config: ModelConfig
) -> None
Source code in vllm/pooling_params.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
def verify(self, task: PoolingTask, model_config: "ModelConfig") -> None:
    if self.task is None:
        self.task = task
    elif self.task != task:
        msg = f"You cannot overwrite {self.task=!r} with {task=!r}!"
        raise ValueError(msg)

    # 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

    if self.dimensions is not None:
        if not model_config.is_matryoshka:
            raise ValueError(
                f'Model "{model_config.served_model_name}" does not '
                f'support matryoshka representation, '
                f'changing output dimensions will lead to poor results.')

        mds = model_config.matryoshka_dimensions
        if mds is not None:
            if self.dimensions not in mds:
                raise ValueError(
                    f'Model "{model_config.served_model_name}" '
                    f'only supports {str(mds)} matryoshka dimensions, '
                    f'use other output dimensions will '
                    f'lead to poor results.')
        elif self.dimensions < 1:
            raise ValueError("Dimensions must be greater than 0")