| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import pytest |
| 4 | import torch |
| 5 | |
| 6 | from vllm.multimodal.inputs import PlaceholderRange |
| 7 | |
| 8 | |
| 9 | @pytest.mark.parametrize( |
| 10 | "is_embed,expected", |
| 11 | [ |
| 12 | (None, 5), |
| 13 | (torch.tensor([True, True, True, True, True]), 5), |
| 14 | (torch.tensor([False, False, False, False, False]), 0), |
| 15 | (torch.tensor([True, False, True, False, True]), 3), |
| 16 | (torch.tensor([True]), 1), |
| 17 | ], |
| 18 | ) |
| 19 | def test_placeholder_range_get_num_embeds(is_embed, expected): |
| 20 | length = len(is_embed) if is_embed is not None else 5 |
| 21 | pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed) |
| 22 | assert pr.get_num_embeds() == expected |
| 23 | |
| 24 | |
| 25 | @pytest.mark.parametrize( |
| 26 | "is_embed,expected", |
| 27 | [ |
| 28 | (None, None), |
| 29 | (torch.tensor([False, True, False, True, True]), [0, 1, 1, 2, 3]), |
| 30 | (torch.tensor([True, True, True]), [1, 2, 3]), |
| 31 | ], |
| 32 | ) |
| 33 | def test_placeholder_range_embeds_cumsum(is_embed, expected): |
| 34 | length = len(is_embed) if is_embed is not None else 5 |
| 35 | pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed) |
| 36 | |
| 37 | if expected is None: |
| 38 | assert pr.embeds_cumsum is None |
| 39 | return |
| 40 | |
| 41 | assert pr.embeds_cumsum == expected |
| 42 | # cached_property should return the same object on repeated access |
| 43 | assert pr.embeds_cumsum is pr.embeds_cumsum |