| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | from torch import nn |
| 9 | |
| 10 | from vllm.model_executor.models.openpangu_vl import ( |
| 11 | OpenPanguVLForConditionalGeneration, |
| 12 | ) |
| 13 | |
| 14 | |
| 15 | class _LanguageModel(nn.Module): |
| 16 | def __init__(self, hidden_size: int) -> None: |
| 17 | super().__init__() |
| 18 | self.hidden_size = hidden_size |
| 19 | |
| 20 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 21 | return input_ids.float().unsqueeze(-1).expand(-1, self.hidden_size).clone() |
| 22 | |
| 23 | |
| 24 | class _OpenPanguVL(OpenPanguVLForConditionalGeneration): |
| 25 | def __init__(self, hidden_size: int) -> None: |
| 26 | nn.Module.__init__(self) |
| 27 | self.config = SimpleNamespace(image_token_id=41, video_token_id=42) |
| 28 | self.language_model = _LanguageModel(hidden_size) |
| 29 | |
| 30 | def get_language_model(self) -> _LanguageModel: |
| 31 | return self.language_model |
| 32 | |
| 33 | |
| 34 | @pytest.mark.parametrize("device_type", ["cpu", "cuda"]) |
| 35 | def test_get_input_embeddings_merges_image_and_video_embeddings( |
| 36 | device_type: str, |
| 37 | ) -> None: |
| 38 | if device_type == "cuda" and not torch.cuda.is_available(): |
| 39 | pytest.skip("CUDA is not available") |
| 40 | |
| 41 | device = torch.device(device_type) |
| 42 | hidden_size = 4 |
| 43 | model = _OpenPanguVL(hidden_size).to(device) |
| 44 | input_ids = torch.tensor([7, 41, 41, 8, 42, 9], device=device) |
| 45 | multimodal_embeddings = ( |
| 46 | torch.full((2, hidden_size), -1.0, device=device), |
| 47 | torch.full((1, hidden_size), -2.0, device=device), |
| 48 | ) |
| 49 | |
| 50 | actual = model.get_input_embeddings(input_ids, multimodal_embeddings) |
| 51 | |
| 52 | expected = model.language_model.embed_input_ids(input_ids) |
| 53 | expected[1:3] = -1.0 |
| 54 | expected[4] = -2.0 |
| 55 | torch.testing.assert_close(actual, expected) |