Skip to content

Commit

Permalink
Working on tool call test
Browse files Browse the repository at this point in the history
  • Loading branch information
DePasqualeOrg committed Dec 12, 2024
1 parent aa60859 commit 19b93e1
Showing 1 changed file with 83 additions and 78 deletions.
161 changes: 83 additions & 78 deletions Tests/ChatTemplateTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -312,84 +312,6 @@ final class ChatTemplateTests: XCTestCase {
],
target: "<s>\n<|start_header_id|>system<|end_header_id|>\n\nCutting Knowledge Date: December 2023\nToday Date: 26 Jul 2024\n\n<|eot_id|><|start_header_id|>user<|end_header_id|>\n\nWhat's in this image?<|image|><|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
),
Test(
name: "Llama-3.2-11B-Vision-Instruct: function calling",
chatTemplate: llama3_2visionChatTemplate,
data: [
"messages": [
[
"role": "user",
"content": [
[
"type": "text",
"text": "What's the weather like in San Francisco?"
] as [String: Any]
] as [[String: Any]]
] as [String: Any],
[
"role": "assistant",
"tool_calls": [
[
"function": [
"name": "get_weather",
"arguments": "{\"location\":\"San Francisco\",\"unit\":\"celsius\"}"
]
]
],
"content": [
[
"type": "text",
"text": ""
]
]
] as [String: Any],
[
"role": "tool",
"content": [
[
"text": "{\"temperature\": 22, \"condition\": \"sunny\"}",
"type": "text"
]
]
] as [String: Any],
[
"role": "assistant",
"content": [
[
"type": "text",
"text": "The weather in San Francisco is sunny with a temperature of 22°C."
] as [String: Any]
] as [[String: Any]]
] as [String: Any]
] as [[String: Any]] as Any,
"bos_token": "<s>" as Any,
"date_string": "26 Jul 2024" as Any,
"tools_in_user_message": false as Any,
"tools": [
[
"name": "get_weather",
"description": "Get the current weather for a location",
"parameters": [
"type": "object",
"required": ["location", "unit"],
"properties": [
"location": [
"description": "The location to get weather for",
"type": "string"
],
"unit": [
"description": "The unit of temperature",
"type": "string",
"enum": ["celsius", "fahrenheit"]
]
]
]
]
] as [[String: Any]] as Any,
"add_generation_prompt": true as Any
],
target: "<s>\n<|start_header_id|>system<|end_header_id|>\n\nEnvironment: ipython\nCutting Knowledge Date: December 2023\nToday Date: 26 Jul 2024\n\nYou have access to the following functions. To call a function, please respond with JSON for a function call.Respond in the format {\"name\": function name, \"parameters\": dictionary of argument name and its value}.Do not use variables.\n\n{\n \"name\": \"get_weather\",\n \"description\": \"Get the current weather for a location\",\n \"parameters\": {\n \"type\": \"object\",\n \"required\": [\n \"location\",\n \"unit\"\n ],\n \"properties\": {\n \"location\": {\n \"description\": \"The location to get weather for\",\n \"type\": \"string\"\n },\n \"unit\": {\n \"description\": \"The unit of temperature\",\n \"type\": \"string\",\n \"enum\": [\n \"celsius\",\n \"fahrenheit\"\n ]\n }\n }\n }\n}\n\n<|eot_id|><|start_header_id|>user<|end_header_id|>\n\nWhat's the weather like in San Francisco?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n{\"name\": \"get_weather\", \"parameters\": \"{\\\"location\\\":\\\"San Francisco\\\",\\\"unit\\\":\\\"celsius\\\"}\"}<|eot_id|><|start_header_id|>ipython<|end_header_id|>\n\n[{\"text\": \"{\\\"temperature\\\": 22, \\\"condition\\\": \\\"sunny\\\"}\", \"type\": \"text\"}]<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\nThe weather in San Francisco is sunny with a temperature of 22°C.<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
),
// Qwen2-VL-7B-Instruct
// Test(
// chatTemplate:
Expand All @@ -411,4 +333,87 @@ final class ChatTemplateTests: XCTestCase {
XCTAssertEqual(result.debugDescription, test.target.debugDescription)
}
}

func testLlama32ToolCalls() throws {
let tools = [
[
"name": "get_current_weather",
"description": "Get the current weather in a given location",
"parameters": [
"type": "object",
"properties": [
"location": [
"type": "string",
"description": "The city and state, e.g. San Francisco, CA"
],
"unit": [
"type": "string",
"enum": ["celsius", "fahrenheit"]
]
],
"required": ["location"]
]
]
]

let messages: [[String: Any]] = [
[
"role": "user",
"content": "What's the weather like in San Francisco?"
],
[
"role": "assistant",
"tool_calls": [
[
"function": [
"name": "get_current_weather",
"arguments": "{\"location\": \"San Francisco, CA\", \"unit\": \"celsius\"}"
]
]
]
],
[
"role": "tool",
"content": "{\"temperature\": 22, \"unit\": \"celsius\", \"description\": \"Sunny\"}"
],
[
"role": "assistant",
"content": "The weather in San Francisco is sunny with a temperature of 22°C."
]
]

let template = try Template(llama3_2visionChatTemplate)
let result = try template.render([
"messages": messages,
"tools": tools,
"bos_token": "<s>",
"date_string": "26 Jul 2024",
"add_generation_prompt": true
])

print(result) // Debugging for comparison with expected

// TODO: Replace with printed result if it works
let expected = """
<s>
<|start_header_id|>system<|end_header_id|>
Environment: ipython
Cutting Knowledge Date: December 2023
Today Date: 26 Jul 2024
<|eot_id|><|start_header_id|>user<|end_header_id|>
What's the weather like in San Francisco?<|eot_id|><|start_header_id|>assistant<|end_header_id|>
{"name": "get_current_weather", "parameters": {"location": "San Francisco, CA", "unit": "celsius"}}<|eot_id|><|start_header_id|>ipython<|end_header_id|>
{"temperature": 22, "unit": "celsius", "description": "Sunny"}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
The weather in San Francisco is sunny with a temperature of 22°C.<|eot_id|><|start_header_id|>assistant<|end_header_id|>
"""

XCTAssertEqual(result, expected)
}
}

0 comments on commit 19b93e1

Please sign in to comment.