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

4
5
6
from abc import ABC, abstractmethod
from collections import UserDict
from collections.abc import Callable, Iterator, Mapping, Sequence
7
8
from typing import (TYPE_CHECKING, Any, Generic, Literal, NamedTuple, Optional,
                    TypeVar, Union)
9
10
11
12
13

import numpy as np
import torch
from typing_extensions import TypeAlias, TypeGuard, assert_never

14
from vllm.utils import LazyLoader, is_list_of
15

16
from .audio import AudioResampler
17
from .inputs import (AudioItem, HfAudioItem, HfImageItem, HfVideoItem,
18
                     ImageItem, ModalityData, MultiModalDataDict,
19
                     MultiModalFieldConfig, MultiModalKwargsItems, VideoItem)
20
21
22
23

_T = TypeVar("_T")
_I = TypeVar("_I")

24
25
26
27
28
if TYPE_CHECKING:
    import PIL.Image as PILImage
else:
    PILImage = LazyLoader("PILImage", globals(), "PIL.Image")

29
30

class ModalityDataItems(ABC, Generic[_T, _I]):
31
    """
32
33
    Represents data items for a modality in
    [`MultiModalDataItems`][vllm.multimodal.parse.MultiModalDataItems].
34
    """
35

36
    def __init__(self, data: _T, modality: str) -> None:
37
38
39
        super().__init__()

        self.data = data
40
41
42
43
44
        self.modality = modality

    def __repr__(self) -> str:
        return (f"{type(self).__name__}(modality={self.modality!r}, "
                f"len={len(self)})")
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

    def __len__(self) -> int:
        return self.get_count()

    def __getitem__(self, index: int) -> _I:
        return self.get(index)

    if TYPE_CHECKING:
        # Auto-generated
        def __iter__(self) -> Iterator[_I]:
            ...

    @abstractmethod
    def get_count(self) -> int:
        """Get the number of data items."""
        raise NotImplementedError

    @abstractmethod
    def get(self, index: int) -> _I:
        """Get a data item by its index."""
        raise NotImplementedError

    def get_all(self) -> list[_I]:
        """Get all data items."""
        return [self.get(idx) for idx in range(self.get_count())]

    @abstractmethod
    def get_processor_data(self) -> Mapping[str, object]:
        """Get the data to pass to the HF processor."""
        raise NotImplementedError

    @abstractmethod
    def get_passthrough_data(self) -> Mapping[str, object]:
        """Get the data to pass directly to the model."""
        raise NotImplementedError


class ProcessorBatchItems(ModalityDataItems[Sequence[_T], _T]):
83
    """Base class for data items that are arranged in a list."""
84
85
86
87
88
89
90
91
92
93
94
95
96
97

    def get_count(self) -> int:
        return len(self.data)

    def get(self, index: int) -> _T:
        return self.data[index]

    def get_processor_data(self) -> Mapping[str, object]:
        return {f"{self.modality}s": self.data}

    def get_passthrough_data(self) -> Mapping[str, object]:
        return {}


98
99
100
101
102
103
class EmbeddingItems(ModalityDataItems[Union[torch.Tensor, list[torch.Tensor]],
                                       torch.Tensor]):
    """
    Base class for data items that are expressed as a batched embedding tensor,
    or a list of embedding tensors (one per item).
    """
104
105
106
107

    def get_count(self) -> int:
        return len(self.data)

108
    def get(self, index: int) -> torch.Tensor:
109
110
111
112
113
114
115
116
        return self.data[index]

    def get_processor_data(self) -> Mapping[str, object]:
        return {}

    def get_passthrough_data(self) -> Mapping[str, object]:
        return {f"{self.modality}_embeds": self.data}

117
118
119
    def get_feature_size(self, item_idx: int) -> int:
        return len(self.get(item_idx))

120

121
122
123
124
125
126
127
128
129
130
131
132
133
class DictEmbeddingItems(ModalityDataItems[Mapping[str, torch.Tensor],
                                           Mapping[str, torch.Tensor]]):
    """
    Base class for data items that are expressed as a dictionary of tensors.

    Usually, the dictionary keys correspond to the outputs of HF processor.
    """

    def __init__(
        self,
        data: Mapping[str, torch.Tensor],
        modality: str,
        required_fields: set[str],
134
135
136
137
        fields_factory: Callable[
            [Mapping[str, torch.Tensor]],
            Mapping[str, MultiModalFieldConfig],
        ],
138
    ) -> None:
139
140
        from transformers.feature_extraction_utils import BatchFeature

141
142
143
144
145
146
147
148
149
        super().__init__(data, modality)

        missing_required_data_keys = required_fields - data.keys()
        if missing_required_data_keys:
            data_keys = set(data.keys())
            msg = (f"The data should contain the fields: {required_fields}, "
                   f"but only found the following keys: {data_keys}")
            raise ValueError(msg)

150
151
152
153
154
155
156
        fields_config = fields_factory(data)
        missing_required_fields = required_fields - fields_config.keys()
        if missing_required_fields:
            fields = set(fields_config.keys())
            msg = f"{required_fields=} should be a subset of {fields=}"
            raise ValueError(msg)

157
158
159
        self.fields_config = fields_config
        self.required_fields = required_fields

160
        self._kwargs = MultiModalKwargsItems.from_hf_inputs(
161
162
163
164
165
            BatchFeature(dict(data)),
            fields_config,
        )

    def get_count(self) -> int:
166
        return len(self._kwargs[self.modality])
167
168

    def get(self, index: int) -> Mapping[str, torch.Tensor]:
169
        return self._kwargs[self.modality][index].get_data()
170
171
172
173
174
175
176
177

    def get_processor_data(self) -> Mapping[str, object]:
        return {}

    def get_passthrough_data(self) -> Mapping[str, object]:
        return self.data


178
179
180
181
182
class AudioProcessorItems(ProcessorBatchItems[HfAudioItem]):

    def __init__(self, data: Sequence[HfAudioItem]) -> None:
        super().__init__(data, "audio")

183
184
185
186
    def get_audio_length(self, item_idx: int) -> int:
        audio = self.get(item_idx)
        return len(audio)

187
188
189

class AudioEmbeddingItems(EmbeddingItems):

190
    def __init__(self, data: Union[torch.Tensor, list[torch.Tensor]]) -> None:
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
        super().__init__(data, "audio")


class ImageSize(NamedTuple):
    width: int
    height: int


class ImageProcessorItems(ProcessorBatchItems[HfImageItem]):

    def __init__(self, data: Sequence[HfImageItem]) -> None:
        super().__init__(data, "image")

    def get_image_size(self, item_idx: int) -> ImageSize:
        image = self.get(item_idx)

207
        if isinstance(image, PILImage.Image):
208
209
210
211
212
213
214
215
216
217
            return ImageSize(*image.size)
        if isinstance(image, (np.ndarray, torch.Tensor)):
            _, h, w = image.shape
            return ImageSize(w, h)

        assert_never(image)


class ImageEmbeddingItems(EmbeddingItems):

218
    def __init__(self, data: Union[torch.Tensor, list[torch.Tensor]]) -> None:
219
220
221
222
223
        super().__init__(data, "image")


class VideoProcessorItems(ProcessorBatchItems[HfVideoItem]):

224
225
226
227
228
229
    def __init__(
        self,
        data: Sequence[HfVideoItem],
        metadata: Optional[Union[dict[str, Any],
                                 list[Optional[dict[str, Any]]]]] = None,
    ) -> None:
230
        super().__init__(data, "video")
231
        self.metadata = metadata
232

233
234
235
236
237
238
    def get_num_frames(self, item_idx: int) -> int:
        return len(self.get(item_idx))

    def get_frame_size(self, item_idx: int) -> ImageSize:
        image = self.get(item_idx)[0]  # Assume that the video isn't empty

239
        if isinstance(image, PILImage.Image):
240
241
242
243
244
245
246
            return ImageSize(*image.size)
        if isinstance(image, (np.ndarray, torch.Tensor)):
            _, h, w = image.shape
            return ImageSize(w, h)

        assert_never(image)

247
248
249

class VideoEmbeddingItems(EmbeddingItems):

250
    def __init__(self, data: Union[torch.Tensor, list[torch.Tensor]]) -> None:
251
252
253
254
255
256
257
258
        super().__init__(data, "video")


_D = TypeVar("_D", bound=ModalityDataItems[Any, Any])


class MultiModalDataItems(UserDict[str, ModalityDataItems[Any, Any]]):
    """
259
260
    As [`MultiModalDataDict`][vllm.multimodal.inputs.MultiModalDataDict], but
    normalized such that each entry corresponds to a list.
261
262
263
264
265
    """

    def get_count(self, modality: str, *, strict: bool = True) -> int:
        """
        Get the number of data items belonging to a modality.
266

267
        If `strict=False`, return `0` instead of raising [`KeyError`][]
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
        even if the modality is not found.
        """
        if modality not in self:
            if strict:
                available_modalities = set(self.keys())
                raise KeyError(f"Modality {modality!r} not found. "
                               f"Available modalities: {available_modalities}")

            return 0

        return self[modality].get_count()

    def get_all_counts(self) -> Mapping[str, int]:
        """Get the number of items belonging to each modality."""
        return {m: items.get_count() for m, items in self.items()}

    def get_items(
        self,
        modality: str,
287
        typ: Union[type[_D], tuple[type[_D], ...]],
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
    ) -> _D:
        """
        Get the data items belonging to a modality,
        requiring that they belong to a certain type.
        """
        if modality not in self:
            available_modalities = set(self.keys())
            raise KeyError(f"Modality {modality!r} not found. "
                           f"Available modalities: {available_modalities}")

        items = self[modality]
        if not isinstance(items, typ):
            raise TypeError(f"Invalid type of data items for {modality=}. "
                            f"Expected type: {typ}, but "
                            f"found type: {type(items)}")

304
        return items  # type: ignore[return-value]
305
306
307


ModalityDataParser: TypeAlias = Callable[[ModalityData[Any]],
308
                                         Optional[ModalityDataItems[Any, Any]]]
309
310
311
312


class MultiModalDataParser:
    """
313
314
    Parses [`MultiModalDataDict`][vllm.multimodal.inputs.MultiModalDataDict]
    into [`MultiModalDataItems`][vllm.multimodal.parse.MultiModalDataItems].
315
316
317
318

    Args:
        target_sr (float, optional): Enables automatic resampling of audio
            items to the model's expected sampling rate.
319
320
    """

321
322
323
324
325
    def __init__(
        self,
        *,
        target_sr: Optional[float] = None,
        audio_resample_method: Literal["librosa", "scipy"] = "librosa",
326
        video_needs_metadata: bool = False,
327
    ) -> None:
328
329
        super().__init__()

330
331
332
333
        self.audio_resampler = AudioResampler(
            target_sr=target_sr,
            method=audio_resample_method,
        )
334
        self.video_needs_metadata = video_needs_metadata
335

336
337
338
    def _is_embeddings(
            self, data: object
    ) -> TypeGuard[Union[torch.Tensor, list[torch.Tensor]]]:
339
340
341
        if isinstance(data, torch.Tensor):
            return data.ndim == 3
        if is_list_of(data, torch.Tensor):
342
343
344
345
346
347
348
349
350
            return data[0].ndim == 2

        return False

    def _is_empty(self, data: object) -> TypeGuard[None]:
        if isinstance(data, list):
            return len(data) == 0
        if isinstance(data, (np.ndarray, torch.Tensor)):
            return data.size == 0
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368

        return False

    def _get_audio_with_sr(
        self,
        audio: AudioItem,
    ) -> tuple[np.ndarray, Optional[float]]:
        if isinstance(audio, tuple):
            return audio
        if isinstance(audio, list):
            return np.array(audio), None
        if isinstance(audio, np.ndarray):
            return audio, None
        if isinstance(audio, torch.Tensor):
            return audio.numpy(), None

        assert_never(audio)

369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
    def _get_video_with_metadata(
        self,
        video: VideoItem,
    ) -> tuple[np.ndarray, Optional[dict[str, Any]]]:
        if isinstance(video, tuple):
            return video
        if isinstance(video, list):
            return np.array(video), None
        if isinstance(video, np.ndarray):
            return video, None
        if isinstance(video, torch.Tensor):
            return video.numpy(), None

        assert_never(video)

