chat_utils.py 20.3 KB
Newer Older
1
import asyncio
2
import codecs
3
import json
4
from abc import ABC, abstractmethod
5
from collections import defaultdict
6
from functools import lru_cache, partial
7
from pathlib import Path
8
from typing import (Any, Awaitable, Dict, Generic, Iterable, List, Literal,
9
                    Mapping, Optional, Tuple, TypeVar, Union, cast)
10

11
12
# yapf conflicts with isort for this block
# yapf: disable
13
14
from openai.types.chat import (ChatCompletionAssistantMessageParam,
                               ChatCompletionContentPartImageParam)
15
16
from openai.types.chat import (
    ChatCompletionContentPartParam as OpenAIChatCompletionContentPartParam)
17
18
from openai.types.chat import (ChatCompletionContentPartRefusalParam,
                               ChatCompletionContentPartTextParam)
19
20
from openai.types.chat import (
    ChatCompletionMessageParam as OpenAIChatCompletionMessageParam)
21
22
from openai.types.chat import (ChatCompletionMessageToolCallParam,
                               ChatCompletionToolMessageParam)
23
24
# yapf: enable
# pydantic needs the TypedDict from typing_extensions
25
from pydantic import ConfigDict
26
from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
27
from typing_extensions import Required, TypeAlias, TypedDict
28

29
from vllm.config import ModelConfig
30
31
from vllm.logger import init_logger
from vllm.multimodal import MultiModalDataDict
32
from vllm.multimodal.utils import (async_get_and_parse_audio,
33
34
                                   async_get_and_parse_image,
                                   get_and_parse_audio, get_and_parse_image)
35
from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer
36
from vllm.utils import print_warning_once
37
38
39
40

logger = init_logger(__name__)


41
42
43
44
45
46
47
48
49
50
51
52
53
54
class AudioURL(TypedDict, total=False):
    url: Required[str]
    """
    Either a URL of the audio or a data URL with base64 encoded audio data.
    """


class ChatCompletionContentPartAudioParam(TypedDict, total=False):
    audio_url: Required[AudioURL]

    type: Required[Literal["audio_url"]]
    """The type of the content part."""


55
56
57
58
59
60
61
class CustomChatCompletionContentPartParam(TypedDict, total=False):
    __pydantic_config__ = ConfigDict(extra="allow")  # type: ignore

    type: Required[str]
    """The type of the content part."""


62
63
ChatCompletionContentPartParam: TypeAlias = Union[
    OpenAIChatCompletionContentPartParam, ChatCompletionContentPartAudioParam,
64
65
    ChatCompletionContentPartRefusalParam,
    CustomChatCompletionContentPartParam]
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82


class CustomChatCompletionMessageParam(TypedDict, total=False):
    """Enables custom roles in the Chat Completion API."""
    role: Required[str]
    """The role of the message's author."""

    content: Union[str, List[ChatCompletionContentPartParam]]
    """The contents of the message."""

    name: str
    """An optional name for the participant.

    Provides the model information to differentiate between participants of the
    same role.
    """

83
84
85
86
87
88
    tool_call_id: Optional[str]
    """Tool call that this message is responding to."""

    tool_calls: Optional[Iterable[ChatCompletionMessageToolCallParam]]
    """The tool calls generated by the model, such as function calls."""

89
90
91
92
93

ChatCompletionMessageParam = Union[OpenAIChatCompletionMessageParam,
                                   CustomChatCompletionMessageParam]


94
# TODO: Make fields ReadOnly once mypy supports it
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
class ConversationMessage(TypedDict, total=False):
    role: Required[str]
    """The role of the message's author."""

    content: Optional[str]
    """The contents of the message"""

    tool_call_id: Optional[str]
    """Tool call that this message is responding to."""

    name: Optional[str]
    """The name of the function to call"""

    tool_calls: Optional[Iterable[ChatCompletionMessageToolCallParam]]
    """The tool calls generated by the model, such as function calls."""
110
111


112
ModalityStr = Literal["image", "audio", "video"]
113
114
115
116
_T = TypeVar("_T")


class BaseMultiModalItemTracker(ABC, Generic[_T]):
117
118
119
120
121
122
123
    """
    Tracks multi-modal items in a given request and ensures that the number
    of multi-modal items in a given request does not exceed the configured
    maximum per prompt.
    """

    def __init__(self, model_config: ModelConfig, tokenizer: AnyTokenizer):
