Merge pull request #13 from DePasqualeOrg/fix-objectvalue

Bugfix: Persist member values in `ObjectValue`
johnmai-dev · Jan 27, 2025 · bbddb92 · bbddb92
2 parents 9c0bbbc + 2a6a0f9
commit bbddb92
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 16 deletions.
diff --git a/Sources/Runtime.swift b/Sources/Runtime.swift
@@ -90,7 +90,7 @@ struct TupleValue: RuntimeValue {
     }
 }
 
-struct ObjectValue: RuntimeValue, Sequence {
+class ObjectValue: RuntimeValue, Sequence {
     var storage: OrderedDictionary<String, any RuntimeValue>
     var builtins: [String: any RuntimeValue]
 
@@ -139,7 +139,7 @@ struct ObjectValue: RuntimeValue, Sequence {
         ]
     }
 
-    mutating func setValue(key: String, value: any RuntimeValue) {
+    func setValue(key: String, value: any RuntimeValue) {
         storage[key] = value
     }
 

diff --git a/Tests/Templates/ChatTemplateTests.swift b/Tests/Templates/ChatTemplateTests.swift
@@ -25,13 +25,12 @@ final class ChatTemplateTests: XCTestCase {
         ],
     ]
 
-    lazy var messagesWithSystemPrompt: [[String: String]] =
-        [
-            [
-                "role": "system",
-                "content": "You are a friendly chatbot who always responds in the style of a pirate",
-            ]
-        ] + messages
+    let systemPromptMessage: [String: String] = [
+        "role": "system",
+        "content": "You are a friendly chatbot who always responds in the style of a pirate",
+    ]
+
+    lazy var messagesWithSystemPrompt: [[String: String]] = [systemPromptMessage] + messages
 
     func testGenericChatTemplate() throws {
         let chatTemplate =
@@ -600,15 +599,16 @@ final class ChatTemplateTests: XCTestCase {
         XCTAssertEqual(result, target)
     }
 
-    func testDeepSeekQwen() throws {
+    let deepSeekR1chatTemplate = """
+        {% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='') %}{%- for message in messages %}{%- if message['role'] == 'system' %}{% set ns.system_prompt = message['content'] %}{%- endif %}{%- endfor %}{{bos_token}}{{ns.system_prompt}}{%- for message in messages %}{%- if message['role'] == 'user' %}{%- set ns.is_tool = false -%}{{'<｜User｜>' + message['content']}}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is none %}{%- set ns.is_tool = false -%}{%- for tool in message['tool_calls']%}{%- if not ns.is_first %}{{'<｜Assistant｜><｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\\n' + '```json' + '\\n' + tool['function']['arguments'] + '\\n' + '```' + '<｜tool▁call▁end｜>'}}{%- set ns.is_first = true -%}{%- else %}{{'\\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\\n' + '```json' + '\\n' + tool['function']['arguments'] + '\\n' + '```' + '<｜tool▁call▁end｜>'}}{{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}}{%- endif %}{%- endfor %}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is not none %}{%- if ns.is_tool %}{{'<｜tool▁outputs▁end｜>' + message['content'] + '<｜end▁of▁sentence｜>'}}{%- set ns.is_tool = false -%}{%- else %}{% set content = message['content'] %}{% if '</think>' in content %}{% set content = content.split('</think>')[-1] %}{% endif %}{{'<｜Assistant｜>' + content + '<｜end▁of▁sentence｜>'}}{%- endif %}{%- endif %}{%- if message['role'] == 'tool' %}{%- set ns.is_tool = true -%}{%- if ns.is_output_first %}{{'<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- set ns.is_output_first = false %}{%- else %}{{'\\n<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- endif %}{%- endif %}{%- endfor -%}{% if ns.is_tool %}{{'<｜tool▁outputs▁end｜>'}}{% endif %}{% if add_generation_prompt and not ns.is_tool %}{{'<｜Assistant｜>'}}{% endif %}
+        """
+
+    func testDeepSeekR1() throws {
         let userMessage = [
             "role": "user",
             "content": "What is the weather in Paris today?",
         ]
-        let chatTemplate = """
-            {% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='') %}{%- for message in messages %}{%- if message['role'] == 'system' %}{% set ns.system_prompt = message['content'] %}{%- endif %}{%- endfor %}{{bos_token}}{{ns.system_prompt}}{%- for message in messages %}{%- if message['role'] == 'user' %}{%- set ns.is_tool = false -%}{{'<｜User｜>' + message['content']}}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is none %}{%- set ns.is_tool = false -%}{%- for tool in message['tool_calls']%}{%- if not ns.is_first %}{{'<｜Assistant｜><｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\\n' + '```json' + '\\n' + tool['function']['arguments'] + '\\n' + '```' + '<｜tool▁call▁end｜>'}}{%- set ns.is_first = true -%}{%- else %}{{'\\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\\n' + '```json' + '\\n' + tool['function']['arguments'] + '\\n' + '```' + '<｜tool▁call▁end｜>'}}{{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}}{%- endif %}{%- endfor %}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is not none %}{%- if ns.is_tool %}{{'<｜tool▁outputs▁end｜>' + message['content'] + '<｜end▁of▁sentence｜>'}}{%- set ns.is_tool = false -%}{%- else %}{% set content = message['content'] %}{% if '</think>' in content %}{% set content = content.split('</think>')[-1] %}{% endif %}{{'<｜Assistant｜>' + content + '<｜end▁of▁sentence｜>'}}{%- endif %}{%- endif %}{%- if message['role'] == 'tool' %}{%- set ns.is_tool = true -%}{%- if ns.is_output_first %}{{'<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- set ns.is_output_first = false %}{%- else %}{{'\\n<｜tool▁output▁begin｜>' + message['content'] + '<｜tool▁output▁end｜>'}}{%- endif %}{%- endif %}{%- endfor -%}{% if ns.is_tool %}{{'<｜tool▁outputs▁end｜>'}}{% endif %}{% if add_generation_prompt and not ns.is_tool %}{{'<｜Assistant｜>'}}{% endif %}
-            """
-        let template = try Template(chatTemplate)
+        let template = try Template(deepSeekR1chatTemplate)
         let result = try template.render([
             "messages": [userMessage],
             "bos_token": "<|begin_of_text|>",
@@ -619,4 +619,21 @@ final class ChatTemplateTests: XCTestCase {
             """
         XCTAssertEqual(result, target)
     }
+
+    func testDeepSeekR1WitihSystemPrompt() throws {
+        let userMessage = [
+            "role": "user",
+            "content": "What is the weather in Paris today?",
+        ]
+        let template = try Template(deepSeekR1chatTemplate)
+        let result = try template.render([
+            "messages": [systemPromptMessage, userMessage],
+            "bos_token": "<|begin_of_text|>",
+            "add_generation_prompt": true,
+        ])
+        let target = """
+            <|begin_of_text|>You are a friendly chatbot who always responds in the style of a pirate<｜User｜>What is the weather in Paris today?<｜Assistant｜>
+            """
+        XCTAssertEqual(result, target)
+    }
 }
diff --git a/Tests/Templates/VisionTests.swift b/Tests/Templates/VisionTests.swift
@@ -80,8 +80,14 @@ final class VisionTests: XCTestCase {
             "bos_token": "<s>" as Any,
             "add_generation_prompt": true as Any,
         ])
-        let 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"
+        let target = """
+            <s>
+            <|start_header_id|>user<|end_header_id|>
+
+            What's in this image?<|image|><|eot_id|><|start_header_id|>assistant<|end_header_id|>
+
+
+            """
         XCTAssertEqual(result, target)
     }