registry.py 6.44 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import time
4
from collections.abc import Mapping
5
from dataclasses import dataclass
6
from typing import TYPE_CHECKING, Any, Union
7

8
import torch
9
from transformers import BatchFeature, PretrainedConfig, ProcessorMixin
10
from typing_extensions import TypeVar
11

12
from vllm.logger import init_logger
13
from vllm.transformers_utils.processor import cached_processor_from_config
14
from vllm.utils import get_allowed_kwarg_only_overrides
15
from vllm.utils.jsontree import JSONTree, json_map_leaves
16
17

if TYPE_CHECKING:
18
    from vllm.config import ModelConfig
19
20
21
22
    from vllm.transformers_utils.tokenizer import AnyTokenizer
else:
    ModelConfig = Any
    AnyTokenizer = Any
23

24
25
26
_T = TypeVar("_T")
_C = TypeVar("_C", bound=PretrainedConfig, default=PretrainedConfig)
_P = TypeVar("_P", bound=ProcessorMixin, default=ProcessorMixin)
27

28
29
logger = init_logger(__name__)

30

31
32
33
34
35
36
37
@dataclass(frozen=True)
class InputContext:
    """
    Contains information about the model which may be used to
    modify the inputs.
    """

38
    model_config: ModelConfig
39
40
    """The configuration of the model."""

41
42
    def get_hf_config(
        self,
43
        typ: Union[type[_C], tuple[type[_C], ...]] = PretrainedConfig,
44
        /,
45
    ) -> _C:
46
47
        """
        Get the HuggingFace configuration
48
        (`transformers.PretrainedConfig`) of the model,
49
50
51
        additionally checking its type.

        Raises:
52
            TypeError: If the configuration is not of the specified type.
53
54
        """
        hf_config = self.model_config.hf_config
55
        if not isinstance(hf_config, typ):
56
            raise TypeError("Invalid type of HuggingFace config. "
57
                            f"Expected type: {typ}, but "
58
59
60
61
                            f"found type: {type(hf_config)}")

        return hf_config

62
    def get_hf_image_processor_config(self) -> dict[str, Any]:
63
64
65
66
67
        """
        Get the HuggingFace image processor configuration of the model.
        """
        return self.model_config.hf_image_processor_config

68
69
70
71
72
73
74
75
76
77
78
79
80
    def get_mm_config(self):
        """
        Get the multimodal config of the model.

        Raises:
            RuntimeError: If the model is not a multimodal model.
        """
        mm_config = self.model_config.multimodal_config
        if mm_config is None:
            raise RuntimeError("Not a multimodal model")

        return mm_config

81
82
    def get_hf_processor(
        self,
83
        typ: Union[type[_P], tuple[type[_P], ...]] = ProcessorMixin,
84
85
        /,
        **kwargs: object,
86
    ) -> _P:
87
88
        """
        Get the HuggingFace processor
89
        (`transformers.ProcessorMixin`) of the model,
90
91
92
93
94
        additionally checking its type.

        Raises:
            TypeError: If the processor is not of the specified type.
        """
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
        return cached_processor_from_config(
            self.model_config,
            processor_cls=typ,
            **kwargs,
        )

    def init_processor(
        self,
        typ: type[_T],
        /,
        **kwargs: object,
    ) -> _T:
        """
        Initialize a HuggingFace-like processor class, merging the
        keyword arguments with those in the model's configuration.
        """
111
112
        mm_config = self.model_config.get_multimodal_config()
        base_kwargs = mm_config.mm_processor_kwargs
113
114
115
116
117
        if base_kwargs is None:
            base_kwargs = {}

        merged_kwargs = {**base_kwargs, **kwargs}

118
        return typ(**merged_kwargs)
119

120

121
122
123
124
125
@dataclass(frozen=True)
class InputProcessingContext(InputContext):
    tokenizer: AnyTokenizer
    """The tokenizer used to tokenize the inputs."""

126
127
    def get_hf_processor(
        self,
128
        typ: Union[type[_P], tuple[type[_P], ...]] = ProcessorMixin,
129
130
        /,
        **kwargs: object,
131
    ) -> _P:
132
133
        return super().get_hf_processor(
            typ,
134
            tokenizer=self.tokenizer,
135
            **kwargs,
136
137
        )

138
    def call_hf_processor(
139
140
        self,
        hf_processor: ProcessorMixin,
141
142
        data: Mapping[str, object],
        kwargs: Mapping[str, object] = {},
143
144
145
        *,
        num_tries: int = 1,
        max_tries: int = 5,
146
    ) -> Union[BatchFeature, JSONTree]:
147
        """
148
149
        Call `hf_processor` on the prompt `data`
        (text, image, audio...) with configurable options `kwargs`.
150
        """
151
152
        assert callable(hf_processor)

153
        mm_config = self.model_config.get_multimodal_config()
154
        merged_kwargs = mm_config.merge_mm_processor_kwargs(kwargs)
155

156
        allowed_kwargs = get_allowed_kwarg_only_overrides(
157
            hf_processor,
158
            merged_kwargs,
159
160
            requires_kw_only=False,
            allow_var_kwargs=True,
161
        )
162

163
164
165
166
167
168
        def maybe_cast_dtype(x):
            # This mimics the behavior of transformers.BatchFeature
            if isinstance(x, torch.Tensor) and x.is_floating_point():
                return x.to(dtype=self.model_config.dtype)
            return x

169
        try:
170
171
172
            output = hf_processor(**data,
                                  **allowed_kwargs,
                                  return_tensors="pt")
173
174
            # this emulates output.to(dtype=self.model_config.dtype)
            if isinstance(output, BatchFeature):
175
                cast_output = json_map_leaves(maybe_cast_dtype, output.data)
176
177
                return BatchFeature(cast_output)

178
179
            cast_output = json_map_leaves(maybe_cast_dtype, output)

180
181
182
183
184
185
            logger.warning_once(
                f"{type(hf_processor).__name__} did not return `BatchFeature`. "
                "Make sure to match the behaviour of `ProcessorMixin` when "
                "implementing custom processors.")
            return cast_output

186
        except Exception as exc:
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
            # See https://github.com/huggingface/tokenizers/issues/537
            if (isinstance(exc, RuntimeError) and exc
                    and exc.args[0] == "Already borrowed"
                    and num_tries < max_tries):
                logger.warning(
                    "Failed to acquire tokenizer in current thread. "
                    "Retrying (%d/%d)...", num_tries, max_tries)
                time.sleep(0.5)
                return self.call_hf_processor(
                    hf_processor,
                    data,
                    kwargs,
                    num_tries=num_tries + 1,
                    max_tries=max_tries,
                )

203
            msg = (f"Failed to apply {type(hf_processor).__name__} "
204
                   f"on data={data} with kwargs={allowed_kwargs}")
205

206
            raise ValueError(msg) from exc