124
125
        super().__init__()

126
127
128
129
130
        self._model_config = model_config
        self._tokenizer = tokenizer
        self._allowed_items = (model_config.multimodal_config.limit_per_prompt
                               if model_config.multimodal_config else {})
        self._consumed_items = {k: 0 for k in self._allowed_items}
131
132

        self._items: List[_T] = []
133
134
135

    @staticmethod
    @lru_cache(maxsize=None)
136
    def _cached_token_str(tokenizer: AnyTokenizer, token_index: int) -> str:
137
138
        return tokenizer.decode(token_index)

139
140
    def _placeholder_str(self, modality: ModalityStr,
                         current_count: int) -> Optional[str]:
141
142
        # TODO: Let user specify how to insert image tokens into prompt
        # (similar to chat template)
143
144
145
        hf_config = self._model_config.hf_config
        model_type = hf_config.model_type

146
147
148
149
150
151
        if modality == "image":
            if model_type == "phi3_v":
                # Workaround since this token is not defined in the tokenizer
                return f"<|image_{current_count}|>"
            if model_type == "minicpmv":
                return "(<image>./</image>)"
Patrick von Platen's avatar
Patrick von Platen committed
152
153
            if model_type in ("blip-2", "chatglm", "fuyu", "paligemma",
                              "pixtral"):
154
155
                # These models do not use image tokens in the prompt
                return None
156
157
            if model_type == "qwen":
                return f"Picture {current_count}: <img></img>"
158
            if model_type.startswith("llava"):
159
160
                return self._cached_token_str(self._tokenizer,
                                              hf_config.image_token_index)
161
            if model_type in ("chameleon", "internvl_chat", "NVLM_D"):
162
                return "<image>"
163
164
            if model_type == "mllama":
                return "<|image|>"
165
166
            if model_type == "qwen2_vl":
                return "<|vision_start|><|image_pad|><|vision_end|>"
167
168
            if model_type == "molmo":
                return ""
169

170
            raise TypeError(f"Unknown {modality} model type: {model_type}")
171
172
173
        elif modality == "audio":
            if model_type == "ultravox":
                return "<|reserved_special_token_0|>"
174
            raise TypeError(f"Unknown {modality} model type: {model_type}")
175
176
177
        elif modality == "video":
            if model_type == "qwen2_vl":
                return "<|vision_start|><|video_pad|><|vision_end|>"
178
            raise TypeError(f"Unknown {modality} model type: {model_type}")
179
180
181
182
        else:
            raise TypeError(f"Unknown modality: {modality}")

    @staticmethod
183
    def _combine(items: List[MultiModalDataDict]) -> MultiModalDataDict:
184
185
186
        mm_lists: Mapping[str, List[object]] = defaultdict(list)

        # Merge all the multi-modal items
187
        for single_mm_data in items:
188
189
190
191
192
193
194
195
196
197
198
199
            for mm_key, mm_item in single_mm_data.items():
                if isinstance(mm_item, list):
                    mm_lists[mm_key].extend(mm_item)
                else:
                    mm_lists[mm_key].append(mm_item)

        # Unpack any single item lists for models that don't expect multiple.
        return {
            mm_key: mm_list[0] if len(mm_list) == 1 else mm_list
            for mm_key, mm_list in mm_lists.items()
        }

200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
    def add(self, modality: ModalityStr, item: _T) -> Optional[str]:
        """
        Add a multi-modal item to the current prompt and returns the
        placeholder string to use, if any.
        """
        allowed_count = self._allowed_items.get(modality, 1)
        current_count = self._consumed_items.get(modality, 0) + 1
        if current_count > allowed_count:
            raise ValueError(
                f"At most {allowed_count} {modality}(s) may be provided in "
                "one request.")

        self._consumed_items[modality] = current_count
        self._items.append(item)

        return self._placeholder_str(modality, current_count)

    @abstractmethod
    def create_parser(self) -> "BaseMultiModalContentParser":
        raise NotImplementedError


class MultiModalItemTracker(BaseMultiModalItemTracker[MultiModalDataDict]):

    def all_mm_data(self) -> Optional[MultiModalDataDict]:
        return self._combine(self._items) if self._items else None

    def create_parser(self) -> "BaseMultiModalContentParser":
        return MultiModalContentParser(self)


class AsyncMultiModalItemTracker(
        BaseMultiModalItemTracker[Awaitable[MultiModalDataDict]]):

    async def all_mm_data(self) -> Optional[MultiModalDataDict]:
        if self._items:
            items = await asyncio.gather(*self._items)
            return self._combine(items)

        return None

    def create_parser(self) -> "BaseMultiModalContentParser":
        return AsyncMultiModalContentParser(self)


class BaseMultiModalContentParser(ABC):

    def __init__(self) -> None:
        super().__init__()

        # multimodal placeholder_string : count
        self._placeholder_counts: Dict[str, int] = defaultdict(lambda: 0)

    def _add_placeholder(self, placeholder: Optional[str]):
        if placeholder:
            self._placeholder_counts[placeholder] += 1

    def mm_placeholder_counts(self) -> Dict[str, int]:
        return dict(self._placeholder_counts)

    @abstractmethod
    def parse_image(self, image_url: str) -> None:
        raise NotImplementedError

    @abstractmethod
    def parse_audio(self, audio_url: str) -> None:
        raise NotImplementedError


class MultiModalContentParser(BaseMultiModalContentParser):

    def __init__(self, tracker: MultiModalItemTracker) -> None:
        super().__init__()

        self._tracker = tracker

    def parse_image(self, image_url: str) -> None:
        image = get_and_parse_image(image_url)

        placeholder = self._tracker.add("image", image)
        self._add_placeholder(placeholder)

    def parse_audio(self, audio_url: str) -> None:
        audio = get_and_parse_audio(audio_url)

        placeholder = self._tracker.add("audio", audio)
        self._add_placeholder(placeholder)


class AsyncMultiModalContentParser(BaseMultiModalContentParser):

    def __init__(self, tracker: AsyncMultiModalItemTracker) -> None:
        super().__init__()

        self._tracker = tracker

    def parse_image(self, image_url: str) -> None:
        image_coro = async_get_and_parse_image(image_url)

        placeholder = self._tracker.add("image", image_coro)
        self._add_placeholder(placeholder)

    def parse_audio(self, audio_url: str) -> None:
        audio_coro = async_get_and_parse_audio(audio_url)

        placeholder = self._tracker.add("audio", audio_coro)
        self._add_placeholder(placeholder)
307
308


309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
def validate_chat_template(chat_template: Optional[Union[Path, str]]):
    """Raises if the provided chat template appears invalid."""
    if chat_template is None:
        return

    elif isinstance(chat_template, Path) and not chat_template.exists():
        raise FileNotFoundError(
            "the supplied chat template path doesn't exist")

    elif isinstance(chat_template, str):
        JINJA_CHARS = "{}\n"
        if not any(c in chat_template
                   for c in JINJA_CHARS) and not Path(chat_template).exists():
            raise ValueError(
                f"The supplied chat template string ({chat_template}) "
                f"appears path-like, but doesn't exist!")

    else:
        raise TypeError(
            f"{type(chat_template)} is not a valid chat template type")


331
332
def load_chat_template(
        chat_template: Optional[Union[Path, str]]) -> Optional[str]:
333
334
335
336
337
338
    if chat_template is None:
        return None
    try:
        with open(chat_template, "r") as f:
            resolved_chat_template = f.read()
    except OSError as e:
339
340
341
        if isinstance(chat_template, Path):
            raise

342
343
344
345
346
347
        JINJA_CHARS = "{}\n"
        if not any(c in chat_template for c in JINJA_CHARS):
            msg = (f"The supplied chat template ({chat_template}) "
                   f"looks like a file path, but it failed to be "
                   f"opened. Reason: {e}")
            raise ValueError(msg) from e
348

349
350
351
        # If opening a file fails, set chat template to be args to
        # ensure we decode so our escape are interpreted correctly
        resolved_chat_template = codecs.decode(chat_template, "unicode_escape")
352

353
354
    logger.info("Using supplied chat template:\n%s", resolved_chat_template)
    return resolved_chat_template
355
356


357
# TODO: Let user specify how to insert multimodal tokens into prompt
358
# (similar to chat template)
359
def _get_full_multimodal_text_prompt(placeholder_counts: Dict[str, int],
360
                                     text_prompt: str) -> str:
361
    """Combine multimodal prompts for a multimodal language model."""
362

363
    # Look through the text prompt to check for missing placeholders
364
    missing_placeholders: List[str] = []
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
    for placeholder in placeholder_counts:

        # For any existing placeholder in the text prompt, we leave it as is
        placeholder_counts[placeholder] -= text_prompt.count(placeholder)

        if placeholder_counts[placeholder] < 0:
            raise ValueError(
                f"Found more '{placeholder}' placeholders in input prompt than "
                "actual multimodal data items.")

        missing_placeholders.extend([placeholder] *
                                    placeholder_counts[placeholder])

    # NOTE: For now we always add missing placeholders at the front of
    # the prompt. This may change to be customizable in the future.
    return "\n".join(missing_placeholders + [text_prompt])
381
382


383
384
385
386
387
# No need to validate using Pydantic again
_TextParser = partial(cast, ChatCompletionContentPartTextParam)
_ImageParser = partial(cast, ChatCompletionContentPartImageParam)
_AudioParser = partial(cast, ChatCompletionContentPartAudioParam)
_RefusalParser = partial(cast, ChatCompletionContentPartRefusalParam)
388
MODEL_KEEP_MULTI_MODAL_CONTENT = {'mllama'}
389
390


391
392
393
def _parse_chat_message_content_parts(
    role: str,
    parts: Iterable[ChatCompletionContentPartParam],
394
    mm_tracker: BaseMultiModalItemTracker,
395
) -> List[ConversationMessage]:
396
    texts: List[str] = []
397

398
    mm_parser = mm_tracker.create_parser()
399
400
401
    keep_multimodal_content = \
        mm_tracker._model_config.hf_config.model_type in \
            MODEL_KEEP_MULTI_MODAL_CONTENT
402

403
    has_image = False
404
405
406
    for part in parts:
        part_type = part["type"]
        if part_type == "text":
407
            text = _TextParser(part)["text"]
408
409
            texts.append(text)
        elif part_type == "image_url":
410
            image_url = _ImageParser(part)["image_url"]
411
412
413
414
415
416

            if image_url.get("detail", "auto") != "auto":
                logger.warning(
                    "'image_url.detail' is currently not supported and "
                    "will be ignored.")

417
            mm_parser.parse_image(image_url["url"])
418
            has_image = True
419
        elif part_type == "audio_url":
420
            audio_url = _AudioParser(part)["audio_url"]
421
422

            mm_parser.parse_audio(audio_url["url"])
423
424
425
        elif part_type == "refusal":
            text = _RefusalParser(part)["refusal"]
            texts.append(text)
426
427
428
429
        else:
            raise NotImplementedError(f"Unknown part type: {part_type}")

    text_prompt = "\n".join(texts)
430
431
432
433
434
435
436
437
438
439
440
441
442
443
    if keep_multimodal_content:
        text_prompt = "\n".join(texts)
        role_content = [{'type': 'text', 'text': text_prompt}]

        if has_image:
            role_content = [{'type': 'image'}] + role_content
        return [ConversationMessage(role=role,
                                    content=role_content)]  # type: ignore
    else:
        mm_placeholder_counts = mm_parser.mm_placeholder_counts()
        if mm_placeholder_counts:
            text_prompt = _get_full_multimodal_text_prompt(
                mm_placeholder_counts, text_prompt)
        return [ConversationMessage(role=role, content=text_prompt)]
444
445


446
447
448
449
450
# No need to validate using Pydantic again
_AssistantParser = partial(cast, ChatCompletionAssistantMessageParam)
_ToolParser = partial(cast, ChatCompletionToolMessageParam)


451
def _parse_chat_message_content(
452
453
454
    message: ChatCompletionMessageParam,
    mm_tracker: BaseMultiModalItemTracker,
) -> List[ConversationMessage]:
455
456
457
458
    role = message["role"]
    content = message.get("content")

    if content is None:
459
460
461
462
463
        content = []
    elif isinstance(content, str):
        content = [
            ChatCompletionContentPartTextParam(type="text", text=content)
        ]
464

