Merge pull request docker#48 from google/test-polyfills

ochafik · web-flow · commit 78bf4a51c6cc · 2025-02-04T11:10:25.000Z
Add test-polyfills
diff --git a/include/minja/chat-template.hpp b/include/minja/chat-template.hpp
@@ -249,11 +249,10 @@ class chat_template {
                     inputs.add_generation_prompt = false;
                     full = apply(inputs);
                 }
-
-                if (full.find(prefix) != 0) {
-                    if (prefix.rfind(eos_token_) == prefix.size() - eos_token_.size()) {
-                        prefix = prefix.substr(0, prefix.size() - eos_token_.size());
-                    }
+                auto eos_pos_last = full.rfind(eos_token_);
+                if (eos_pos_last == prefix.size() - eos_token_.size() ||
+                      (full[full.size() - 1] == '\n' && (eos_pos_last == full.size() - eos_token_.size() - 1))) {
+                    full = full.substr(0, eos_pos_last);
                 }
                 if (full.find(prefix) != 0) {
                     fprintf(stderr, "Failed to infer a tool call example (possible template bug)\n");
@@ -363,7 +362,7 @@ class chat_template {
             if (polyfill_tools) {
                 adjusted_messages = add_system(inputs.messages,
                     "You can call any of the following tools to satisfy the user's requests: " + minja::Value(inputs.tools).dump(2, /* to_json= */ true) +
-                    (!polyfill_tool_call_example || tool_call_example_.empty() ? "" : "\n\nExample tool call syntax:\n\n" + tool_call_example_));
+                    (!polyfill_tool_call_example || tool_call_example_.empty() ? "" : "\n\nExample tool call syntax:\n\n" + tool_call_example_ + "\n\n"));
             } else {
                 adjusted_messages = inputs.messages;
             }
diff --git a/scripts/fetch_templates_and_goldens.py b/scripts/fetch_templates_and_goldens.py
@@ -278,7 +278,7 @@ def apply(self, context):
             if has_tools and not caps.supports_tools:
                 add_system(context['messages'],
                     f"You can call any of the following tools to satisfy the user's requests: {json.dumps(context['tools'], indent=2)}" +
-                    ("\n\nExample tool call syntax:\n\n" + self.tool_call_example if self.tool_call_example is not None else ""))
+                    ("\n\nExample tool call syntax:\n\n" + self.tool_call_example + "\n\n" if self.tool_call_example is not None else ""))
 
             for message in context['messages']:
                 if 'tool_calls' in message:
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
@@ -17,8 +17,25 @@ target_link_libraries(test-syntax PRIVATE
     gtest_main
     gmock
 )
+
+add_executable(test-polyfills test-polyfills.cpp)
+target_compile_features(test-polyfills PUBLIC cxx_std_17)
+if (CMAKE_SYSTEM_NAME STREQUAL "Windows" AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
+    target_compile_definitions(test-polyfills PUBLIC _CRT_SECURE_NO_WARNINGS)
+    target_compile_options(gtest PRIVATE -Wno-language-extension-token)
+endif()
+target_link_libraries(test-polyfills PRIVATE
+    nlohmann_json::nlohmann_json
+    gtest_main
+    gmock
+)
+if (NOT CMAKE_CROSSCOMPILING)
+    gtest_discover_tests(test-syntax)
+endif()
+
 if (NOT CMAKE_CROSSCOMPILING)
     gtest_discover_tests(test-syntax)
+    gtest_discover_tests(test-polyfills)
 endif()
 
 add_executable(test-capabilities test-capabilities.cpp)
@@ -54,7 +71,7 @@ set(MODEL_IDS
     # minja implementation on the same template and context, and compare the output with the golden.
     #
     # For Gated models, you'll need to run `huggingface-cli login` (and be granted access) to download their template.
-    
+
     abacusai/Fewshot-Metamath-OrcaVicuna-Mistral
     bofenghuang/vigogne-2-70b-chat
     CohereForAI/c4ai-command-r-plus # Gated
diff --git a/tests/test-polyfills.cpp b/tests/test-polyfills.cpp
