diff --git a/examples/pooling/embed/vision_embedding_offline.py b/examples/pooling/embed/vision_embedding_offline.py index 3190e00c47..2c82d0374d 100644 --- a/examples/pooling/embed/vision_embedding_offline.py +++ b/examples/pooling/embed/vision_embedding_offline.py @@ -154,7 +154,13 @@ def run_siglip(seed: int): ) print("Text embedding output:") - outputs = llm.embed(text, use_tqdm=False) + # SigLIP is trained with padding to a fixed length and no attention mask, so + # text embeddings only line up with image embeddings when they are padded. + outputs = llm.embed( + text, + tokenization_kwargs={"padding": "max_length"}, + use_tqdm=False, + ) print_embeddings(outputs[0].outputs.embedding) print("Image embedding output:") diff --git a/examples/pooling/embed/vision_embedding_online.py b/examples/pooling/embed/vision_embedding_online.py index fb9e09ead4..79d6e397c8 100644 --- a/examples/pooling/embed/vision_embedding_online.py +++ b/examples/pooling/embed/vision_embedding_online.py @@ -35,6 +35,7 @@ def create_chat_embeddings( encoding_format: Literal["base64", "float"] | NotGiven = NOT_GIVEN, continue_final_message: bool = False, add_special_tokens: bool = False, + padding: str | NotGiven = NOT_GIVEN, ) -> CreateEmbeddingResponse: """ Convenience function for accessing vLLM's Chat Embeddings API, @@ -49,6 +50,7 @@ def create_chat_embeddings( "encoding_format": encoding_format, "continue_final_message": continue_final_message, "add_special_tokens": add_special_tokens, + "padding": padding, }, ) @@ -288,6 +290,8 @@ def run_siglip(client: OpenAI, model: str): print("Image embedding output:", response.data[0].embedding) + # SigLIP is trained with padding to a fixed length and no attention mask, so + # text embeddings only line up with image embeddings when they are padded. response = create_chat_embeddings( client, messages=[ @@ -300,6 +304,7 @@ def run_siglip(client: OpenAI, model: str): ], model=model, encoding_format="float", + padding="max_length", ) print("Text embedding output:", response.data[0].embedding) diff --git a/tests/entrypoints/pooling/embed/test_online.py b/tests/entrypoints/pooling/embed/test_online.py index 10c9bfd3f5..6cc0320d0b 100644 --- a/tests/entrypoints/pooling/embed/test_online.py +++ b/tests/entrypoints/pooling/embed/test_online.py @@ -295,6 +295,37 @@ async def test_truncate_prompt_tokens(client: openai.AsyncOpenAI, model_name: st ) +@pytest.mark.asyncio +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +async def test_padding(client: openai.AsyncOpenAI, model_name: str): + input_texts = ["Hello"] + + async def prompt_tokens(**extra_body): + response = await client.embeddings.create( + model=model_name, + input=input_texts, + extra_body=extra_body, + ) + embeddings = EmbeddingResponse.model_validate(response.model_dump(mode="json")) + assert len(embeddings.data) == 1 + return embeddings.usage.prompt_tokens + + unpadded = await prompt_tokens() + padded = await prompt_tokens(padding="max_length") + + # Models trained with a fixed sequence length and no attention mask need + # this; without it their embeddings are not comparable. + assert padded > unpadded + assert await prompt_tokens(padding="do_not_pad") == unpadded + + with pytest.raises(openai.BadRequestError): + await client.embeddings.create( + model=model_name, + input=input_texts, + extra_body={"padding": "longest"}, + ) + + @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) async def test_chat_request( diff --git a/vllm/entrypoints/pooling/base/protocol.py b/vllm/entrypoints/pooling/base/protocol.py index 6e592e30f8..b4f0ff21b8 100644 --- a/vllm/entrypoints/pooling/base/protocol.py +++ b/vllm/entrypoints/pooling/base/protocol.py @@ -39,6 +39,17 @@ class PoolingBasicRequestMixin(OpenAIBaseModel): # --8<-- [start:pooling-common-extra-params] truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None + padding: Literal["max_length", "do_not_pad"] | None = Field( + default=None, + description=( + "Whether to pad the prompt, using the same names as the " + "Transformers tokenizer. 'max_length' pads to the maximum input " + "length, 'do_not_pad' leaves the prompt as is. Models trained " + "with a fixed sequence length and no attention mask, such as " + "SigLIP, need 'max_length' to match training, otherwise their " + "embeddings are not comparable." + ), + ) truncation_side: Literal["left", "right"] | None = Field( default=None, description=( @@ -94,7 +105,7 @@ class PoolingBasicRequestMixin(OpenAIBaseModel): ) -> TokenizeParams: encoder_config = model_config.encoder_config or {} if max_output_tokens_param is None: - return TokenizeParams( + tok_params = TokenizeParams( max_total_tokens=max_total_tokens, max_output_tokens=max_output_tokens, truncate_prompt_tokens=self.truncate_prompt_tokens, @@ -103,17 +114,22 @@ class PoolingBasicRequestMixin(OpenAIBaseModel): add_special_tokens=add_special_tokens, max_total_tokens_param=max_total_tokens_param, ) + else: + tok_params = TokenizeParams( + max_total_tokens=max_total_tokens, + max_output_tokens=max_output_tokens, + truncate_prompt_tokens=self.truncate_prompt_tokens, + truncation_side=self.truncation_side, + do_lower_case=encoder_config.get("do_lower_case", False), + add_special_tokens=add_special_tokens, + max_total_tokens_param=max_total_tokens_param, + max_output_tokens_param=max_output_tokens_param, + ) - return TokenizeParams( - max_total_tokens=max_total_tokens, - max_output_tokens=max_output_tokens, - truncate_prompt_tokens=self.truncate_prompt_tokens, - truncation_side=self.truncation_side, - do_lower_case=encoder_config.get("do_lower_case", False), - add_special_tokens=add_special_tokens, - max_total_tokens_param=max_total_tokens_param, - max_output_tokens_param=max_output_tokens_param, - ) + if self.padding is None: + return tok_params + + return tok_params.with_kwargs(padding=self.padding) class PoolingTokenizeParamsMixin: