conftest.py 41.5 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
import json
4
import os
5
import tempfile
6
from collections import UserList
7
from enum import Enum
8
9
from typing import (Any, Callable, Dict, List, Optional, Tuple, Type,
                    TypedDict, TypeVar, Union)
10
11
import pytest
import pytest_html
Woosuk Kwon's avatar
Woosuk Kwon committed
12

13
import numpy as np
Woosuk Kwon's avatar
Woosuk Kwon committed
14
15
import pytest
import torch
16
import torch.nn as nn
17
import torch.nn.functional as F
18
from huggingface_hub import snapshot_download
19
from PIL import Image
20
from transformers import (AutoModelForCausalLM, AutoTokenizer, BatchEncoding,
21
                          BatchFeature)
22
from transformers.models.auto.auto_factory import _BaseAutoModelClass
Woosuk Kwon's avatar
Woosuk Kwon committed
23

24
25
from tests.models.utils import (TokensTextLogprobs,
                                TokensTextLogprobsPromptLogprobs)
Woosuk Kwon's avatar
Woosuk Kwon committed
26
from vllm import LLM, SamplingParams
27
from vllm.assets.image import ImageAsset
28
from vllm.assets.video import VideoAsset
29
from vllm.config import LoadFormat, TaskOption, TokenizerPoolConfig
30
from vllm.connections import global_http_connection
31
from vllm.distributed import (cleanup_dist_env_and_memory,
32
33
                              init_distributed_environment,
                              initialize_model_parallel)
34
from vllm.inputs import (ExplicitEncoderDecoderPrompt, TextPrompt,
youkaichao's avatar
youkaichao committed
35
36
                         TokensPrompt, to_enc_dec_tuple_list,
                         zip_enc_dec_prompts)
37
from vllm.logger import init_logger
38
from vllm.outputs import RequestOutput
39
from vllm.sampling_params import BeamSearchParams
40
from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, cuda_device_count_stateless,
youkaichao's avatar
youkaichao committed
41
                        identity, is_list_of)
zhuwenwen's avatar
zhuwenwen committed
42
from .utils import models_path_prefix
43

44
logger = init_logger(__name__)
Woosuk Kwon's avatar
Woosuk Kwon committed
45

46
47
48
_TEST_DIR = os.path.dirname(__file__)
_TEST_PROMPTS = [os.path.join(_TEST_DIR, "prompts", "example.txt")]
_LONG_PROMPTS = [os.path.join(_TEST_DIR, "prompts", "summary.txt")]
49
_SYS_MSG = os.path.join(_TEST_DIR, "system_messages", "sonnet3.5_nov2024.txt")
50

Cyrus Leung's avatar
Cyrus Leung committed
51
_M = TypeVar("_M")
52
53
54
55
56
57
58
59
60
61
62

MODELS_ON_S3 = [
    "distilbert/distilgpt2",
    "meta-llama/Llama-2-7b-hf",
    "meta-llama/Meta-Llama-3-8B",
    "meta-llama/Llama-3.2-1B",
    "meta-llama/Llama-3.2-1B-Instruct",
    "openai-community/gpt2",
    "ArthurZ/Ilama-3.2-1B",
    "llava-hf/llava-1.5-7b-hf",
    "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
    "ai21labs/Jamba-tiny-random",
    "neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV",
    "nm-testing/Phi-3-mini-128k-instruct-FP8",
    "nm-testing/Qwen2-0.5B-Instruct-FP8-SkipQKV",
    "neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV",
    "nm-testing/Qwen2-1.5B-Instruct-FP8-K-V",
    "ModelCloud/Qwen1.5-1.8B-Chat-GPTQ-4bits-dynamic-cfg-with-lm_head-symTrue",
    "ModelCloud/Qwen1.5-1.8B-Chat-GPTQ-4bits-dynamic-cfg-with-lm_head-symFalse",
    "AMead10/Llama-3.2-1B-Instruct-AWQ",
    "shuyuej/Llama-3.2-1B-Instruct-GPTQ",
    "ModelCloud/Qwen1.5-1.8B-Chat-GPTQ-4bits-dynamic-cfg-with-lm_head",
    "ModelCloud/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit-10-25-2024",
    "TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ",
    "neuralmagic/Meta-Llama-3-8B-Instruct-FP8",
    "amd/Llama-3.1-8B-Instruct-FP8-KV-Quark-test",
    "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change",
    "nm-testing/tinyllama-oneshot-w8-channel-a8-tensor",
    "nm-testing/asym-w8w8-int8-static-per-tensor-tiny-llama",
    "neuralmagic/Llama-3.2-1B-quantized.w8a8",
    "nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Dynamic-Asym",
    "nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Static-Per-Tensor-Sym",
    "nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Static-Per-Tensor-Asym",
    "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change",
    "nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2",
    "nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2-asym",
    "nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2",
    "nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2-asym",
    "nm-testing/tinyllama-oneshot-w4a16-channel-v2",
    "nm-testing/tinyllama-oneshot-w4a16-group128-v2",
    "nm-testing/tinyllama-oneshot-w8a16-per-channel",
    "nm-testing/llama7b-one-shot-2_4-w4a16-marlin24-t",
    "nm-testing/Meta-Llama-3-8B-FP8-compressed-tensors-test",
    "nm-testing/TinyLlama-1.1B-compressed-tensors-kv-cache-scheme",
    "nm-testing/Meta-Llama-3-8B-Instruct-FP8-Dynamic-2of4-testing",
    "nm-testing/Meta-Llama-3-8B-Instruct-FP8-Static-Per-Tensor-testing",
    "nm-testing/Meta-Llama-3-8B-Instruct-FP8-Static-testing",
    "nm-testing/Meta-Llama-3-8B-Instruct-FP8-Dynamic-IA-Per-Tensor-Weight-testing",
    "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-chnl_wts_per_tok_dyn_act_fp8-BitM",
    "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-chnl_wts_tensor_act_fp8-BitM",
    "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-tensor_wts_per_tok_dyn_act_fp8-BitM",
    "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-tensor_wts_tensor_act_fp8-BitM",
    "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-chnl_wts_per_tok_dyn_act_int8-BitM",
    "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-chnl_wts_tensor_act_int8-BitM",
    "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-tensor_wts_per_tok_dyn_act_int8-BitM",
    "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-tensor_wts_tensor_act_int8-BitM",
    "nm-testing/TinyLlama-1.1B-Chat-v1.0-INT8-Dynamic-IA-Per-Channel-Weight-testing",
    "nm-testing/TinyLlama-1.1B-Chat-v1.0-INT8-Static-testing",
    "nm-testing/TinyLlama-1.1B-Chat-v1.0-INT8-Dynamic-IA-Per-Tensor-Weight-testing",
    "nm-testing/TinyLlama-1.1B-Chat-v1.0-2of4-Sparse-Dense-Compressor",
    "nm-testing/llama2.c-stories42M-pruned2.4-compressed",
113
114
]

zhuwenwen's avatar
zhuwenwen committed
115
MODEL_WEIGHTS_S3_BUCKET = models_path_prefix
116

Cyrus Leung's avatar
Cyrus Leung committed
117
118
119
120
121
_PromptMultiModalInput = Union[List[_M], List[List[_M]]]

PromptImageInput = _PromptMultiModalInput[Image.Image]
PromptAudioInput = _PromptMultiModalInput[Tuple[np.ndarray, int]]
PromptVideoInput = _PromptMultiModalInput[np.ndarray]
122

123

124
def _read_prompts(filename: str) -> List[str]:
125
    with open(filename) as f:
126
127
        prompts = f.readlines()
        return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
128
129


130
131
132
class _ImageAssetPrompts(TypedDict):
    stop_sign: str
    cherry_blossom: str
133
134


135
136
class _ImageAssetsBase(UserList[ImageAsset]):
    pass
137

138
139

class _ImageAssets(_ImageAssetsBase):
140
141

    def __init__(self) -> None:
142
143
144
145
        super().__init__([
            ImageAsset("stop_sign"),
            ImageAsset("cherry_blossom"),
        ])
146
147
148
149
150
151
152
153

    def prompts(self, prompts: _ImageAssetPrompts) -> List[str]:
        """
        Convenience method to define the prompt for each test image.

        The order of the returned prompts matches the order of the
        assets when iterating through this object.
        """
154
        return [prompts["stop_sign"], prompts["cherry_blossom"]]
155
156


157
158
159
160
class _VideoAssetPrompts(TypedDict):
    sample_demo_1: str


161
162
class _VideoAssetsBase(UserList[VideoAsset]):
    pass
163
164
165
166
167
168
169
170
171
172
173
174
175


class _VideoAssets(_VideoAssetsBase):

    def __init__(self) -> None:
        super().__init__([
            VideoAsset("sample_demo_1.mp4"),
        ])

    def prompts(self, prompts: _VideoAssetPrompts) -> List[str]:
        return [prompts["sample_demo_1"]]


176
177
IMAGE_ASSETS = _ImageAssets()
"""Singleton instance of :class:`_ImageAssets`."""
178
179
VIDEO_ASSETS = _VideoAssets()
"""Singleton instance of :class:`_VideoAssets`."""
180
181


Joe Runde's avatar
Joe Runde committed
182
@pytest.fixture(params=[True, False])
183
def run_with_both_engines(request, monkeypatch):
Joe Runde's avatar
Joe Runde committed
184
185
186
187
188
189
190
191
    # Automatically runs tests twice, once with V1 and once without
    use_v1 = request.param
    # Tests decorated with `@skip_v1` are only run without v1
    skip_v1 = request.node.get_closest_marker("skip_v1")

    if use_v1:
        if skip_v1:
            pytest.skip("Skipping test on vllm V1")
192
        monkeypatch.setenv('VLLM_USE_V1', '1')
Joe Runde's avatar
Joe Runde committed
193
    else:
194
195
196
        monkeypatch.setenv('VLLM_USE_V1', '0')

    yield
Joe Runde's avatar
Joe Runde committed
197
198


199
200
201
202
203
204
205
@pytest.fixture(autouse=True)
def init_test_http_connection():
    # pytest_asyncio may use a different event loop per test
    # so we need to make sure the async client is created anew
    global_http_connection.reuse_client = False


206
207
208
209
210
211
212
213
214
215
216
217
@pytest.fixture
def dist_init():
    temp_file = tempfile.mkstemp()[1]
    init_distributed_environment(
        world_size=1,
        rank=0,
        distributed_init_method=f"file://{temp_file}",
        local_rank=0,
        backend="nccl",
    )
    initialize_model_parallel(1, 1)
    yield
218
    cleanup_dist_env_and_memory()
219
220


221
@pytest.fixture()
222
def should_do_global_cleanup_after_test(request) -> bool:
223
224
225
226
    """Allow subdirectories to skip global cleanup by overriding this fixture.
    This can provide a ~10x speedup for non-GPU unit tests since they don't need
    to initialize torch.
    """
227

228
    return not request.node.get_closest_marker("skip_global_cleanup")
229
230


231
@pytest.fixture(autouse=True)
232
def cleanup_fixture(should_do_global_cleanup_after_test: bool):
233
    yield
234
    if should_do_global_cleanup_after_test:
235
        cleanup_dist_env_and_memory()
236
237


238
239
240
241
242
243
@pytest.fixture(autouse=True)
def dynamo_reset():
    yield
    torch._dynamo.reset()


