test_formatter.py 9.81 KB
Newer Older
chenych's avatar
chenych committed
1
# Copyright 2025 the LlamaFactory team.
chenych's avatar
chenych committed
2
3
4
5
6
7
8
9
10
11
12
13
14
15
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
luopl's avatar
luopl committed
16
from datetime import datetime
chenych's avatar
chenych committed
17
18
19
20

from llamafactory.data.formatter import EmptyFormatter, FunctionFormatter, StringFormatter, ToolFormatter


luopl's avatar
luopl committed
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
FUNCTION = {"name": "tool_name", "arguments": {"foo": "bar", "size": 10}}

TOOLS = [
    {
        "name": "test_tool",
        "description": "tool_desc",
        "parameters": {
            "type": "object",
            "properties": {
                "foo": {"type": "string", "description": "foo_desc"},
                "bar": {"type": "number", "description": "bar_desc"},
            },
            "required": ["foo"],
        },
    }
]


chenych's avatar
chenych committed
39
40
41
42
43
44
45
46
47
48
49
def test_empty_formatter():
    formatter = EmptyFormatter(slots=["\n"])
    assert formatter.apply() == ["\n"]


def test_string_formatter():
    formatter = StringFormatter(slots=["<s>", "Human: {{content}}\nAssistant:"])
    assert formatter.apply(content="Hi") == ["<s>", "Human: Hi\nAssistant:"]


def test_function_formatter():
luopl's avatar
luopl committed
50
51
    formatter = FunctionFormatter(slots=["{{content}}", "</s>"], tool_format="default")
    tool_calls = json.dumps(FUNCTION)
chenych's avatar
chenych committed
52
    assert formatter.apply(content=tool_calls) == [
luopl's avatar
luopl committed
53
54
        """Action: tool_name\nAction Input: {"foo": "bar", "size": 10}\n""",
        "</s>",
chenych's avatar
chenych committed
55
56
57
58
    ]


def test_multi_function_formatter():
luopl's avatar
luopl committed
59
60
    formatter = FunctionFormatter(slots=["{{content}}", "</s>"], tool_format="default")
    tool_calls = json.dumps([FUNCTION] * 2)
chenych's avatar
chenych committed
61
    assert formatter.apply(content=tool_calls) == [
luopl's avatar
luopl committed
62
63
64
        """Action: tool_name\nAction Input: {"foo": "bar", "size": 10}\n"""
        """Action: tool_name\nAction Input: {"foo": "bar", "size": 10}\n""",
        "</s>",
chenych's avatar
chenych committed
65
66
67
68
69
    ]


def test_default_tool_formatter():
    formatter = ToolFormatter(tool_format="default")
luopl's avatar
luopl committed
70
    assert formatter.apply(content=json.dumps(TOOLS)) == [
chenych's avatar
chenych committed
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
        "You have access to the following tools:\n"
        "> Tool Name: test_tool\n"
        "Tool Description: tool_desc\n"
        "Tool Args:\n"
        "  - foo (string, required): foo_desc\n"
        "  - bar (number): bar_desc\n\n"
        "Use the following format if using a tool:\n"
        "```\n"
        "Action: tool name (one of [test_tool])\n"
        "Action Input: the input to the tool, in a JSON format representing the kwargs "
        """(e.g. ```{"input": "hello world", "num_beams": 5}```)\n"""
        "```\n"
    ]


def test_default_tool_extractor():
    formatter = ToolFormatter(tool_format="default")
    result = """Action: test_tool\nAction Input: {"foo": "bar", "size": 10}\n"""
    assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")]


def test_default_multi_tool_extractor():
    formatter = ToolFormatter(tool_format="default")
    result = (
        """Action: test_tool\nAction Input: {"foo": "bar", "size": 10}\n"""
        """Action: another_tool\nAction Input: {"foo": "job", "size": 2}\n"""
    )
    assert formatter.extract(result) == [
        ("test_tool", """{"foo": "bar", "size": 10}"""),
        ("another_tool", """{"foo": "job", "size": 2}"""),
    ]


