[Frontend] Move chat utils (#6602)

Co-authored-by: Roger Wang <ywang@roblox.com>
This commit is contained in:
Cyrus Leung 2024-07-21 08:38:17 +08:00 committed by GitHub
parent 082ecd80d5
commit d7f4178dd9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 53 additions and 48 deletions

View File

@ -3,7 +3,7 @@ import pathlib
import pytest import pytest
from vllm.entrypoints.openai.chat_utils import load_chat_template from vllm.entrypoints.chat_utils import load_chat_template
from vllm.entrypoints.openai.protocol import ChatCompletionRequest from vllm.entrypoints.openai.protocol import ChatCompletionRequest
from vllm.transformers_utils.tokenizer import get_tokenizer from vllm.transformers_utils.tokenizer import get_tokenizer

View File

@ -1,15 +1,23 @@
import codecs import codecs
from dataclasses import dataclass, field from dataclasses import dataclass, field
from functools import lru_cache from functools import lru_cache
from typing import Awaitable, Iterable, List, Optional, TypedDict, cast, final from typing import Awaitable, Iterable, List, Optional, Union, cast, final
from openai.types.chat import (ChatCompletionContentPartImageParam, # yapf conflicts with isort for this block
ChatCompletionContentPartTextParam) # yapf: disable
from openai.types.chat import ChatCompletionContentPartImageParam
from openai.types.chat import (
ChatCompletionContentPartParam as OpenAIChatCompletionContentPartParam)
from openai.types.chat import ChatCompletionContentPartTextParam
from openai.types.chat import (
ChatCompletionMessageParam as OpenAIChatCompletionMessageParam)
# yapf: enable
# pydantic needs the TypedDict from typing_extensions
from pydantic import ConfigDict
from transformers import PreTrainedTokenizer from transformers import PreTrainedTokenizer
from typing_extensions import Required, TypedDict
from vllm.config import ModelConfig from vllm.config import ModelConfig
from vllm.entrypoints.openai.protocol import (ChatCompletionContentPartParam,
ChatCompletionMessageParam)
from vllm.logger import init_logger from vllm.logger import init_logger
from vllm.multimodal import MultiModalDataDict from vllm.multimodal import MultiModalDataDict
from vllm.multimodal.utils import async_get_and_parse_image from vllm.multimodal.utils import async_get_and_parse_image
@ -17,6 +25,37 @@ from vllm.multimodal.utils import async_get_and_parse_image
logger = init_logger(__name__) logger = init_logger(__name__)
class CustomChatCompletionContentPartParam(TypedDict, total=False):
__pydantic_config__ = ConfigDict(extra="allow") # type: ignore
type: Required[str]
"""The type of the content part."""
ChatCompletionContentPartParam = Union[OpenAIChatCompletionContentPartParam,
CustomChatCompletionContentPartParam]
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.
"""
ChatCompletionMessageParam = Union[OpenAIChatCompletionMessageParam,
CustomChatCompletionMessageParam]
@final # So that it should be compatible with Dict[str, str] @final # So that it should be compatible with Dict[str, str]
class ConversationMessage(TypedDict): class ConversationMessage(TypedDict):
role: str role: str

View File

@ -3,50 +3,16 @@
import time import time
from typing import Any, Dict, List, Literal, Optional, Union from typing import Any, Dict, List, Literal, Optional, Union
import openai.types.chat
import torch import torch
from pydantic import BaseModel, ConfigDict, Field, model_validator from pydantic import BaseModel, ConfigDict, Field, model_validator
# pydantic needs the TypedDict from typing_extensions from typing_extensions import Annotated
from typing_extensions import Annotated, Required, TypedDict
from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
from vllm.pooling_params import PoolingParams from vllm.pooling_params import PoolingParams
from vllm.sampling_params import SamplingParams from vllm.sampling_params import SamplingParams
from vllm.utils import random_uuid from vllm.utils import random_uuid
class CustomChatCompletionContentPartParam(TypedDict, total=False):
__pydantic_config__ = ConfigDict(extra="allow") # type: ignore
type: Required[str]
"""The type of the content part."""
ChatCompletionContentPartParam = Union[
openai.types.chat.ChatCompletionContentPartParam,
CustomChatCompletionContentPartParam]
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.
"""
ChatCompletionMessageParam = Union[
openai.types.chat.ChatCompletionMessageParam,
CustomChatCompletionMessageParam]
class OpenAIBaseModel(BaseModel): class OpenAIBaseModel(BaseModel):
# OpenAI API does not allow extra fields # OpenAI API does not allow extra fields
model_config = ConfigDict(extra="forbid") model_config = ConfigDict(extra="forbid")

View File

@ -9,9 +9,9 @@ from transformers import PreTrainedTokenizer
from vllm.config import ModelConfig from vllm.config import ModelConfig
from vllm.engine.async_llm_engine import AsyncLLMEngine from vllm.engine.async_llm_engine import AsyncLLMEngine
from vllm.entrypoints.openai.chat_utils import (ConversationMessage, from vllm.entrypoints.chat_utils import (ConversationMessage,
load_chat_template, load_chat_template,
parse_chat_message_content) parse_chat_message_content)
from vllm.entrypoints.openai.protocol import ( from vllm.entrypoints.openai.protocol import (
ChatCompletionLogProb, ChatCompletionLogProbs, ChatCompletionLogProb, ChatCompletionLogProbs,
ChatCompletionLogProbsContent, ChatCompletionNamedToolChoiceParam, ChatCompletionLogProbsContent, ChatCompletionNamedToolChoiceParam,

View File

@ -2,9 +2,9 @@ from typing import List, Optional
from vllm.config import ModelConfig from vllm.config import ModelConfig
from vllm.engine.async_llm_engine import AsyncLLMEngine from vllm.engine.async_llm_engine import AsyncLLMEngine
from vllm.entrypoints.openai.chat_utils import (ConversationMessage, from vllm.entrypoints.chat_utils import (ConversationMessage,
load_chat_template, load_chat_template,
parse_chat_message_content) parse_chat_message_content)
from vllm.entrypoints.openai.protocol import (DetokenizeRequest, from vllm.entrypoints.openai.protocol import (DetokenizeRequest,
DetokenizeResponse, DetokenizeResponse,
TokenizeRequest, TokenizeRequest,