Woosuk Kwon's avatar
Woosuk Kwon committed
244
245
@pytest.fixture
def example_prompts() -> List[str]:
246
247
    prompts = []
    for filename in _TEST_PROMPTS:
248
        prompts += _read_prompts(filename)
249
250
251
    return prompts


252
253
254
255
256
257
@pytest.fixture
def example_system_message() -> str:
    with open(_SYS_MSG) as f:
        return f.read()


258
259
260
261
262
263
264
class DecoderPromptType(Enum):
    """For encoder/decoder models only."""
    CUSTOM = 1
    NONE = 2
    EMPTY_STR = 3


265
@pytest.fixture
266
267
def example_encoder_decoder_prompts(
) -> Dict[DecoderPromptType, List[ExplicitEncoderDecoderPrompt]]:
268
269
270
271
272
273
    '''
    Returns an encoder prompt list and a decoder prompt list, wherein each pair
    of same-index entries in both lists corresponds to an (encoder prompt,
    decoder prompt) tuple.

    Returns:
274

275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
    * Encoder prompt list
    * Decoder prompt list (reverse of encoder prompt list)
    '''

    encoder_prompts = []
    for filename in _TEST_PROMPTS:
        encoder_prompts += _read_prompts(filename)

    custom_decoder_prompts = encoder_prompts[::-1]
    empty_str_decoder_prompts = [""] * len(encoder_prompts)
    none_decoder_prompts = [None] * len(encoder_prompts)

    # NONE decoder prompt type
    return {
        DecoderPromptType.NONE:
290
        zip_enc_dec_prompts(encoder_prompts, none_decoder_prompts),
291
        DecoderPromptType.EMPTY_STR:
292
        zip_enc_dec_prompts(encoder_prompts, empty_str_decoder_prompts),
293
        DecoderPromptType.CUSTOM:
294
        zip_enc_dec_prompts(encoder_prompts, custom_decoder_prompts),
295
296
297
    }


298
299
300
301
@pytest.fixture
def example_long_prompts() -> List[str]:
    prompts = []
    for filename in _LONG_PROMPTS:
302
        prompts += _read_prompts(filename)
303
    return prompts
Woosuk Kwon's avatar
Woosuk Kwon committed
304
305


306
307
308
309
310
@pytest.fixture(scope="session")
def image_assets() -> _ImageAssets:
    return IMAGE_ASSETS


311
312
313
314
315
@pytest.fixture(scope="session")
def video_assets() -> _VideoAssets:
    return VIDEO_ASSETS


316
_T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
317
_R = TypeVar("_R")
318

Woosuk Kwon's avatar
Woosuk Kwon committed
319
320
321

class HfRunner:

322
    def wrap_device(self, x: _T, device: Optional[str] = None) -> _T:
323
        from vllm.platforms import current_platform
324
325
326
        if x is None or isinstance(x, (bool, )):
            return x

327
        if device is None:
328
            device = "cpu" if current_platform.is_cpu() else "cuda"
329

330
331
        if isinstance(x, dict):
            return {k: self.wrap_device(v, device) for k, v in x.items()}
332

333
334
335
336
        if hasattr(x, "device") and x.device.type == device:
            return x

        return x.to(device)
337

Woosuk Kwon's avatar
Woosuk Kwon committed
338
339
340
341
    def __init__(
        self,
        model_name: str,
        dtype: str = "half",
342
        *,
343
        model_kwargs: Optional[Dict[str, Any]] = None,
344
        is_sentence_transformer: bool = False,
345
        is_cross_encoder: bool = False,
346
        skip_tokenizer_init: bool = False,
347
        auto_cls: Type[_BaseAutoModelClass] = AutoModelForCausalLM,
348
        postprocess_inputs: Callable[..., BatchEncoding] = identity,
Woosuk Kwon's avatar
Woosuk Kwon committed
349
    ) -> None:
350
        torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[dtype]
351

352
        self.model_name = model_name
353

354
        if is_sentence_transformer:
355
356
            # Lazy init required for AMD CI
            from sentence_transformers import SentenceTransformer
357
358
359
360
            self.model = self.wrap_device(
                SentenceTransformer(
                    model_name,
                    device="cpu",
361
                    trust_remote_code=True,
362
                ).to(dtype=torch_dtype))
363
364
365
366
367
368
369
370
        elif is_cross_encoder:
            # Lazy init required for AMD CI
            from sentence_transformers import CrossEncoder
            self.model = CrossEncoder(model_name,
                                      device="cpu",
                                      trust_remote_code=True)
            self.model.model = self.wrap_device(self.model.model)\
                .to(dtype=torch_dtype)
371
        else:
372
            model_kwargs = model_kwargs if model_kwargs is not None else {}
373
            self.model = self.wrap_device(
374
                auto_cls.from_pretrained(
375
376
377
                    model_name,
                    torch_dtype=torch_dtype,
                    trust_remote_code=True,
378
                    **model_kwargs,
379
                ))
380

381
382
383
384
385
386
        if not skip_tokenizer_init:
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_name,
                torch_dtype=torch_dtype,
                trust_remote_code=True,
            )
387

388
389
390
391
392
393
394
395
        # don't put this import at the top level
        # it will call torch.cuda.device_count()
        from transformers import AutoProcessor  # noqa: F401
        self.processor = AutoProcessor.from_pretrained(
            model_name,
            torch_dtype=torch_dtype,
            trust_remote_code=True,
        )
396
397
        if skip_tokenizer_init:
            self.tokenizer = self.processor.tokenizer
Woosuk Kwon's avatar
Woosuk Kwon committed
398

399
        self.dtype = dtype
400
401
        self.postprocess_inputs = postprocess_inputs

402
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
403
404
        self,
        prompts: List[str],
405
        images: Optional[PromptImageInput] = None,
