Convert formatting to use ruff instead of yapf + isort (#26247)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
This commit is contained in:
@@ -10,17 +10,18 @@ import openai
|
||||
import pytest
|
||||
|
||||
from tests.conftest import HfRunner
|
||||
from tests.models.language.pooling.embed_utils import (
|
||||
run_embedding_correctness_test)
|
||||
from tests.models.language.pooling.embed_utils import run_embedding_correctness_test
|
||||
from tests.models.utils import EmbedModelInfo
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.entrypoints.openai.protocol import EmbeddingResponse
|
||||
|
||||
MODELS = [
|
||||
EmbedModelInfo("intfloat/multilingual-e5-small", is_matryoshka=False),
|
||||
EmbedModelInfo("Snowflake/snowflake-arctic-embed-m-v1.5",
|
||||
is_matryoshka=True,
|
||||
matryoshka_dimensions=[256]),
|
||||
EmbedModelInfo(
|
||||
"Snowflake/snowflake-arctic-embed-m-v1.5",
|
||||
is_matryoshka=True,
|
||||
matryoshka_dimensions=[256],
|
||||
),
|
||||
]
|
||||
|
||||
input_texts = [
|
||||
@@ -48,15 +49,14 @@ def server(model_info, dtype: str):
|
||||
dtype,
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"512"
|
||||
"512",
|
||||
]
|
||||
|
||||
if model_info.name == "Snowflake/snowflake-arctic-embed-m-v1.5":
|
||||
# Manually enable Matryoshka Embeddings
|
||||
args.extend([
|
||||
"--trust_remote_code", "--hf_overrides",
|
||||
'{"matryoshka_dimensions":[256]}'
|
||||
])
|
||||
args.extend(
|
||||
["--trust_remote_code", "--hf_overrides", '{"matryoshka_dimensions":[256]}']
|
||||
)
|
||||
|
||||
with RemoteOpenAIServer(model_info.name, args) as remote_server:
|
||||
yield remote_server
|
||||
@@ -64,14 +64,16 @@ def server(model_info, dtype: str):
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def hf_model(hf_runner, model_info, dtype: str):
|
||||
with hf_runner(model_info.name, dtype=dtype,
|
||||
is_sentence_transformer=True) as hf_model:
|
||||
with hf_runner(
|
||||
model_info.name, dtype=dtype, is_sentence_transformer=True
|
||||
) as hf_model:
|
||||
yield hf_model
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_matryoshka(model_info: EmbedModelInfo,
|
||||
server: RemoteOpenAIServer, hf_model: HfRunner):
|
||||
async def test_matryoshka(
|
||||
model_info: EmbedModelInfo, server: RemoteOpenAIServer, hf_model: HfRunner
|
||||
):
|
||||
client = server.get_async_client()
|
||||
|
||||
async def make_request_and_correctness_test(dimensions):
|
||||
@@ -84,7 +86,8 @@ async def test_matryoshka(model_info: EmbedModelInfo,
|
||||
encoding_format="float",
|
||||
)
|
||||
embeddings = EmbeddingResponse.model_validate(
|
||||
embedding_response.model_dump(mode="json"))
|
||||
embedding_response.model_dump(mode="json")
|
||||
)
|
||||
|
||||
assert embeddings.id is not None
|
||||
assert len(embeddings.data) == 3
|
||||
@@ -97,8 +100,7 @@ async def test_matryoshka(model_info: EmbedModelInfo,
|
||||
assert len(embeddings.data[0].embedding) == dimensions
|
||||
|
||||
vllm_outputs = [d.embedding for d in embeddings.data]
|
||||
run_embedding_correctness_test(hf_model, prompts, vllm_outputs,
|
||||
dimensions)
|
||||
run_embedding_correctness_test(hf_model, prompts, vllm_outputs, dimensions)
|
||||
|
||||
if model_info.is_matryoshka:
|
||||
valid_dimensions: list[Optional[int]] = [None]
|
||||
|
||||
Reference in New Issue
Block a user