test_jina.py 3.98 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from functools import partial

5
6
import pytest

7
from vllm import PoolingParams
8

9
10
from ...utils import EmbedModelInfo, RerankModelInfo
from .embed_utils import (check_embeddings_close,
11
                          correctness_test_embed_models, matryoshka_fy)
12
from .mteb_utils import mteb_test_embed_models, mteb_test_rerank_models
13

14
EMBEDDING_MODELS = [
15
16
    EmbedModelInfo("jinaai/jina-embeddings-v3",
                   architecture="XLMRobertaModel",
17
                   is_matryoshka=True)
18
19
]

20
RERANK_MODELS = [
21
22
    RerankModelInfo("jinaai/jina-reranker-v2-base-multilingual",
                    architecture="XLMRobertaForSequenceClassification")
23
]
24
25


26
27
28
29
30
31
32
33
@pytest.fixture(autouse=True)
def v1(run_with_both_engines):
    # Simple autouse wrapper to run both engines for each test
    # This can be promoted up to conftest.py to run for every
    # test in a package
    pass


34
35
36
@pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
def test_embed_models_mteb(hf_runner, vllm_runner,
                           model_info: EmbedModelInfo) -> None:
37

38
39
    def hf_model_callback(model):
        model.encode = partial(model.encode, task="text-matching")
40

41
42
43
44
    mteb_test_embed_models(hf_runner,
                           vllm_runner,
                           model_info,
                           hf_model_callback=hf_model_callback)
45
46


47
48
49
50
@pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
def test_embed_models_correctness(hf_runner, vllm_runner,
                                  model_info: EmbedModelInfo,
                                  example_prompts) -> None:
51

52
53
    def hf_model_callback(model):
        model.encode = partial(model.encode, task="text-matching")
54

55
56
57
58
59
    correctness_test_embed_models(hf_runner,
                                  vllm_runner,
                                  model_info,
                                  example_prompts,
                                  hf_model_callback=hf_model_callback)
60
61


62
63
64
65
66
67
@pytest.mark.parametrize("model_info", RERANK_MODELS)
def test_rerank_models_mteb(hf_runner, vllm_runner,
                            model_info: RerankModelInfo) -> None:
    mteb_test_rerank_models(hf_runner, vllm_runner, model_info)


68
@pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
69
70
71
72
73
@pytest.mark.parametrize("dtype", ["half"])
@pytest.mark.parametrize("dimensions", [16, 32])
def test_matryoshka(
    hf_runner,
    vllm_runner,
74
    model_info,
75
76
    dtype: str,
    dimensions: int,
77
    example_prompts,
78
79
    monkeypatch,
) -> None:
80
81
    if not model_info.is_matryoshka:
        pytest.skip("Model is not matryoshka")
82

83
84
    # ST will strip the input texts, see test_embedding.py
    example_prompts = [str(s).strip() for s in example_prompts]
85
86

    with hf_runner(
87
            model_info.name,
88
89
90
91
92
93
            dtype=dtype,
            is_sentence_transformer=True,
    ) as hf_model:
        hf_outputs = hf_model.encode(example_prompts, task="text-matching")
        hf_outputs = matryoshka_fy(hf_outputs, dimensions)

94
95
96
    with vllm_runner(model_info.name,
                     task="embed",
                     dtype=dtype,
97
                     max_model_len=None) as vllm_model:
98
        assert vllm_model.llm.llm_engine.model_config.is_matryoshka
99

100
        matryoshka_dimensions = (
101
            vllm_model.llm.llm_engine.model_config.matryoshka_dimensions)
102
103
104
105
        assert matryoshka_dimensions is not None

        if dimensions not in matryoshka_dimensions:
            with pytest.raises(ValueError):
106
                vllm_model.embed(
107
108
109
                    example_prompts,
                    pooling_params=PoolingParams(dimensions=dimensions))
        else:
110
            vllm_outputs = vllm_model.embed(
111
112
113
114
115
116
117
118
119
120
                example_prompts,
                pooling_params=PoolingParams(dimensions=dimensions))

            check_embeddings_close(
                embeddings_0_lst=hf_outputs,
                embeddings_1_lst=vllm_outputs,
                name_0="hf",
                name_1="vllm",
                tol=1e-2,
            )