406
407
408
409
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
    ) -> List[BatchEncoding]:
        if images is not None:
410
            assert len(prompts) == len(images)
411

412
413
414
415
416
417
418
        if videos is not None:
            assert len(prompts) == len(videos)

        if audios is not None:
            assert len(prompts) == len(audios)

        all_inputs: List[BatchEncoding] = []
419
        for i, prompt in enumerate(prompts):
420
421
422
423
            processor_kwargs: Dict[str, Any] = {
                "text": prompt,
                "return_tensors": "pt",
            }
Cyrus Leung's avatar
Cyrus Leung committed
424
425
426
427
428
429
            if images is not None and (image := images[i]) is not None:
                processor_kwargs["images"] = image
            if videos is not None and (video := videos[i]) is not None:
                processor_kwargs["videos"] = video
            if audios is not None and (audio_tuple := audios[i]) is not None:
                audio, sr = audio_tuple
430
431
                processor_kwargs["audio"] = audio
                processor_kwargs["sampling_rate"] = sr
432
433

            inputs = self.processor(**processor_kwargs)
434
            inputs = self.postprocess_inputs(inputs, dtype=self.dtype)
435

436
437
438
439
            all_inputs.append(inputs)

        return all_inputs

440
441
442
443
444
445
446
447
448
449
450
    def classify(self, prompts: List[str]) -> List[str]:
        # output is final logits
        all_inputs = self.get_inputs(prompts)
        outputs = []
        for inputs in all_inputs:
            output = self.model(**self.wrap_device(inputs))
            logits = output.logits.softmax(dim=-1)[0].tolist()
            outputs.append(logits)

        return outputs

451
452
453
454
    def generate(
        self,
        prompts: List[str],
        images: Optional[PromptImageInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
455
        videos: Optional[PromptVideoInput] = None,
456
457
458
459
460
461
462
        audios: Optional[PromptAudioInput] = None,
        **kwargs: Any,
    ) -> List[Tuple[List[List[int]], List[str]]]:
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)
463

464
465
        outputs: List[Tuple[List[List[int]], List[str]]] = []
        for inputs in all_inputs:
Woosuk Kwon's avatar
Woosuk Kwon committed
466
            output_ids = self.model.generate(
467
                **self.wrap_device(inputs, device=self.model.device.type),
Woosuk Kwon's avatar
Woosuk Kwon committed
468
469
470
                use_cache=True,
                **kwargs,
            )
471
            output_str = self.processor.batch_decode(
Woosuk Kwon's avatar
Woosuk Kwon committed
472
473
474
                output_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
475
476
            )
            output_ids = output_ids.cpu().tolist()
