serving_tokenization.py 5.11 KB
Newer Older
1
from typing import List, Optional, Union
2
3

from vllm.config import ModelConfig
4
from vllm.engine.protocol import AsyncEngineClient
5
6
from vllm.entrypoints.chat_utils import (apply_chat_template,
                                         load_chat_template,
7
                                         parse_chat_messages_futures)
8
from vllm.entrypoints.logger import RequestLogger
9
10
# yapf conflicts with isort for this block
# yapf: disable
11
12
from vllm.entrypoints.openai.protocol import (DetokenizeRequest,
                                              DetokenizeResponse,
13
14
                                              ErrorResponse,
                                              TokenizeChatRequest,
15
16
                                              TokenizeRequest,
                                              TokenizeResponse)
17
# yapf: enable
18
19
from vllm.entrypoints.openai.serving_engine import (LoRAModulePath,
                                                    OpenAIServing)
20
from vllm.logger import init_logger
21
from vllm.utils import random_uuid
22

23
24
logger = init_logger(__name__)

25
26
27

class OpenAIServingTokenization(OpenAIServing):

28
29
    def __init__(
        self,
30
        async_engine_client: AsyncEngineClient,
31
32
33
34
35
36
37
        model_config: ModelConfig,
        served_model_names: List[str],
        *,
        lora_modules: Optional[List[LoRAModulePath]],
        request_logger: Optional[RequestLogger],
        chat_template: Optional[str],
    ):
38
        super().__init__(async_engine_client=async_engine_client,
39
40
                         model_config=model_config,
                         served_model_names=served_model_names,
41
42
43
                         lora_modules=lora_modules,
                         prompt_adapters=None,
                         request_logger=request_logger)
44

45
        # If this is None we use the tokenizer's default chat template
46
47
48
49
50
        # the list of commonly-used chat template names for HF named templates
        hf_chat_templates: List[str] = ['default', 'tool_use']
        self.chat_template = chat_template \
            if chat_template in hf_chat_templates \
            else load_chat_template(chat_template)
51

52
53
54
55
    async def create_tokenize(
        self,
        request: TokenizeRequest,
    ) -> Union[TokenizeResponse, ErrorResponse]:
56
57
58
59
        error_check_ret = await self._check_model(request)
        if error_check_ret is not None:
            return error_check_ret

60
        request_id = f"tokn-{random_uuid()}"
61

62
63
64
65
        (
            lora_request,
            prompt_adapter_request,
        ) = self._maybe_get_adapters(request)
66

67
        tokenizer = await self.async_engine_client.get_tokenizer(lora_request)
68
69
70
71

        if isinstance(request, TokenizeChatRequest):
            model_config = self.model_config

72
            conversation, mm_data_future = parse_chat_messages_futures(
73
                request.messages, model_config, tokenizer)
74

75
76
            mm_data = await mm_data_future
            if mm_data:
77
78
                logger.warning(
                    "Multi-modal inputs are ignored during tokenization")
79

80
81
            prompt = apply_chat_template(
                tokenizer,
82
                conversation=conversation,
83
84
85
                chat_template=self.chat_template,
                add_generation_prompt=request.add_generation_prompt,
            )
86
87
88
89
90
91
92
93
        else:
            prompt = request.prompt

        self._log_inputs(request_id,
                         prompt,
                         params=None,
                         lora_request=lora_request,
                         prompt_adapter_request=prompt_adapter_request)
94

95
96
97
        # Silently ignore prompt adapter since it does not affect tokenization

        prompt_input = self._tokenize_prompt_input(
98
            request,
99
            tokenizer,
100
101
102
103
            prompt,
            add_special_tokens=request.add_special_tokens,
        )
        input_ids = prompt_input["prompt_token_ids"]
104
105
106
107
108
109

        return TokenizeResponse(tokens=input_ids,
                                count=len(input_ids),
                                max_model_len=self.max_model_len)

    async def create_detokenize(
110
111
112
        self,
        request: DetokenizeRequest,
    ) -> Union[DetokenizeResponse, ErrorResponse]:
113
114
115
116
        error_check_ret = await self._check_model(request)
        if error_check_ret is not None:
            return error_check_ret

117
118
119
120
121
122
123
        request_id = f"tokn-{random_uuid()}"

        (
            lora_request,
            prompt_adapter_request,
        ) = self._maybe_get_adapters(request)

124
        tokenizer = await self.async_engine_client.get_tokenizer(lora_request)
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

        self._log_inputs(request_id,
                         request.tokens,
                         params=None,
                         lora_request=lora_request,
                         prompt_adapter_request=prompt_adapter_request)

        if prompt_adapter_request is not None:
            raise NotImplementedError("Prompt adapter is not supported "
                                      "for tokenization")

        prompt_input = self._tokenize_prompt_input(
            request,
            tokenizer,
            request.tokens,
        )
        input_text = prompt_input["prompt"]
142
143

        return DetokenizeResponse(prompt=input_text)