luopl's avatar
luopl committed
104
105
106
107
108
109
def test_glm4_function_formatter():
    formatter = FunctionFormatter(slots=["{{content}}"], tool_format="glm4")
    tool_calls = json.dumps(FUNCTION)
    assert formatter.apply(content=tool_calls) == ["""tool_name\n{"foo": "bar", "size": 10}"""]


chenych's avatar
chenych committed
110
111
def test_glm4_tool_formatter():
    formatter = ToolFormatter(tool_format="glm4")
luopl's avatar
luopl committed
112
    assert formatter.apply(content=json.dumps(TOOLS)) == [
chenych's avatar
chenych committed
113
114
        "你是一个名为 ChatGLM 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的,"
        "你的任务是针对用户的问题和要求提供适当的答复和支持。# 可用工具\n\n"
chenych's avatar
chenych committed
115
116
        f"## test_tool\n\n{json.dumps(TOOLS[0], indent=4, ensure_ascii=False)}\n"
        "在调用上述函数时,请使用 Json 格式表示调用的参数。"
chenych's avatar
chenych committed
117
118
119
120
121
122
123
    ]


def test_glm4_tool_extractor():
    formatter = ToolFormatter(tool_format="glm4")
    result = """test_tool\n{"foo": "bar", "size": 10}\n"""
    assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")]
luopl's avatar
luopl committed
124
125
126


def test_llama3_function_formatter():
chenych's avatar
chenych committed
127
    formatter = FunctionFormatter(slots=["{{content}}<|eot_id|>"], tool_format="llama3")
luopl's avatar
luopl committed
128
129
    tool_calls = json.dumps({"name": "tool_name", "arguments": {"foo": "bar", "size": 10}})
    assert formatter.apply(content=tool_calls) == [
chenych's avatar
chenych committed
130
        """{"name": "tool_name", "parameters": {"foo": "bar", "size": 10}}<|eot_id|>"""
luopl's avatar
luopl committed
131
132
133
134
135
136
137
138
139
    ]


def test_llama3_tool_formatter():
    formatter = ToolFormatter(tool_format="llama3")
    date = datetime.now().strftime("%d %b %Y")
    wrapped_tool = {"type": "function", "function": TOOLS[0]}
    assert formatter.apply(content=json.dumps(TOOLS)) == [
        f"Cutting Knowledge Date: December 2023\nToday Date: {date}\n\n"
chenych's avatar
chenych committed
140
141
        "You have access to the following functions. "
        "To call a function, please respond with JSON for a function call. "
luopl's avatar
luopl committed
142
143
144
145
146
147
148
149
150
151
152
153
        """Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}. """
        f"Do not use variables.\n\n{json.dumps(wrapped_tool, indent=4, ensure_ascii=False)}\n\n"
    ]


def test_llama3_tool_extractor():
    formatter = ToolFormatter(tool_format="llama3")
    result = """{"name": "test_tool", "parameters": {"foo": "bar", "size": 10}}\n"""
    assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")]


def test_mistral_function_formatter():
chenych's avatar
chenych committed
154
    formatter = FunctionFormatter(slots=["[TOOL_CALLS] {{content}}", "</s>"], tool_format="mistral")
luopl's avatar
luopl committed
155
156
    tool_calls = json.dumps(FUNCTION)
    assert formatter.apply(content=tool_calls) == [
chenych's avatar
chenych committed
157
        "[TOOL_CALLS] " """[{"name": "tool_name", "arguments": {"foo": "bar", "size": 10}}]""",
luopl's avatar
luopl committed
158
159
160
161
162
        "</s>",
    ]


def test_mistral_multi_function_formatter():
chenych's avatar
chenych committed
163
    formatter = FunctionFormatter(slots=["[TOOL_CALLS] {{content}}", "</s>"], tool_format="mistral")
luopl's avatar
luopl committed
164
165
    tool_calls = json.dumps([FUNCTION] * 2)
    assert formatter.apply(content=tool_calls) == [
chenych's avatar
chenych committed
166
        "[TOOL_CALLS] "
luopl's avatar
luopl committed
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
        """[{"name": "tool_name", "arguments": {"foo": "bar", "size": 10}}, """
        """{"name": "tool_name", "arguments": {"foo": "bar", "size": 10}}]""",
        "</s>",
    ]