Woosuk Kwon's avatar
Woosuk Kwon committed
477
478
479
480
481
482
483
            outputs.append((output_ids, output_str))
        return outputs

    def generate_greedy(
        self,
        prompts: List[str],
        max_tokens: int,
484
        images: Optional[PromptImageInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
485
        videos: Optional[PromptVideoInput] = None,
486
        audios: Optional[PromptAudioInput] = None,
487
        **kwargs: Any,
Woosuk Kwon's avatar
Woosuk Kwon committed
488
    ) -> List[Tuple[List[int], str]]:
489
490
        outputs = self.generate(prompts,
                                do_sample=False,
491
                                max_new_tokens=max_tokens,
Chang Su's avatar
Chang Su committed
492
                                images=images,
493
494
                                videos=videos,
                                audios=audios,
Chang Su's avatar
Chang Su committed
495
                                **kwargs)
496
497
498

        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
499
500
501
502
503
504

    def generate_beam_search(
        self,
        prompts: List[str],
        beam_width: int,
        max_tokens: int,
505
    ) -> List[Tuple[List[List[int]], List[str]]]:
506
507
508
509
510
511
512
513
514
515
516
517
518
519
        outputs = self.generate(prompts,
                                do_sample=False,
                                max_new_tokens=max_tokens,
                                num_beams=beam_width,
                                num_return_sequences=beam_width)
        for i in range(len(outputs)):
            output_ids, output_str = outputs[i]
            for j in range(len(output_ids)):
                output_ids[j] = [
                    x for x in output_ids[j]
                    if x != self.tokenizer.pad_token_id
                ]
            outputs[i] = (output_ids, output_str)
        return outputs
Woosuk Kwon's avatar
Woosuk Kwon committed
520

521
522
523
524
    def generate_greedy_logprobs(
        self,
        prompts: List[str],
        max_tokens: int,
525
        images: Optional[PromptImageInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
526
        videos: Optional[PromptVideoInput] = None,
527
        audios: Optional[PromptAudioInput] = None,
528
        **kwargs: Any,
529
    ) -> List[List[torch.Tensor]]:
530
531
532
533
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)
534

535
536
        all_logprobs: List[List[torch.Tensor]] = []
        for inputs in all_inputs:
537
            output = self.model.generate(
538
                **self.wrap_device(inputs, device=self.model.device.type),
539
540
541
542
543
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
544
                **kwargs,
545
            )
546
547
            seq_logprobs = self._hidden_states_to_seq_logprobs(
                output.hidden_states)
548
549
550
            all_logprobs.append(seq_logprobs)
        return all_logprobs

551
    def _hidden_states_to_seq_logprobs(
552
        self,
553
554
555
556
        hidden_states: Tuple[Tuple[torch.Tensor, ...], ...],
    ) -> List[torch.Tensor]:
        output_embeddings = self.model.get_output_embeddings()

557
558
559
560
        seq_logprobs: List[torch.Tensor] = []
        for _, hidden_state in enumerate(hidden_states):
            last_hidden_states = hidden_state[-1][0]
            logits = torch.matmul(
561
562
                last_hidden_states.to(output_embeddings.weight.device),
                output_embeddings.weight.t(),
563
            )
564
565
            if getattr(output_embeddings, "bias", None) is not None:
                logits += output_embeddings.bias.unsqueeze(0)
566
567
568
            logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
            seq_logprobs.append(logprobs)

569
570
571
572
573
574
575
576
577
578
        return seq_logprobs

    def _hidden_states_to_logprobs(
        self,
        hidden_states: Tuple[Tuple[torch.Tensor, ...], ...],
        num_logprobs: int,
    ) -> Tuple[List[Dict[int, float]], int]:
        seq_logprobs = self._hidden_states_to_seq_logprobs(hidden_states)
        output_len = len(hidden_states)

579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
        # convert to dict
        seq_logprobs_lst: List[Dict[int, float]] = []
        for tok_idx, tok_logprobs in enumerate(seq_logprobs):
            # drop prompt logprobs
            if tok_idx == 0:
                tok_logprobs = tok_logprobs[-1, :].reshape(1, -1)
            topk = tok_logprobs.topk(num_logprobs)

            tok_logprobs_dct = {}
            for token_id, logprob in zip(topk.indices[0], topk.values[0]):
                tok_logprobs_dct[token_id.item()] = logprob.item()

            seq_logprobs_lst.append(tok_logprobs_dct)

        return (
            seq_logprobs_lst,
            output_len,
        )

598
599
600
601
602
    def generate_greedy_logprobs_limit(
        self,
        prompts: List[str],
        max_tokens: int,
        num_logprobs: int,
603
604
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
Cyrus Leung's avatar
Cyrus Leung committed
605
        videos: Optional[PromptVideoInput] = None,
606
        **kwargs: Any,
607
    ) -> List[TokensTextLogprobs]:
608
609
610
611
612
        all_inputs = self.get_inputs(prompts,
                                     images=images,
                                     videos=videos,
                                     audios=audios)

613
614
615
        all_logprobs: List[List[Dict[int, float]]] = []
        all_output_ids: List[List[int]] = []
        all_output_strs: List[str] = []
616

617
        for inputs in all_inputs:
618
            output = self.model.generate(
619
                **self.wrap_device(inputs, device=self.model.device.type),
620
621
622
623
624
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
625
                **kwargs,
626
627
            )

628
629
630
631
632
633
634
635
636
637
638
639
            (
                seq_logprobs_lst,
                output_len,
            ) = self._hidden_states_to_logprobs(output.hidden_states,
                                                num_logprobs)

            all_logprobs.append(seq_logprobs_lst)
            seq_ids = output.sequences[0]
            output_len = len(seq_logprobs_lst)
            output_ids = seq_ids[-output_len:]
            all_output_ids.append(output_ids.tolist())
            all_output_strs.append(self.tokenizer.decode(output_ids))
640

641
642
643
644
645
646
        outputs = zip(all_output_ids, all_output_strs, all_logprobs)
        return [(output_ids, output_str, output_logprobs)
                for output_ids, output_str, output_logprobs in outputs]

    def generate_encoder_decoder_greedy_logprobs_limit(
        self,
647
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
648
649
        max_tokens: int,
        num_logprobs: int,
650
        images: Optional[PromptImageInput] = None,
651
        **kwargs: Any,
652
    ) -> List[TokensTextLogprobs]:
653
654
655
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''
656

657
658
659
        all_logprobs: List[List[Dict[int, float]]] = []
        all_output_ids: List[List[int]] = []
        all_output_strs: List[str] = []
660

661
662
663
664
665
666
667
668
        for i, (encoder_prompt, decoder_prompt) in enumerate(
                to_enc_dec_tuple_list(encoder_decoder_prompts)):
            processor_kwargs: Dict[str, Any] = {
                "text": encoder_prompt,
                "return_tensors": "pt",
            }
            if images is not None and images[i] is not None:
                processor_kwargs["images"] = images[i]
669

670
            encoder_input_ids = self.wrap_device(
671
                self.processor(**processor_kwargs).input_ids,
672
673
674
675
676
677
678
                device=self.model.device.type,
            )

            if decoder_prompt is None:
                decoder_input_ids = None
            else:
                decoder_input_ids = self.wrap_device(
679
                    self.tokenizer(decoder_prompt,
680
681
682
                                   return_tensors="pt").input_ids,
                    device=self.model.device.type,
                )
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699

            output = self.model.generate(
                encoder_input_ids,
                decoder_input_ids=decoder_input_ids,
                use_cache=True,
                do_sample=False,
                max_new_tokens=max_tokens,
                output_hidden_states=True,
                return_dict_in_generate=True,
                **kwargs,
            )

            (
                seq_logprobs_lst,
                output_len,
            ) = self._hidden_states_to_logprobs(output.decoder_hidden_states,
                                                num_logprobs)
700
701
702
703
704
705
706
707
708
709
710

            all_logprobs.append(seq_logprobs_lst)
            seq_ids = output.sequences[0]
            output_ids = seq_ids[-output_len:]
            all_output_ids.append(output_ids.tolist())
            all_output_strs.append(self.tokenizer.decode(output_ids))

        outputs = zip(all_output_ids, all_output_strs, all_logprobs)
        return [(output_ids, output_str, output_logprobs)
                for output_ids, output_str, output_logprobs in outputs]

711
712
713
    def encode(self, prompts: List[str]) -> List[List[torch.Tensor]]:
        return self.model.encode(prompts)

714
715
716
    def predict(self, prompts: List[List[str]]) -> torch.Tensor:
        return self.model.predict(prompts, convert_to_tensor=True)

717
718
719
720
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
721
        del self.model
722
        cleanup_dist_env_and_memory()
723

Woosuk Kwon's avatar
Woosuk Kwon committed
724

Cyrus Leung's avatar
Cyrus Leung committed
725
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
726
727
728
729
730
731
732
733
734
def hf_runner():
    return HfRunner


class VllmRunner:

    def __init__(
        self,
        model_name: str,
735
        task: TaskOption = "auto",
Woosuk Kwon's avatar
Woosuk Kwon committed
736
        tokenizer_name: Optional[str] = None,
737
        tokenizer_mode: str = "auto",
738
739
        # Use smaller max model length, otherwise bigger model cannot run due
        # to kv cache size limit.
740
        max_model_len: int = 1024,
Woosuk Kwon's avatar
Woosuk Kwon committed
741
        dtype: str = "half",
742
        disable_log_stats: bool = True,
743
        tensor_parallel_size: int = 1,
744
745
        block_size: int = 16,
        enable_chunked_prefill: bool = False,
746
        swap_space: int = 4,
747
        enforce_eager: Optional[bool] = False,
748
        load_format: Optional[LoadFormat] = None,
749
        **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
750
    ) -> None:
751
        if model_name in MODELS_ON_S3 and not load_format:
752
            model_name = (f"{MODEL_WEIGHTS_S3_BUCKET}/{model_name}")
753
754
755
            load_format = LoadFormat.RUNAI_STREAMER
        if not load_format:
            load_format = LoadFormat.AUTO
Woosuk Kwon's avatar
Woosuk Kwon committed
756
757
        self.model = LLM(
            model=model_name,
758
            task=task,
Woosuk Kwon's avatar
Woosuk Kwon committed
759
            tokenizer=tokenizer_name,
760
            tokenizer_mode=tokenizer_mode,
Woosuk Kwon's avatar
Woosuk Kwon committed
761
762
            trust_remote_code=True,
            dtype=dtype,
763
            swap_space=swap_space,
Cyrus Leung's avatar
Cyrus Leung committed
764
            enforce_eager=enforce_eager,
765
            disable_log_stats=disable_log_stats,
766
            tensor_parallel_size=tensor_parallel_size,
767
            max_model_len=max_model_len,
768
769
            block_size=block_size,
            enable_chunked_prefill=enable_chunked_prefill,
770
            load_format=load_format,
771
            **kwargs,
Woosuk Kwon's avatar
Woosuk Kwon committed
772
773
        )

774
    def get_inputs(
Woosuk Kwon's avatar
Woosuk Kwon committed
775
776
        self,
        prompts: List[str],
777
        images: Optional[PromptImageInput] = None,
778
779
780
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
    ) -> List[TextPrompt]:
781
        if images is not None:
782
            assert len(prompts) == len(images)
783

784
785
786
787
788
789
        if videos is not None:
            assert len(prompts) == len(videos)

        if audios is not None:
            assert len(prompts) == len(audios)

790
791
792
        inputs = [TextPrompt(prompt=prompt) for prompt in prompts]
        if images is not None:
            for i, image in enumerate(images):
Cyrus Leung's avatar
Cyrus Leung committed
793
794
                if image is not None:
                    inputs[i]["multi_modal_data"] = {"image": image}
795

796
797
        if videos is not None:
            for i, video in enumerate(videos):
Cyrus Leung's avatar
Cyrus Leung committed
798
799
                if video is not None:
                    inputs[i]["multi_modal_data"] = {"video": video}
800
801
802

        if audios is not None:
            for i, audio in enumerate(audios):
Cyrus Leung's avatar
Cyrus Leung committed
803
804
                if audio is not None:
                    inputs[i]["multi_modal_data"] = {"audio": audio}
805
806
807
808
809
810
811
812
813
814

        return inputs

    def generate(
        self,
        prompts: List[str],
        sampling_params: SamplingParams,
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
815
        **kwargs: Any,
816
817
818
819
820
    ) -> List[Tuple[List[List[int]], List[str]]]:
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)
821

