test_jinavl_reranker.py 11 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from typing import cast
4
5
6
7

import pytest
from transformers import AutoModel

8
9
10
11
12
from vllm.entrypoints.chat_utils import (
    ChatCompletionContentPartImageEmbedsParam,
    ChatCompletionContentPartImageParam,
    ChatCompletionContentPartTextParam,
)
13
from vllm.entrypoints.pooling.score.utils import ScoreMultiModalParam
14
15
16

from ....conftest import HfRunner, VllmRunner

17
MODELS = ["jinaai/jina-reranker-m0"]
18

19
MM_PROCESSOR_KWARGS = {
20
21
22
23
    "min_pixels": 3136,
    "max_pixels": 602112,
}

24
LIMIT_MM_PER_PROMPT = {"image": 2}
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
83
84
85
86
87
88
89
90
CHECKPOINT_TO_HF_MAPPER = {
    "visual.": "model.visual.",
    "model.": "model.language_model.",
}

# Shared long text for test data
LONG_TEXT_DOC = """We present ReaderLM-v2, a compact 1.5 billion parameter language model designed for efficient
web content extraction. Our model processes documents up to 512K tokens, transforming messy HTML
into clean Markdown or JSON formats with high accuracy -- making it an ideal tool for grounding
large language models. The models effectiveness results from two key innovations: (1) a three-stage
data synthesis pipeline that generates high quality, diverse training data by iteratively drafting,
refining, and critiquing web content extraction; and (2) a unified training framework combining
continuous pre-training with multi-objective optimization. Intensive evaluation demonstrates that
ReaderLM-v2 outperforms GPT-4o-2024-08-06 and other larger models by 15-20% on carefully curated
benchmarks, particularly excelling at documents exceeding 100K tokens, while maintaining significantly
lower computational requirements."""  # noqa: E501

# Test data for different scenarios
TEXT_IMAGE_TEST_DATA = {
    "query": [{"text": "slm markdown"}],
    "documents": [
        {
            "image": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"
        },
        {
            "image": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
        },
    ],
}

TEXT_TEXT_TEST_DATA = {
    "query": [{"text": "slm markdown"}],
    "documents": [
        {"text": LONG_TEXT_DOC},
        {"text": "数据提取么?为什么不用正则啊,你用正则不就全解决了么?"},
    ],
}

IMAGE_TEXT_TEST_DATA = {
    "query": [
        {
            "image": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
        }
    ],
    "documents": [
        {"text": LONG_TEXT_DOC},
        {"text": "数据提取么?为什么不用正则啊,你用正则不就全解决了么?"},
    ],
}

IMAGE_IMAGE_TEST_DATA = {
    "query": [
        {
            "image": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
        }
    ],
    "documents": [
        {
            "image": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"
        },
        {
            "image": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
        },
    ],
}
91

92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
TEXT_MIXED_DOCS_TEST_DATA = {
    "query": [{"text": "slm markdown"}],
    "documents": [
        {"text": LONG_TEXT_DOC},
        {
            "image": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
        },
        {"text": "数据提取么?为什么不用正则啊,你用正则不就全解决了么?"},
        {
            "image": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"
        },
    ],
}


def _normalize_image(image_val: str) -> str:
    """Normalize image value to proper format for HF model."""
    return (
        image_val
        if image_val.startswith(("http://", "https://"))
        else f"data:image/png;base64,{image_val}"
    )


def create_score_multimodal_param(
    content_parts: list[dict],
) -> ScoreMultiModalParam:
    """
    Create a ScoreMultiModalParam from a list of content dictionaries.

    Each dict supports the following formats:
    - Text: {'text': 'content'}
    - Image URL: {'image': 'https://...'}
    - Image Base64: {'image': 'base64_str'}
    """
    formatted_content = []

    for part in content_parts:
        if "text" in part:
            formatted_content.append(
                ChatCompletionContentPartTextParam(
                    type="text",
                    text=part["text"],
                )
            )
        elif "image" in part:
            image_val = part["image"]
            if image_val.startswith(("http://", "https://")):
                formatted_content.append(
                    ChatCompletionContentPartImageParam(
                        type="image_url",
                        image_url={"url": image_val},
                    )
                )
            else:
                formatted_content.append(
                    ChatCompletionContentPartImageEmbedsParam(
                        type="image_embeds", image_embeds=image_val
                    )
                )

    return ScoreMultiModalParam(content=formatted_content)


def _run_vllm(
157
    vllm_runner: type[VllmRunner],
158
    model: str,
159
    dtype: str,
160
161
162
163
164
165
    query_strs: list[dict[str, str]],
    document_strs: list[dict[str, str]],
) -> list[float]:
    """Run vLLM reranker and return scores."""
    query = create_score_multimodal_param(query_strs)
    documents = create_score_multimodal_param(document_strs)
166
167

    with vllm_runner(
168
        model,
169
170
171
172
        runner="pooling",
        dtype=dtype,
        max_num_seqs=2,
        max_model_len=2048,
173
174
        mm_processor_kwargs=MM_PROCESSOR_KWARGS,
        limit_mm_per_prompt=LIMIT_MM_PER_PROMPT,
175
    ) as vllm_model:
176
        outputs = vllm_model.llm.score(query, documents)
177
178
179
180

    return [output.outputs.score for output in outputs]


