test_inputs.py 3.1 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import pytest
5
6
import torch

7
from vllm.multimodal.inputs import MultiModalKwargs, NestedTensors
8

9
10
pytestmark = pytest.mark.cpu_test

11

12
def assert_nested_tensors_equal(expected: NestedTensors, actual: NestedTensors):
13
    assert type(expected) == type(actual)  # noqa: E721
14
15
16
17
18
19
20
    if isinstance(expected, torch.Tensor):
        assert torch.equal(expected, actual)
    else:
        for expected_item, actual_item in zip(expected, actual):
            assert_nested_tensors_equal(expected_item, actual_item)


21
22
23
def assert_multimodal_inputs_equal(
    expected: MultiModalKwargs, actual: MultiModalKwargs
):
24
25
26
27
28
29
30
    assert set(expected.keys()) == set(actual.keys())
    for key in expected:
        assert_nested_tensors_equal(expected[key], actual[key])


def test_multimodal_input_batch_single_tensor():
    t = torch.rand([1, 2])
31
    result = MultiModalKwargs.batch([{"image": t}])
32
33
34
35
36
37
38
    assert_multimodal_inputs_equal(result, {"image": t.unsqueeze(0)})


def test_multimodal_input_batch_multiple_tensors():
    a = torch.rand([1, 1, 2])
    b = torch.rand([1, 1, 2])
    c = torch.rand([1, 1, 2])
39
    result = MultiModalKwargs.batch([{"image": a}, {"image": b}, {"image": c}])
40
41
42
43
44
45
46
    assert_multimodal_inputs_equal(result, {"image": torch.stack([a, b, c])})


def test_multimodal_input_batch_multiple_heterogeneous_tensors():
    a = torch.rand([1, 2, 2])
    b = torch.rand([1, 3, 2])
    c = torch.rand([1, 4, 2])
47
    result = MultiModalKwargs.batch([{"image": a}, {"image": b}, {"image": c}])
48
49
50
51
52
53
54
    assert_multimodal_inputs_equal(result, {"image": [a, b, c]})


def test_multimodal_input_batch_nested_tensors():
    a = torch.rand([2, 3])
    b = torch.rand([2, 3])
    c = torch.rand([2, 3])
55
56
57
58
    result = MultiModalKwargs.batch([{"image": [a]}, {"image": [b]}, {"image": [c]}])
    assert_multimodal_inputs_equal(
        result, {"image": torch.stack([a.unsqueeze(0), b.unsqueeze(0), c.unsqueeze(0)])}
    )
59
60
61
62
63
64


def test_multimodal_input_batch_heterogeneous_lists():
    a = torch.rand([1, 2, 3])
    b = torch.rand([1, 2, 3])
    c = torch.rand([1, 2, 3])
65
    result = MultiModalKwargs.batch([{"image": [a, b]}, {"image": [c]}])
66
    assert_multimodal_inputs_equal(
67
68
        result, {"image": [torch.stack([a, b]), c.unsqueeze(0)]}
    )
69
70
71
72
73
74
75


def test_multimodal_input_batch_multiple_batchable_lists():
    a = torch.rand([1, 2, 3])
    b = torch.rand([1, 2, 3])
    c = torch.rand([1, 2, 3])
    d = torch.rand([1, 2, 3])
76
    result = MultiModalKwargs.batch([{"image": [a, b]}, {"image": [c, d]}])
77
    assert_multimodal_inputs_equal(
78
79
        result, {"image": torch.stack([torch.stack([a, b]), torch.stack([c, d])])}
    )
80
81
82
83
84
85
86


def test_multimodal_input_batch_mixed_stacking_depths():
    a = torch.rand([1, 2, 3])
    b = torch.rand([1, 3, 3])
    c = torch.rand([1, 4, 3])

87
    result = MultiModalKwargs.batch([{"image": [a, b]}, {"image": [c]}])
88
89
    assert_multimodal_inputs_equal(result, {"image": [[a, b], c.unsqueeze(0)]})

90
    result = MultiModalKwargs.batch([{"image": [a]}, {"image": [b, c]}])
91
    assert_multimodal_inputs_equal(result, {"image": [a.unsqueeze(0), [b, c]]})