822
        req_outputs = self.model.generate(inputs,
823
824
                                          sampling_params=sampling_params,
                                          **kwargs)
825
826

        outputs: List[Tuple[List[List[int]], List[str]]] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
827
828
829
        for req_output in req_outputs:
            prompt_str = req_output.prompt
            prompt_ids = req_output.prompt_token_ids
830
831
            req_sample_output_ids: List[List[int]] = []
            req_sample_output_strs: List[str] = []
832
833
            for sample in req_output.outputs:
                output_str = sample.text
834
                output_ids = list(sample.token_ids)
835
836
837
                req_sample_output_ids.append(prompt_ids + output_ids)
                req_sample_output_strs.append(prompt_str + output_str)
            outputs.append((req_sample_output_ids, req_sample_output_strs))
Woosuk Kwon's avatar
Woosuk Kwon committed
838
839
        return outputs

840
    @staticmethod
841
842
    def _final_steps_generate_w_logprobs(
        req_outputs: List[RequestOutput],
843
844
    ) -> List[TokensTextLogprobsPromptLogprobs]:
        outputs: List[TokensTextLogprobsPromptLogprobs] = []
845
        for req_output in req_outputs:
846
            assert len(req_output.outputs) > 0
847
848
            for sample in req_output.outputs:
                output_str = sample.text
849
                output_ids = list(sample.token_ids)
850
                output_logprobs = sample.logprobs
851
852
            outputs.append((output_ids, output_str, output_logprobs,
                            req_output.prompt_logprobs))
853
854
        return outputs

855
856
857
858
    def generate_w_logprobs(
        self,
        prompts: List[str],
        sampling_params: SamplingParams,
859
860
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
861
        videos: Optional[PromptVideoInput] = None,
862
        **kwargs: Any,
863
864
    ) -> Union[List[TokensTextLogprobs],
               List[TokensTextLogprobsPromptLogprobs]]:
865
866
867
868
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)
869

