| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm.model_executor.models.ernie45_vl import ( |
| 9 | Ernie4_5_VLMoeForConditionalGeneration, |
| 10 | ) |
| 11 | from vllm.multimodal.inputs import ( |
| 12 | MultiModalFeatureSpec, |
| 13 | MultiModalFieldElem, |
| 14 | MultiModalKwargsItem, |
| 15 | PlaceholderRange, |
| 16 | ) |
| 17 | |
| 18 | pytestmark = pytest.mark.skip_global_cleanup |
| 19 | |
| 20 | |
| 21 | @pytest.fixture(autouse=True, scope="module") |
| 22 | def _force_cpu_default_device(): |
| 23 | original = torch.get_default_device() |
| 24 | torch.set_default_device("cpu") |
| 25 | yield |
| 26 | torch.set_default_device(original) |
| 27 | |
| 28 | |
| 29 | @dataclass |
| 30 | class DummyConfig: |
| 31 | spatial_conv_size: int = 2 |
| 32 | temporal_conv_size: int = 2 |
| 33 | |
| 34 | |
| 35 | def make_model(config: DummyConfig) -> Ernie4_5_VLMoeForConditionalGeneration: |
| 36 | model = object.__new__(Ernie4_5_VLMoeForConditionalGeneration) |
| 37 | model.config = config |
| 38 | return model |
| 39 | |
| 40 | |
| 41 | def make_mm_feature( |
| 42 | *, |
| 43 | modality: str, |
| 44 | offset: int, |
| 45 | length: int, |
| 46 | grid_thw: tuple[int, int, int], |
| 47 | ) -> MultiModalFeatureSpec: |
| 48 | field_name = "image_grid_thw" if modality == "image" else "video_grid_thw" |
| 49 | return MultiModalFeatureSpec( |
| 50 | data=MultiModalKwargsItem( |
| 51 | { |
| 52 | field_name: MultiModalFieldElem( |
| 53 | data=torch.tensor(grid_thw), |
| 54 | field=None, # HACK. |
| 55 | ), |
| 56 | } |
| 57 | ), |
| 58 | modality=modality, |
| 59 | identifier="DUMMY", |
| 60 | mm_position=PlaceholderRange(offset=offset, length=length), |
| 61 | ) |
| 62 | |
| 63 | |
| 64 | def test_get_mrope_input_positions_text_only(): |
| 65 | model = make_model(DummyConfig()) |
| 66 | |
| 67 | positions, delta = model.get_mrope_input_positions( |
| 68 | input_tokens=[11, 12, 13, 14, 15], |
| 69 | mm_features=[], |
| 70 | ) |
| 71 | |
| 72 | expected = torch.tensor( |
| 73 | [ |
| 74 | [0, 1, 2, 3, 4], |
| 75 | [0, 1, 2, 3, 4], |
| 76 | [0, 1, 2, 3, 4], |
| 77 | ] |
| 78 | ) |
| 79 | |
| 80 | assert torch.equal(positions, expected) |
| 81 | assert delta == 0 |
| 82 | |
| 83 | |
| 84 | def test_get_mrope_input_positions_single_image(): |
| 85 | model = make_model(DummyConfig()) |
| 86 | mm_features = [ |
| 87 | make_mm_feature( |
| 88 | modality="image", |
| 89 | offset=1, |
| 90 | length=4, |
| 91 | grid_thw=(1, 4, 4), |
| 92 | ) |
| 93 | ] |
| 94 | |
| 95 | positions, delta = model.get_mrope_input_positions( |
| 96 | input_tokens=[10, 20, 21, 22, 23, 30, 31], |
| 97 | mm_features=mm_features, |
| 98 | ) |
| 99 | |
| 100 | expected = torch.tensor( |
| 101 | [ |
| 102 | [0, 1, 1, 1, 1, 3, 4], |
| 103 | [0, 1, 1, 2, 2, 3, 4], |
| 104 | [0, 1, 2, 1, 2, 3, 4], |
| 105 | ] |
| 106 | ) |
| 107 | |
| 108 | assert torch.equal(positions, expected) |
| 109 | assert delta == -2 |
| 110 | |
| 111 | |
| 112 | def test_get_mrope_input_positions_interleaved_image_and_video(): |
| 113 | model = make_model(DummyConfig()) |
| 114 | mm_features = [ |
| 115 | make_mm_feature( |
| 116 | modality="image", |
| 117 | offset=1, |
| 118 | length=4, |
| 119 | grid_thw=(1, 4, 4), |
| 120 | ), |
| 121 | make_mm_feature( |
| 122 | modality="video", |
| 123 | offset=7, |
| 124 | length=2, |
| 125 | grid_thw=(2, 4, 2), |
| 126 | ), |
| 127 | ] |
| 128 | |
| 129 | positions, delta = model.get_mrope_input_positions( |
| 130 | input_tokens=[10, 20, 21, 22, 23, 30, 31, 40, 41, 50, 51], |
| 131 | mm_features=mm_features, |
| 132 | ) |
| 133 | |
| 134 | expected = torch.tensor( |
| 135 | [ |
| 136 | [0, 1, 1, 1, 1, 3, 4, 5, 5, 7, 8], |
| 137 | [0, 1, 1, 2, 2, 3, 4, 5, 6, 7, 8], |
| 138 | [0, 1, 2, 1, 2, 3, 4, 5, 5, 7, 8], |
| 139 | ] |
| 140 | ) |
| 141 | |
| 142 | assert torch.equal(positions, expected) |
| 143 | assert delta == -2 |