181
def _run_hf(
182
    hf_runner: type[HfRunner],
183
    model: str,
184
    dtype: str,
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
    query_strs: list[dict[str, str]],
    document_strs: list[dict[str, str]],
) -> list[float]:
    """Run HuggingFace reranker and return scores."""
    query = query_strs[0]
    if "text" in query:
        query_type = "text"
        query_data = query["text"]
    elif "image" in query:
        query_type = "image"
        query_data = _normalize_image(query["image"])
    else:
        raise ValueError("Unsupported query format")

    # Separate documents by type
    text_docs: list[str] = []
    image_docs: list[str] = []
    text_indices: list[int] = []
    image_indices: list[int] = []

    for idx, doc in enumerate(document_strs):
        if "text" in doc:
            text_docs.append(doc["text"])
            text_indices.append(idx)
        elif "image" in doc:
            image_docs.append(_normalize_image(doc["image"]))
            image_indices.append(idx)
        else:
            raise ValueError(f"Unsupported document format at index {idx}")

    scores: list[None | float] = [None] * len(document_strs)
216

217
    with hf_runner(
218
        model,
219
220
221
        dtype=dtype,
        trust_remote_code=True,
        auto_cls=AutoModel,
222
        model_kwargs={"key_mapping": CHECKPOINT_TO_HF_MAPPER},
223
    ) as hf_model:
224
225
226
227
228
229
230
231
232
233
        # Score text documents
        if text_docs:
            text_scores = hf_model.model.compute_score(
                [[query_data, d] for d in text_docs],
                max_length=2048,
                query_type=query_type,
                doc_type="text",
            )
            for i, s in zip(text_indices, text_scores):
                scores[i] = s
234

235
236
237
238
239
240
241
242
243
244
        # Score image documents
        if image_docs:
            image_scores = hf_model.model.compute_score(
                [[query_data, d] for d in image_docs],
                max_length=2048,
                query_type=query_type,
                doc_type="image",
            )
            for i, s in zip(image_indices, image_scores):
                scores[i] = s
245

246
247
    assert all(s is not None for s in scores)
    return cast(list[float], scores)
248
249


250
251
252
253
254
255
256
257
258
259
260
261
262
def _run_test(
    hf_runner: type[HfRunner],
    vllm_runner: type[VllmRunner],
    model: str,
    dtype: str,
    query_strs: list[dict[str, str]],
    document_strs: list[dict[str, str]],
) -> None:
    """Run comparison test between vLLM and HuggingFace implementations."""
    # NOTE: take care of the order. run vLLM first, and then run HF.
    # vLLM needs a fresh new process without cuda initialization.
    # if we run HF first, the cuda initialization will be done and it
    # will hurt multiprocessing backend with fork method (the default method).
263

264
265
266
267
268
269
    vllm_outputs = _run_vllm(vllm_runner, model, dtype, query_strs, document_strs)
    hf_outputs = _run_hf(hf_runner, model, dtype, query_strs, document_strs)

    # Compare outputs
    assert len(hf_outputs) == len(vllm_outputs), (
        f"Output length mismatch: HF={len(hf_outputs)}, vLLM={len(vllm_outputs)}"
270
    )
271

272
273
274
275
    for i, (hf_score, vllm_score) in enumerate(zip(hf_outputs, vllm_outputs)):
        assert hf_score == pytest.approx(vllm_score, rel=0.02), (
            f"Score mismatch at index {i}: HF={hf_score}, vLLM={vllm_score}"
        )
276
277


278
@pytest.mark.parametrize("model", MODELS)
279
@pytest.mark.parametrize("dtype", ["half"])
280
281
282
283
284
285
286
287
288
289
290
291
292
293
def test_model_text_image(
    hf_runner,
    vllm_runner,
    model: str,
    dtype: str,
) -> None:
    """Visual Documents Reranking"""
    _run_test(
        hf_runner,
        vllm_runner,
        model,
        dtype,
        TEXT_IMAGE_TEST_DATA["query"],
        TEXT_IMAGE_TEST_DATA["documents"],
294
    )
295

296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313

@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["half"])
def test_model_text_text(
    hf_runner,
    vllm_runner,
    model: str,
    dtype: str,
) -> None:
    """Textual Documents Reranking"""
    _run_test(
        hf_runner,
        vllm_runner,
        model,
        dtype,
        TEXT_TEXT_TEST_DATA["query"],
        TEXT_TEXT_TEST_DATA["documents"],
    )
314
315


316
@pytest.mark.parametrize("model", MODELS)
317
@pytest.mark.parametrize("dtype", ["half"])
318
319
320
321
322
323
324
325
326
327
328
329
330
331
def test_model_image_text(
    hf_runner,
    vllm_runner,
    model: str,
    dtype: str,
) -> None:
    """Image Querying for Textual Documents"""
    _run_test(
        hf_runner,
        vllm_runner,
        model,
        dtype,
        IMAGE_TEXT_TEST_DATA["query"],
        IMAGE_TEXT_TEST_DATA["documents"],
332
    )
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350


@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["half"])
def test_model_image_image(
    hf_runner,
    vllm_runner,
    model: str,
    dtype: str,
) -> None:
    """Image Querying for Image Documents"""
    _run_test(
        hf_runner,
        vllm_runner,
        model,
        dtype,
        IMAGE_IMAGE_TEST_DATA["query"],
        IMAGE_IMAGE_TEST_DATA["documents"],
351
    )
352

353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370

@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["half"])
def test_model_text_mixed_documents(
    hf_runner,
    vllm_runner,
    model: str,
    dtype: str,
) -> None:
    """Text Query for Mixed Text and Image Documents"""
    _run_test(
        hf_runner,
        vllm_runner,
        model,
        dtype,
        TEXT_MIXED_DOCS_TEST_DATA["query"],
        TEXT_MIXED_DOCS_TEST_DATA["documents"],
    )