870
        req_outputs = self.model.generate(inputs,
871
872
                                          sampling_params=sampling_params,
                                          **kwargs)
873
874
875
876
877
878
879

        toks_str_logsprobs_prompt_logprobs = (
            self._final_steps_generate_w_logprobs(req_outputs))
        # Omit prompt logprobs if not required by sampling params
        return ([x[0:-1] for x in toks_str_logsprobs_prompt_logprobs]
                if sampling_params.prompt_logprobs is None else
                toks_str_logsprobs_prompt_logprobs)
880
881
882

    def generate_encoder_decoder_w_logprobs(
        self,
883
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
884
        sampling_params: SamplingParams,
885
886
    ) -> Union[List[TokensTextLogprobs],
               List[TokensTextLogprobsPromptLogprobs]]:
887
888
889
890
891
892
893
        '''
        Logprobs generation for vLLM encoder/decoder models
        '''

        assert sampling_params.logprobs is not None
        req_outputs = self.model.generate(encoder_decoder_prompts,
                                          sampling_params=sampling_params)
894
895
896
897
898
899
        toks_str_logsprobs_prompt_logprobs = (
            self._final_steps_generate_w_logprobs(req_outputs))
        # Omit prompt logprobs if not required by sampling params
        return ([x[0:-1] for x in toks_str_logsprobs_prompt_logprobs]
                if sampling_params.prompt_logprobs is None else
                toks_str_logsprobs_prompt_logprobs)
900

Woosuk Kwon's avatar
Woosuk Kwon committed
901
902
903
904
    def generate_greedy(
        self,
        prompts: List[str],
        max_tokens: int,
905
        images: Optional[PromptImageInput] = None,
906
907
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
908
        **kwargs: Any,
Woosuk Kwon's avatar
Woosuk Kwon committed
909
910
    ) -> List[Tuple[List[int], str]]:
        greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
911
912
913
914
        outputs = self.generate(prompts,
                                greedy_params,
                                images=images,
                                videos=videos,
915
916
                                audios=audios,
                                **kwargs)
917
918
        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]
919

920
921
922
923
924
    def generate_greedy_logprobs(
        self,
        prompts: List[str],
        max_tokens: int,
        num_logprobs: int,
925
        num_prompt_logprobs: Optional[int] = None,
926
927
        images: Optional[PromptImageInput] = None,
        audios: Optional[PromptAudioInput] = None,
928
        videos: Optional[PromptVideoInput] = None,
929
        stop_token_ids: Optional[List[int]] = None,
930
        stop: Optional[List[str]] = None,
931
        **kwargs: Any,
932
933
934
935
936
937
    ) -> Union[List[TokensTextLogprobs],
               List[TokensTextLogprobsPromptLogprobs]]:
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
938
            prompt_logprobs=num_prompt_logprobs,
939
940
            stop_token_ids=stop_token_ids,
            stop=stop)
941
942
943
944
945

        return self.generate_w_logprobs(prompts,
                                        greedy_logprobs_params,
                                        images=images,
                                        audios=audios,
946
947
                                        videos=videos,
                                        **kwargs)
948

949
950
    def generate_encoder_decoder_greedy_logprobs(
        self,
951
        encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
952
953
        max_tokens: int,
        num_logprobs: int,
954
955
956
957
958
959
960
961
962
        num_prompt_logprobs: Optional[int] = None,
    ) -> Union[List[TokensTextLogprobs],
               List[TokensTextLogprobsPromptLogprobs]]:
        greedy_logprobs_params = SamplingParams(
            temperature=0.0,
            max_tokens=max_tokens,
            logprobs=num_logprobs,
            prompt_logprobs=(num_prompt_logprobs),
        )
963
964
965
966
        '''
        Greedy logprobs generation for vLLM encoder/decoder models
        '''

967
        return self.generate_encoder_decoder_w_logprobs(
968
969
            encoder_decoder_prompts, greedy_logprobs_params)

970
    def generate_beam_search(
971
972
973
974
975
        self,
        prompts: Union[List[str], List[List[int]]],
        beam_width: int,
        max_tokens: int,
    ) -> List[Tuple[List[List[int]], List[str]]]:
youkaichao's avatar
youkaichao committed
976
977
978
979
980
981
        if is_list_of(prompts, str, check="all"):
            prompts = [TextPrompt(prompt=prompt) for prompt in prompts]
        else:
            prompts = [
                TokensPrompt(prompt_token_ids=tokens) for tokens in prompts
            ]
982
983
984
        outputs = self.model.beam_search(
            prompts,
            BeamSearchParams(beam_width=beam_width, max_tokens=max_tokens))
985
986
987
988
989
990
991
        returned_outputs = []
        for output in outputs:
            token_ids = [x.tokens for x in output.sequences]
            texts = [x.text for x in output.sequences]
            returned_outputs.append((token_ids, texts))
        return returned_outputs

992
993
994
995
    def classify(self, prompts: List[str]) -> List[List[float]]:
        req_outputs = self.model.classify(prompts)
        return [req_output.outputs.probs for req_output in req_outputs]

Cyrus Leung's avatar
Cyrus Leung committed
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
    def encode(
        self,
        prompts: List[str],
        images: Optional[PromptImageInput] = None,
        videos: Optional[PromptVideoInput] = None,
        audios: Optional[PromptAudioInput] = None,
    ) -> List[List[float]]:
        inputs = self.get_inputs(prompts,
                                 images=images,
                                 videos=videos,
                                 audios=audios)

1008
        req_outputs = self.model.embed(inputs)
Cyrus Leung's avatar
Cyrus Leung committed
1009
        return [req_output.outputs.embedding for req_output in req_outputs]
1010

1011
1012
1013
1014
    def score(
        self,
        text_1: Union[str, List[str]],
        text_2: Union[str, List[str]],
1015
    ) -> List[float]:
1016
        req_outputs = self.model.score(text_1, text_2)
1017
        return [req_output.outputs.score for req_output in req_outputs]
1018

1019
1020
1021
1022
    def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
        executor = self.model.llm_engine.model_executor
        return executor.apply_model(func)

1023
1024
1025
1026
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
1027
        del self.model
1028
        cleanup_dist_env_and_memory()
1029

Woosuk Kwon's avatar
Woosuk Kwon committed
1030

1031
@pytest.fixture(scope="session")
Woosuk Kwon's avatar
Woosuk Kwon committed
1032
1033
def vllm_runner():
    return VllmRunner
1034
1035
1036
1037
1038
1039
1040
1041
1042


def get_tokenizer_pool_config(tokenizer_group_type):
    if tokenizer_group_type is None:
        return None
    if tokenizer_group_type == "ray":
        return TokenizerPoolConfig(pool_size=1,
                                   pool_type="ray",
                                   extra_config={})
1043
1044
1045
1046
    if isinstance(tokenizer_group_type, type):
        return TokenizerPoolConfig(pool_size=1,
                                   pool_type=tokenizer_group_type,
                                   extra_config={})
1047
    raise ValueError(f"Unknown tokenizer_group_type: {tokenizer_group_type}")
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063


@pytest.fixture()
def temporary_enable_log_propagate():
    import logging
    logger = logging.getLogger("vllm")
    logger.propagate = True
    yield
    logger.propagate = False


@pytest.fixture()
def caplog_vllm(temporary_enable_log_propagate, caplog):
    # To capture vllm log, we should enable propagate=True temporarily
    # because caplog depends on logs propagated to the root logger.
    yield caplog
1064
1065
1066
1067
1068
1069
1070


@pytest.fixture(scope="session")
def num_gpus_available():
    """Get number of GPUs without initializing the CUDA context
    in current process."""

1071
    return cuda_device_count_stateless()
1072
1073


1074
# temp_dir = tempfile.gettempdir()
zhuwenwen's avatar
zhuwenwen committed
1075
1076
1077
1078
_dummy_opt_path = os.path.join(models_path_prefix, "dummy_opt")
_dummy_llava_path = os.path.join(models_path_prefix, "dummy_llava")
_dummy_gemma2_embedding_path = os.path.join(models_path_prefix, "dummy_gemma2_embedding")

1079
1080
1081
1082


@pytest.fixture
def dummy_opt_path():
1083
1084
    json_path = os.path.join(_dummy_opt_path, "config.json")
    if not os.path.exists(_dummy_opt_path):
1085
        snapshot_download(repo_id="facebook/opt-125m",
1086
                          local_dir=_dummy_opt_path,
1087
1088
1089
1090
1091
                          ignore_patterns=[
                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
                              "*.msgpack"
                          ])
        assert os.path.exists(json_path)
1092
        with open(json_path) as f:
1093
1094
1095
1096
            config = json.load(f)
        config["architectures"] = ["MyOPTForCausalLM"]
        with open(json_path, "w") as f:
            json.dump(config, f)
1097
1098
    return _dummy_opt_path

1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116

# 定义一个 pytest 钩子,在测试后生成报告
@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_runtest_makereport(item, call):
    # 获取测试结果
    outcome = yield
    result = outcome.get_result()

    # 如果测试失败并且有浏览器实例,添加截图
    if result.when == "call" and result.failed:
        if hasattr(item, "funcargs") and "browser" in item.funcargs:
            browser = item.funcargs["browser"]
            screenshot_path = "screenshot.png"  # 设置截图路径
            browser.save_screenshot(screenshot_path)

            # 如果测试结果有 extra 属性,则添加截图
            if hasattr(result, "extra"):
                result.extra.append(pytest_html.extras.image(screenshot_path))
zhuwenwen's avatar
zhuwenwen committed
1117
1118


1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
@pytest.fixture
def dummy_llava_path():
    json_path = os.path.join(_dummy_llava_path, "config.json")
    if not os.path.exists(_dummy_llava_path):
        snapshot_download(repo_id="llava-hf/llava-1.5-7b-hf",
                          local_dir=_dummy_llava_path,
                          ignore_patterns=[
                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
                              "*.msgpack"
                          ])
        assert os.path.exists(json_path)
1130
        with open(json_path) as f:
1131
1132
1133
1134
1135
            config = json.load(f)
        config["architectures"] = ["MyLlava"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_llava_path
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148


@pytest.fixture
def dummy_gemma2_embedding_path():
    json_path = os.path.join(_dummy_gemma2_embedding_path, "config.json")
    if not os.path.exists(_dummy_gemma2_embedding_path):
        snapshot_download(repo_id="BAAI/bge-multilingual-gemma2",
                          local_dir=_dummy_gemma2_embedding_path,
                          ignore_patterns=[
                              "*.bin", "*.bin.index.json", "*.pt", "*.h5",
                              "*.msgpack"
                          ])
        assert os.path.exists(json_path)
1149
        with open(json_path) as f:
1150
1151
1152
1153
1154
            config = json.load(f)
        config["architectures"] = ["MyGemma2Embedding"]
        with open(json_path, "w") as f:
            json.dump(config, f)
    return _dummy_gemma2_embedding_path
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173


# Add the flag `--optional` to allow run tests
# that are marked with @pytest.mark.optional
def pytest_addoption(parser):
    parser.addoption("--optional",
                     action="store_true",
                     default=False,
                     help="run optional test")


def pytest_collection_modifyitems(config, items):
    if config.getoption("--optional"):
        # --optional given in cli: do not skip optional tests
        return
    skip_optional = pytest.mark.skip(reason="need --optional option to run")
    for item in items:
        if "optional" in item.keywords:
            item.add_marker(skip_optional)