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

4
5
import pytest

6
from vllm.config import ModelConfig
7
8
from vllm.entrypoints.chat_utils import (apply_hf_chat_template,
                                         load_chat_template)
9
from vllm.entrypoints.openai.protocol import ChatCompletionRequest
10
from vllm.transformers_utils.tokenizer import get_tokenizer
11

12
from ...models.registry import HF_EXAMPLE_MODELS
13
from ...utils import VLLM_PATH
14
15

chatml_jinja_path = VLLM_PATH / "examples/template_chatml.jinja"
Simon Mo's avatar
Simon Mo committed
16
17
assert chatml_jinja_path.exists()

18
19
# Define models, templates, and their corresponding expected outputs
MODEL_TEMPLATE_GENERATON_OUTPUT = [
20
    ("facebook/opt-125m", chatml_jinja_path, True, False, """<|im_start|>user
21
22
23
24
25
26
27
Hello<|im_end|>
<|im_start|>assistant
Hi there!<|im_end|>
<|im_start|>user
What is the capital of<|im_end|>
<|im_start|>assistant
"""),
28
    ("facebook/opt-125m", chatml_jinja_path, False, False, """<|im_start|>user
29
30
31
32
Hello<|im_end|>
<|im_start|>assistant
Hi there!<|im_end|>
<|im_start|>user
33
34
35
36
37
38
39
40
41
What is the capital of"""),
    ("facebook/opt-125m", chatml_jinja_path, False, True, """<|im_start|>user
Hello<|im_end|>
<|im_start|>assistant
Hi there!<|im_end|>
<|im_start|>user
What is the capital of<|im_end|>
<|im_start|>assistant
The capital of"""),
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
]

TEST_MESSAGES = [
    {
        'role': 'user',
        'content': 'Hello'
    },
    {
        'role': 'assistant',
        'content': 'Hi there!'
    },
    {
        'role': 'user',
        'content': 'What is the capital of'
    },
]
58
59
60
61
ASSISTANT_MESSAGE_TO_CONTINUE = {
    'role': 'assistant',
    'content': 'The capital of'
}
62
63


64
def test_load_chat_template():
65
    # Testing chatml template
66
    template_content = load_chat_template(chat_template=chatml_jinja_path)
67
68
69
70
71

    # Test assertions
    assert template_content is not None
    # Hard coded value for template_chatml.jinja
    assert template_content == """{% for message in messages %}{{'<|im_start|>' + message['role'] + '\\n' + message['content']}}{% if (loop.last and add_generation_prompt) or not loop.last %}{{ '<|im_end|>' + '\\n'}}{% endif %}{% endfor %}
72
{% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}{{ '<|im_start|>assistant\\n' }}{% endif %}"""  # noqa: E501
73
74


75
def test_no_load_chat_template_filelike():
76
77
    # Testing chatml template
    template = "../../examples/does_not_exist"
78
79

    with pytest.raises(ValueError, match="looks like a file path"):
80
        load_chat_template(chat_template=template)
81
82


83
def test_no_load_chat_template_literallike():
84
85
86
    # Testing chatml template
    template = "{{ messages }}"

87
    template_content = load_chat_template(chat_template=template)
88

89
    assert template_content == template
90
91
92


@pytest.mark.parametrize(
93
    "model,template,add_generation_prompt,continue_final_message,expected_output",
94
    MODEL_TEMPLATE_GENERATON_OUTPUT)
95
def test_get_gen_prompt(model, template, add_generation_prompt,
96
                        continue_final_message, expected_output):
97
98
99
100
101
102
103
104
105
106
107
    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
    model_info.check_available_online(on_fail="skip")

    model_config = ModelConfig(
        model,
        tokenizer=model_info.tokenizer or model,
        tokenizer_mode=model_info.tokenizer_mode,
        trust_remote_code=model_info.trust_remote_code,
        hf_overrides=model_info.hf_overrides,
    )

108
    # Initialize the tokenizer
109
110
111
112
    tokenizer = get_tokenizer(
        tokenizer_name=model_config.tokenizer,
        trust_remote_code=model_config.trust_remote_code,
    )
113
    template_content = load_chat_template(chat_template=template)
114
115
116
117

    # Create a mock request object using keyword arguments
    mock_request = ChatCompletionRequest(
        model=model,
118
119
120
121
122
        messages=TEST_MESSAGES + [ASSISTANT_MESSAGE_TO_CONTINUE]
        if continue_final_message else TEST_MESSAGES,
        add_generation_prompt=add_generation_prompt,
        continue_final_message=continue_final_message,
    )
123
124

    # Call the function and get the result
125
    result = apply_hf_chat_template(
126
        tokenizer=tokenizer,
127
        conversation=mock_request.messages,
128
        chat_template=mock_request.chat_template or template_content,
129
        model_config=model_config,
130
        tools=None,
131
        add_generation_prompt=mock_request.add_generation_prompt,
132
        continue_final_message=mock_request.continue_final_message,
133
    )
134
135

    # Test assertion
136
137
138
    assert result == expected_output, (
        f"The generated prompt does not match the expected output for "
        f"model {model} and template {template}")