465
    result = _parse_chat_message_content_parts(
466
467
        role,
        content,  # type: ignore
468
        mm_tracker,
469
    )
470

471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
    for result_msg in result:
        if role == 'assistant':
            parsed_msg = _AssistantParser(message)

            if "tool_calls" in parsed_msg:
                result_msg["tool_calls"] = list(parsed_msg["tool_calls"])
        elif role == "tool":
            parsed_msg = _ToolParser(message)
            if "tool_call_id" in parsed_msg:
                result_msg["tool_call_id"] = parsed_msg["tool_call_id"]

        if "name" in message and isinstance(message["name"], str):
            result_msg["name"] = message["name"]

    return result

487

488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
def _postprocess_messages(messages: List[ConversationMessage]) -> None:
    # per the Transformers docs & maintainers, tool call arguments in
    # assistant-role messages with tool_calls need to be dicts not JSON str -
    # this is how tool-use chat templates will expect them moving forwards
    # so, for messages that have tool_calls, parse the string (which we get
    # from openAI format) to dict
    for message in messages:
        if (message["role"] == "assistant" and "tool_calls" in message
                and isinstance(message["tool_calls"], list)):

            for item in message["tool_calls"]:
                item["function"]["arguments"] = json.loads(
                    item["function"]["arguments"])


503
504
505
def parse_chat_messages(
    messages: List[ChatCompletionMessageParam],
    model_config: ModelConfig,
506
    tokenizer: AnyTokenizer,
507
) -> Tuple[List[ConversationMessage], Optional[MultiModalDataDict]]:
508
    conversation: List[ConversationMessage] = []
509
    mm_tracker = MultiModalItemTracker(model_config, tokenizer)
510
511

    for msg in messages:
512
        sub_messages = _parse_chat_message_content(msg, mm_tracker)
513

514
        conversation.extend(sub_messages)
515

516
517
    _postprocess_messages(conversation)

518
    return conversation, mm_tracker.all_mm_data()
519
520


521
522
523
524
525
526
527
528
529
530
531
532
533
def parse_chat_messages_futures(
    messages: List[ChatCompletionMessageParam],
    model_config: ModelConfig,
    tokenizer: AnyTokenizer,
) -> Tuple[List[ConversationMessage], Awaitable[Optional[MultiModalDataDict]]]:
    conversation: List[ConversationMessage] = []
    mm_tracker = AsyncMultiModalItemTracker(model_config, tokenizer)

    for msg in messages:
        sub_messages = _parse_chat_message_content(msg, mm_tracker)

        conversation.extend(sub_messages)

534
535
    _postprocess_messages(conversation)

536
537
538
    return conversation, mm_tracker.all_mm_data()


539
540
def apply_hf_chat_template(
    tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
541
542
543
544
545
    conversation: List[ConversationMessage],
    chat_template: Optional[str],
    *,
    tokenize: bool = False,  # Different from HF's default
    **kwargs: Any,
546
) -> str:
547
548
549
550
551
552
    if chat_template is None and tokenizer.chat_template is None:
        raise ValueError(
            "As of transformers v4.44, default chat template is no longer "
            "allowed, so you must provide a chat template if the tokenizer "
            "does not define one.")

553
554
555
556
557
558
    return tokenizer.apply_chat_template(
        conversation=conversation,  # type: ignore[arg-type]
        chat_template=chat_template,
        tokenize=tokenize,
        **kwargs,
    )
559
560


561
562
563
def apply_mistral_chat_template(
    tokenizer: MistralTokenizer,
    messages: List[ChatCompletionMessageParam],
564
    chat_template: Optional[str] = None,
565
566
    **kwargs: Any,
) -> List[int]:
567
    if chat_template is not None:
568
        print_warning_once(
569
            "'chat_template' cannot be overridden for mistral tokenizer.")
570
    if "add_generation_prompt" in kwargs:
571
        print_warning_once(
572
573
574
            "'add_generation_prompt' is not supported for mistral tokenizer, "
            "so it will be ignored.")
    if "continue_final_message" in kwargs:
575
        print_warning_once(
576
577
            "'continue_final_message' is not supported for mistral tokenizer, "
            "so it will be ignored.")
578

579
580
    return tokenizer.apply_chat_template(
        messages=messages,
581
582
        **kwargs,
    )