Codeprocastinator/optimized-tinyllama-covalent
0119
1/*2 Copyright 2024 Google LLC3 4 Use of this source code is governed by an MIT-style5 license that can be found in the LICENSE file or at6 https://opensource.org/licenses/MIT.7*/8// SPDX-License-Identifier: MIT9#pragma once10 11#include "minja.hpp"12 13#include <chrono>14#include <cstddef>15#include <cstdio>16#include <exception>17#include <iomanip>18#include <memory>19#include <sstream>20#include <string>21#include <vector>22 23#include <json.hpp>24 25using json = nlohmann::ordered_json;26 27namespace minja {28 29struct chat_template_caps {30 bool supports_tools = false;31 bool supports_tool_calls = false;32 bool supports_tool_responses = false;33 bool supports_system_role = false;34 bool supports_parallel_tool_calls = false;35 bool supports_tool_call_id = false;36 // meta-llama/Llama-3.1-8B-Instruct expects arguments to be an object.37 // Most other templates (and OpenAI's API) expect the arguments object to be stringified.38 bool requires_object_arguments = false;39 // CohereForAI/c4ai-command-r-plus simple variant40 bool requires_non_null_content = false;41 // MiniMaxAI/MiniMax-Text-01 special42 bool requires_typed_content = false;43};44 45struct chat_template_inputs {46 nlohmann::ordered_json messages;47 nlohmann::ordered_json tools;48 bool add_generation_prompt = true;49 nlohmann::ordered_json extra_context;50 std::chrono::system_clock::time_point now = std::chrono::system_clock::now();51};52 53struct chat_template_options {54 bool apply_polyfills = true;55 bool use_bos_token = true;56 bool use_eos_token = true;57 bool define_strftime_now = true;58 59 bool polyfill_tools = true;60 bool polyfill_tool_call_examples = true;61 bool polyfill_tool_calls = true;62 bool polyfill_tool_responses = true;63 bool polyfill_system_role = true;64 bool polyfill_object_arguments = true;65 bool polyfill_typed_content = true;66};67 68class chat_template {69 70 private:71 chat_template_caps caps_;72 std::string source_;73 std::string bos_token_;74 std::string eos_token_;75 std::shared_ptr<minja::TemplateNode> template_root_;76 std::string tool_call_example_;77 78 std::string try_raw_render(79 const nlohmann::ordered_json & messages,80 const nlohmann::ordered_json & tools,81 bool add_generation_prompt,82 const nlohmann::ordered_json & extra_context = nlohmann::ordered_json()) const83 {84 try {85 chat_template_inputs inputs;86 inputs.messages = messages;87 inputs.tools = tools;88 inputs.add_generation_prompt = add_generation_prompt;89 inputs.extra_context = extra_context;90 // Use fixed date for tests91 inputs.now = std::chrono::system_clock::from_time_t(0);92 93 chat_template_options opts;94 opts.apply_polyfills = false;95 96 auto prompt = apply(inputs, opts);97 // fprintf(stderr, "try_raw_render: %s\n", prompt.c_str());98 return prompt;99 } catch (const std::exception & e) {100 // fprintf(stderr, "try_raw_render error: %s\n", e.what());101 return "";102 }103 }104 105 public:106 107 chat_template(const std::string & source, const std::string & bos_token, const std::string & eos_token)108 : source_(source), bos_token_(bos_token), eos_token_(eos_token)109 {110 template_root_ = minja::Parser::parse(source_, {111 /* .trim_blocks = */ true,112 /* .lstrip_blocks = */ true,113 /* .keep_trailing_newline = */ false,114 });115 116 auto contains = [](const std::string & haystack, const std::string & needle) {117 return haystack.find(needle) != std::string::npos;118 };119 120 const std::string user_needle = "<User Needle>";121 const std::string sys_needle = "<System Needle>";122 const json dummy_str_user_msg = {{"role", "user"}, {"content", user_needle}};123 const json dummy_typed_user_msg = {{"role", "user"}, {"content", json::array({{{"type", "text"}, {"text", user_needle}}})}};124 125 caps_.requires_typed_content =126 !contains(try_raw_render(json::array({dummy_str_user_msg}), {}, false), user_needle)127 && contains(try_raw_render(json::array({dummy_typed_user_msg}), {}, false), user_needle);128 129 const auto dummy_user_msg = caps_.requires_typed_content130 ? dummy_typed_user_msg131 : dummy_str_user_msg;132 const json needle_system_msg = {133 {"role", "system"},134 {"content", caps_.requires_typed_content ? json::array({{{"type", "text"}, {"text", sys_needle}}}) : json(sys_needle)},135 };136 137 caps_.supports_system_role = contains(try_raw_render({needle_system_msg, dummy_user_msg,}, {}, false), sys_needle);138 139 auto out = try_raw_render(json::array({140 dummy_user_msg141 }), json::array({142 {143 {"name", "some_tool"},144 {"type", "function"},145 {"function", {146 {"name", "some_tool"},147 {"description", "Some tool."},148 {"parameters", {149 {"type", "object"},150 {"properties", {151 {"arg", {152 {"type", "string"},153 {"description", "Some argument."},154 }},155 }},156 {"required", json::array({ "arg" })},157 }},158 }},159 },160 }), false);161 caps_.supports_tools = contains(out, "some_tool");162 163 auto make_tool_calls_msg = [&](const json & tool_calls) {164 return json {165 {"role", "assistant"},166 {"content", nullptr},167 {"tool_calls", tool_calls},168 };169 };170 auto make_tool_call = [](const std::string & tool_name, const json & arguments) {171 return json {172 {"id", "call_1___"},173 {"type", "function"},174 {"function", {175 {"arguments", arguments},176 {"name", tool_name},177 }},178 };179 };180 const json dummy_args_obj {{"argument_needle", "print('Hello, World!')"}};181 182 // Note: the arguments are rendered in both cases, but may be double-escaped, which we don't want.183 out = try_raw_render(json::array({184 dummy_user_msg,185 make_tool_calls_msg(json::array({make_tool_call("ipython", dummy_args_obj.dump())})),186 }), {}, false);187 auto tool_call_renders_str_arguments = contains(out, "\"argument_needle\":") || contains(out, "'argument_needle':");188 out = try_raw_render(json::array({189 dummy_user_msg,190 make_tool_calls_msg(json::array({make_tool_call("ipython", dummy_args_obj)})),191 }), {}, false);192 auto tool_call_renders_obj_arguments = contains(out, "\"argument_needle\":") || contains(out, "'argument_needle':");193 194 caps_.supports_tool_calls = tool_call_renders_str_arguments || tool_call_renders_obj_arguments;195 caps_.requires_object_arguments = !tool_call_renders_str_arguments && tool_call_renders_obj_arguments;196 auto out_empty = try_raw_render(json::array({dummy_user_msg, {{"role", "assistant"}, {"content", ""}}}), {}, false);197 auto out_null = try_raw_render(json::array({dummy_user_msg, {{"role", "assistant"}, {"content", nullptr}}}), {}, false);198 caps_.requires_non_null_content = contains(out_empty, user_needle) && !contains(out_null, user_needle);199 200 if (caps_.supports_tool_calls) {201 auto dummy_args = caps_.requires_object_arguments ? dummy_args_obj : json(dummy_args_obj.dump());202 auto tc1 = make_tool_call("test_tool1", dummy_args);203 auto tc2 = make_tool_call("test_tool2", dummy_args);204 auto out = try_raw_render(json::array({205 dummy_user_msg,206 make_tool_calls_msg(json::array({tc1, tc2})),207 }), {}, false);208 caps_.supports_parallel_tool_calls = contains(out, "test_tool1") && contains(out, "test_tool2");209 210 out = try_raw_render(json::array({211 dummy_user_msg,212 make_tool_calls_msg(json::array({tc1})),213 {214 {"role", "tool"},215 {"name", "test_tool1"},216 {"content", "Some response!"},217 {"tool_call_id", "call_911_"},218 }219 }), {}, false);220 caps_.supports_tool_responses = contains(out, "Some response!");221 caps_.supports_tool_call_id = contains(out, "call_911_");222 }223 224 try {225 if (!caps_.supports_tools) {226 const json user_msg {227 {"role", "user"},228 {"content", "Hey"},229 };230 const json args {231 {"arg1", "some_value"},232 };233 const json tool_call_msg {234 {"role", "assistant"},235 {"content", nullptr},236 {"tool_calls", json::array({237 {238 // TODO: detect if requires numerical id or fixed length == 6 like Nemo239 {"id", "call_1___"},240 {"type", "function"},241 {"function", {242 {"name", "tool_name"},243 {"arguments", (caps_.requires_object_arguments ? args : json(minja::Value(args).dump(-1, /* to_json= */ true)))},244 }},245 },246 })},247 };248 std::string prefix, full;249 {250 chat_template_inputs inputs;251 inputs.messages = json::array({user_msg});252 inputs.add_generation_prompt = true;253 prefix = apply(inputs);254 }255 {256 chat_template_inputs inputs;257 inputs.messages = json::array({user_msg, tool_call_msg});258 inputs.add_generation_prompt = false;259 full = apply(inputs);260 }261 auto eos_pos_last = full.rfind(eos_token_);262 if (eos_pos_last == prefix.size() - eos_token_.size() ||263 (full[full.size() - 1] == '\n' && (eos_pos_last == full.size() - eos_token_.size() - 1))) {264 full = full.substr(0, eos_pos_last);265 }266 size_t common_prefix_length = 0;267 for (size_t i = 0; i < prefix.size() && i < full.size(); ++i) {268 if (prefix[i] != full[i]) {269 break;270 }271 if (prefix[i] == '<') {272 // DeepSeek R1's template (as of 20250209) adds a trailing <think> if add_generation_prompt,273 // but it removes thinking tags for past messages.274 // The prefix and full strings diverge at <think> vs. <｜tool▁calls▁begin｜>, we avoid consuming the leading <.275 continue;276 }277 common_prefix_length = i + 1;278 }279 auto example = full.substr(common_prefix_length);280 if (example.find("tool_name") == std::string::npos && example.find("some_value") == std::string::npos) {281 fprintf(stderr, "Failed to infer a tool call example (possible template bug)\n");282 } else {283 tool_call_example_ = example;284 }285 }286 } catch (const std::exception & e) {287 fprintf(stderr, "Failed to generate tool call example: %s\n", e.what());288 }289 }290 291 const std::string & source() const { return source_; }292 const std::string & bos_token() const { return bos_token_; }293 const std::string & eos_token() const { return eos_token_; }294 const chat_template_caps & original_caps() const { return caps_; }295 296 // Deprecated, please use the form with chat_template_inputs and chat_template_options297 std::string apply(298 const nlohmann::ordered_json & messages,299 const nlohmann::ordered_json & tools,300 bool add_generation_prompt,301 const nlohmann::ordered_json & extra_context = nlohmann::ordered_json(),302 bool apply_polyfills = true)303 {304 fprintf(stderr, "[%s] Deprecated!\n", __func__);305 chat_template_inputs inputs;306 inputs.messages = messages;307 inputs.tools = tools;308 inputs.add_generation_prompt = add_generation_prompt;309 inputs.extra_context = extra_context;310 inputs.now = std::chrono::system_clock::now();311 312 chat_template_options opts;313 opts.apply_polyfills = apply_polyfills;314 315 return apply(inputs, opts);316 }317 318 std::string apply(319 const chat_template_inputs & inputs,320 const chat_template_options & opts = chat_template_options()) const321 {322 json actual_messages;323 324 auto has_tools = inputs.tools.is_array() && !inputs.tools.empty();325 auto has_tool_calls = false;326 auto has_tool_responses = false;327 auto has_string_content = false;328 for (const auto & message : inputs.messages) {329 if (message.contains("tool_calls") && !message["tool_calls"].is_null()) {330 has_tool_calls = true;331 }332 if (message.contains("role") && message["role"] == "tool") {333 has_tool_responses = true;334 }335 if (message.contains("content") && message["content"].is_string()) {336 has_string_content = true;337 }338 }339 340 auto polyfill_system_role = opts.polyfill_system_role && !caps_.supports_system_role;341 auto polyfill_tools = opts.polyfill_tools && has_tools && !caps_.supports_tools;342 auto polyfill_tool_call_example = polyfill_tools && opts.polyfill_tool_call_examples;343 auto polyfill_tool_calls = opts.polyfill_tool_calls && has_tool_calls && !caps_.supports_tool_calls;344 auto polyfill_tool_responses = opts.polyfill_tool_responses && has_tool_responses && !caps_.supports_tool_responses;345 auto polyfill_object_arguments = opts.polyfill_object_arguments && has_tool_calls && caps_.requires_object_arguments;346 auto polyfill_typed_content = opts.polyfill_typed_content && has_string_content && caps_.requires_typed_content;347 348 auto needs_polyfills = opts.apply_polyfills && (false349 || polyfill_system_role350 || polyfill_tools351 || polyfill_tool_calls352 || polyfill_tool_responses353 || polyfill_object_arguments354 || polyfill_typed_content355 );356 357 if (needs_polyfills) {358 actual_messages = json::array();359 360 auto add_message = [&](const json & msg) {361 if (polyfill_typed_content && msg.contains("content") && !msg.at("content").is_null() && msg.at("content").is_string()) {362 actual_messages.push_back({363 {"role", msg.at("role")},364 {"content", {{365 {"type", "text"},366 {"text", msg.at("content")},367 }}},368 });369 } else {370 actual_messages.push_back(msg);371 }372 };373 374 std::string pending_system;375 auto flush_sys = [&]() {376 if (!pending_system.empty()) {377 add_message({378 {"role", "user"},379 {"content", pending_system},380 });381 pending_system.clear();382 }383 };384 385 json adjusted_messages;386 if (polyfill_tools) {387 adjusted_messages = add_system(inputs.messages,388 "You can call any of the following tools to satisfy the user's requests: " + minja::Value(inputs.tools).dump(2, /* to_json= */ true) +389 (!polyfill_tool_call_example || tool_call_example_.empty() ? "" : "\n\nExample tool call syntax:\n\n" + tool_call_example_ + "\n\n"));390 } else {391 adjusted_messages = inputs.messages;392 }393 394 for (const auto & message_ : adjusted_messages) {395 auto message = message_;396 if (!message.contains("role") || !message.contains("content")) {397 throw std::runtime_error("message must have 'role' and 'content' fields: " + message.dump());398 }399 std::string role = message.at("role");400 401 if (message.contains("tool_calls")) {402 if (polyfill_object_arguments || polyfill_tool_calls) {403 for (auto & tool_call : message.at("tool_calls")) {404 if (tool_call["type"] == "function") {405 auto & function = tool_call.at("function");406 auto & arguments = function.at("arguments");407 if (arguments.is_string()) {408 try {409 arguments = json::parse(arguments.get<std::string>());410 } catch (const std::exception & ecvt) {411 fprintf(stderr, "Failed to parse arguments: %s\n", ecvt.what());412 }413 }414 }415 }416 }417 if (polyfill_tool_calls) {418 auto content = message.at("content");419 auto tool_calls = json::array();420 for (const auto & tool_call : message.at("tool_calls")) {421 if (tool_call.at("type") != "function") {422 continue;423 }424 const auto & function = tool_call.at("function");425 auto tc = json {426 {"name", function.at("name")},427 {"arguments", function.at("arguments")},428 };429 if (tool_call.contains("id")) {430 tc["id"] = tool_call["id"];431 }432 tool_calls.push_back(tc);433 }434 auto obj = json {435 {"tool_calls", tool_calls},436 };437 if (!content.is_null() && !content.empty()) {438 obj["content"] = content;439 }440 message["content"] = obj.dump(2);441 message.erase("tool_calls");442 }443 }444 if (polyfill_tool_responses && role == "tool") {445 message["role"] = "user";446 auto obj = json {447 {"tool_response", json::object()},448 };449 if (message.contains("name")) {450 obj["tool_response"]["tool"] = message.at("name");451 }452 obj["tool_response"]["content"] = message.at("content");453 if (message.contains("tool_call_id")) {454 obj["tool_response"]["tool_call_id"] = message.at("tool_call_id");455 }456 message["content"] = obj.dump(2);457 message.erase("name");458 }459 460 if (!message["content"].is_null() && polyfill_system_role) {461 std::string content = message.at("content");462 if (role == "system") {463 if (!pending_system.empty()) pending_system += "\n";464 pending_system += content;465 continue;466 } else {467 if (role == "user") {468 if (!pending_system.empty()) {469 message["content"] = pending_system + (content.empty() ? "" : "\n" + content);470 pending_system.clear();471 }472 } else {473 flush_sys();474 }475 }476 }477 add_message(message);478 }479 flush_sys();480 } else {481 actual_messages = inputs.messages;482 }483 484 auto context = minja::Context::make(json({485 {"messages", actual_messages},486 {"add_generation_prompt", inputs.add_generation_prompt},487 }));488 context->set("bos_token", opts.use_bos_token ? bos_token_ : "");489 context->set("eos_token", opts.use_eos_token ? eos_token_ : "");490 if (opts.define_strftime_now) {491 auto now = inputs.now;492 context->set("strftime_now", Value::callable([now](const std::shared_ptr<minja::Context> &, minja::ArgumentsValue & args) {493 args.expectArgs("strftime_now", {1, 1}, {0, 0});494 auto format = args.args[0].get<std::string>();495 496 auto time = std::chrono::system_clock::to_time_t(now);497 auto local_time = *std::localtime(&time);498 std::ostringstream ss;499 ss << std::put_time(&local_time, format.c_str());500 return ss.str();501 }));502 }503 if (!inputs.tools.is_null()) {504 context->set("tools", minja::Value(inputs.tools));505 }506 if (!inputs.extra_context.is_null()) {507 for (auto & kv : inputs.extra_context.items()) {508 context->set(kv.key(), minja::Value(kv.value()));509 }510 }511 512 auto ret = template_root_->render(context);513 // fprintf(stderr, "actual_messages: %s\n", actual_messages.dump(2).c_str());514 // fprintf(stderr, "apply: %s\n\n", ret.c_str());515 return ret;516 }517 518 static nlohmann::ordered_json add_system(const nlohmann::ordered_json & messages, const std::string & system_prompt) {519 json messages_with_system = messages;520 521 if (!messages_with_system.empty() && messages_with_system[0].at("role") == "system") {522 std::string existing_system = messages_with_system.at(0).at("content");523 messages_with_system[0] = json {524 {"role", "system"},525 {"content", existing_system + "\n\n" + system_prompt},526 };527 } else {528 messages_with_system.insert(messages_with_system.begin(), json {529 {"role", "system"},530 {"content", system_prompt},531 });532 }533 return messages_with_system;534 }535};536 537} // namespace minja538