def test_mistral_tool_formatter():
    formatter = ToolFormatter(tool_format="mistral")
    wrapped_tool = {"type": "function", "function": TOOLS[0]}
    assert formatter.apply(content=json.dumps(TOOLS)) == [
        "[AVAILABLE_TOOLS] " + json.dumps([wrapped_tool], ensure_ascii=False) + "[/AVAILABLE_TOOLS]"
    ]


def test_mistral_tool_extractor():
    formatter = ToolFormatter(tool_format="mistral")
    result = """{"name": "test_tool", "arguments": {"foo": "bar", "size": 10}}"""
    assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")]


def test_mistral_multi_tool_extractor():
    formatter = ToolFormatter(tool_format="mistral")
    result = (
        """[{"name": "test_tool", "arguments": {"foo": "bar", "size": 10}}, """
        """{"name": "another_tool", "arguments": {"foo": "job", "size": 2}}]"""
    )
    assert formatter.extract(result) == [
        ("test_tool", """{"foo": "bar", "size": 10}"""),
        ("another_tool", """{"foo": "job", "size": 2}"""),
    ]


def test_qwen_function_formatter():
chenych's avatar
chenych committed
200
    formatter = FunctionFormatter(slots=["{{content}}<|im_end|>\n"], tool_format="qwen")
luopl's avatar
luopl committed
201
202
    tool_calls = json.dumps(FUNCTION)
    assert formatter.apply(content=tool_calls) == [
chenych's avatar
chenych committed
203
        """<tool_call>\n{"name": "tool_name", "arguments": {"foo": "bar", "size": 10}}\n</tool_call><|im_end|>\n"""
luopl's avatar
luopl committed
204
205
206
207
    ]


def test_qwen_multi_function_formatter():
chenych's avatar
chenych committed
208
    formatter = FunctionFormatter(slots=["{{content}}<|im_end|>\n"], tool_format="qwen")
luopl's avatar
luopl committed
209
210
211
    tool_calls = json.dumps([FUNCTION] * 2)
    assert formatter.apply(content=tool_calls) == [
        """<tool_call>\n{"name": "tool_name", "arguments": {"foo": "bar", "size": 10}}\n</tool_call>\n"""
chenych's avatar
chenych committed
212
213
        """<tool_call>\n{"name": "tool_name", "arguments": {"foo": "bar", "size": 10}}\n</tool_call>"""
        "<|im_end|>\n"
luopl's avatar
luopl committed
214
215
216
217
218
219
220
221
222
223
224
225
    ]


def test_qwen_tool_formatter():
    formatter = ToolFormatter(tool_format="qwen")
    wrapped_tool = {"type": "function", "function": TOOLS[0]}
    assert formatter.apply(content=json.dumps(TOOLS)) == [
        "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\n"
        "You are provided with function signatures within <tools></tools> XML tags:\n<tools>"
        f"\n{json.dumps(wrapped_tool, ensure_ascii=False)}"
        "\n</tools>\n\nFor each function call, return a json object with function name and arguments within "
        """<tool_call></tool_call> XML tags:\n<tool_call>\n{"name": <function-name>, """
chenych's avatar
chenych committed
226
        """"arguments": <args-json-object>}\n</tool_call>"""
luopl's avatar
luopl committed
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
    ]


def test_qwen_tool_extractor():
    formatter = ToolFormatter(tool_format="qwen")
    result = """<tool_call>\n{"name": "test_tool", "arguments": {"foo": "bar", "size": 10}}\n</tool_call>"""
    assert formatter.extract(result) == [("test_tool", """{"foo": "bar", "size": 10}""")]


def test_qwen_multi_tool_extractor():
    formatter = ToolFormatter(tool_format="qwen")
    result = (
        """<tool_call>\n{"name": "test_tool", "arguments": {"foo": "bar", "size": 10}}\n</tool_call>\n"""
        """<tool_call>\n{"name": "another_tool", "arguments": {"foo": "job", "size": 2}}\n</tool_call>"""
    )
    assert formatter.extract(result) == [
        ("test_tool", """{"foo": "bar", "size": 10}"""),
        ("another_tool", """{"foo": "job", "size": 2}"""),
    ]