384
385
386
    def _parse_audio_data(
        self,
        data: ModalityData[AudioItem],
387
388
389
390
391
392
    ) -> Optional[ModalityDataItems[Any, Any]]:
        # also check single audio item with sampling rate
        if self._is_empty(data) or (isinstance(data, tuple)
                                    and self._is_empty(data[0])):
            return None

393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
        if self._is_embeddings(data):
            return AudioEmbeddingItems(data)

        if (is_list_of(data, float)
                or isinstance(data,
                              (np.ndarray, torch.Tensor)) and data.ndim == 1
                or isinstance(data, tuple)):
            data_items = [data]
        elif isinstance(data, (np.ndarray, torch.Tensor)):
            data_items = [elem for elem in data]
        else:
            data_items = data

        new_audios = list[np.ndarray]()
        for data_item in data_items:
            audio, orig_sr = self._get_audio_with_sr(data_item)
            if orig_sr is None:
                new_audio = audio
            else:
412
413
                new_audio = self.audio_resampler.resample(audio,
                                                          orig_sr=orig_sr)
414
415
416
417
418
419
420
421

            new_audios.append(new_audio)

        return AudioProcessorItems(new_audios)

    def _parse_image_data(
        self,
        data: ModalityData[ImageItem],
422
423
424
425
    ) -> Optional[ModalityDataItems[Any, Any]]:
        if self._is_empty(data):
            return None

426
427
428
        if self._is_embeddings(data):
            return ImageEmbeddingItems(data)

429
        if (isinstance(data, PILImage.Image)
430
431
432
433
434
435
436
437
438
439
440
441
442
                or isinstance(data,
                              (np.ndarray, torch.Tensor)) and data.ndim == 3):
            data_items = [data]
        elif isinstance(data, (np.ndarray, torch.Tensor)):
            data_items = [elem for elem in data]
        else:
            data_items = data

        return ImageProcessorItems(data_items)

    def _parse_video_data(
        self,
        data: ModalityData[VideoItem],
443
444
445
446
    ) -> Optional[ModalityDataItems[Any, Any]]:
        if self._is_empty(data):
            return None

447
448
449
        if self._is_embeddings(data):
            return VideoEmbeddingItems(data)

450
        if (is_list_of(data, PILImage.Image)
451
452
453
454
455
                or isinstance(data,
                              (np.ndarray, torch.Tensor)) and data.ndim == 4):
            data_items = [data]
        elif isinstance(data, (np.ndarray, torch.Tensor)):
            data_items = [elem for elem in data]
456
457
        elif isinstance(data, tuple) and len(data) == 2:
            data_items = [data]
458
459
460
        else:
            data_items = data

461
462
463
464
465
466
467
468
469
470
471
472
473
474
        new_videos = list[tuple[np.ndarray, Optional[dict[str, Any]]]]()
        metadata_lst: list[Optional[dict[str, Any]]] = []
        for data_item in data_items:
            video, metadata = self._get_video_with_metadata(data_item)
            if self.video_needs_metadata:
                new_videos.append((video, metadata))
                metadata_lst.append(metadata)
            else:
                new_videos.append(video)

        if not self.video_needs_metadata:
            metadata = None

        return VideoProcessorItems(new_videos, metadata=metadata_lst)
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491

    def _get_subparsers(self) -> Mapping[str, ModalityDataParser]:
        return {
            "audio": self._parse_audio_data,
            "image": self._parse_image_data,
            "video": self._parse_video_data,
        }

    def parse_mm_data(self,
                      mm_data: MultiModalDataDict) -> MultiModalDataItems:
        subparsers = self._get_subparsers()

        mm_items = MultiModalDataItems()
        for k, v in mm_data.items():
            if k not in subparsers:
                raise ValueError(f"Unsupported modality: {k}")

492
493
494
            # ignore empty embedding data
            if (parsed_data := subparsers[k](v)) is not None:
                mm_items[k] = parsed_data
495
496

        return mm_items