Skip to content

Commit c64d2be

Browse files
ochafikochafik
authored
1 parent 96f4053 commit c64d2be

File tree

2 files changed

+54
-11
lines changed

2 files changed

+54
-11
lines changed

‎common/chat-template.hpp

+28-9
Original file line numberDiff line numberDiff line change
@@ -25,21 +25,22 @@ class chat_template {
2525
// Meta-Llama-3.1-8B-Instruct's template expects arguments to be an object.
2626
// Most other templates (and OpenAI's API) expect the arguments object to be stringified.
2727
bool requires_object_arguments_ = false;
28+
bool requires_typed_content_ = false;
2829
bool supports_system_role_ = true;
2930
bool supports_parallel_tool_calls_ = false;
3031
std::string source_;
3132
std::string bos_token_;
3233
std::string eos_token_;
3334
std::shared_ptr<minja::TemplateNode> template_root_;
3435

35-
std::string try_render(
36+
std::string try_raw_render(
3637
const nlohmann::ordered_json & messages,
3738
const nlohmann::ordered_json & tools,
3839
bool add_generation_prompt,
3940
const nlohmann::ordered_json & extra_context = nlohmann::ordered_json()) const
4041
{
4142
try {
42-
auto prompt = apply(messages, tools, add_generation_prompt, extra_context);
43+
auto prompt = apply(messages, tools, add_generation_prompt, extra_context, /* adjust_inputs= */ false);
4344
// fprintf(stderr, "Prompt: %s\n", prompt.c_str());
4445
return prompt;
4546
} catch (const std::exception & e) {
@@ -60,7 +61,7 @@ class chat_template {
6061
supports_tools_ = source.find("tools") != std::string::npos;
6162

6263
auto renders_string_arguments =
63-
try_render({
64+
try_raw_render({
6465
{
6566
{"role", "user"},
6667
{"content", "Hey"}
@@ -81,7 +82,7 @@ class chat_template {
8182
}, {}, false).find("{\"code\": \"print") != std::string::npos;
8283
if (!renders_string_arguments) {
8384
auto renders_object_arguments =
84-
try_render({
85+
try_raw_render({
8586
{
8687
{"role", "user"},
8788
{"content", "Hey"}
@@ -106,10 +107,13 @@ class chat_template {
106107
}
107108
supports_parallel_tool_calls_ = source.find("tool_call_id") != std::string::npos;
108109

109-
supports_system_role_ = try_render({
110+
supports_system_role_ = try_raw_render({
110111
{{"role", "system"}, {"content", "<System Needle>"}},
111112
{{"role", "user"}, {"content", "Hey"}}
112113
}, {}, false).find("<System Needle>") != std::string::npos;
114+
115+
requires_typed_content_ = try_raw_render({{{"role", "user"}, {"content", "Hey"}}}, {}, false).find("Hey") == std::string::npos
116+
&& try_raw_render({{{"role", "user"}, {"content", {{{"type", "text"}, {"text", "Hey"}}}}}}, {}, false).find("Hey") != std::string::npos;
113117
}
114118

115119
const std::string & source() const { return source_; }
@@ -122,19 +126,34 @@ class chat_template {
122126
const nlohmann::ordered_json & messages,
123127
const nlohmann::ordered_json & tools,
124128
bool add_generation_prompt,
125-
const nlohmann::ordered_json & extra_context = nlohmann::ordered_json()) const
129+
const nlohmann::ordered_json & extra_context = nlohmann::ordered_json(),
130+
bool adjust_inputs = true) const
126131
{
127132
json actual_messages;
128133

129134
// First, "fix" messages so they have a chance to be rendered correctly by the template
130135

131-
if (requires_object_arguments_ || !supports_system_role_ || !supports_tools_) {
136+
if (adjust_inputs && (requires_object_arguments_ || !supports_system_role_ || !supports_tools_ || requires_typed_content_)) {
132137
actual_messages = json::array();
133138

139+
auto add_message = [&](const json & msg) {
140+
if (requires_typed_content_ && msg.contains("content") && !msg.at("content").is_null() && msg.at("content").is_string()) {
141+
actual_messages.push_back({
142+
{"role", msg.at("role")},
143+
{"content", {{
144+
{"type", "text"},
145+
{"text", msg.at("content")},
146+
}}},
147+
});
148+
} else {
149+
actual_messages.push_back(msg);
150+
}
151+
};
152+
134153
std::string pending_system;
135154
auto flush_sys = [&]() {
136155
if (!pending_system.empty()) {
137-
actual_messages.push_back({
156+
add_message({
138157
{"role", "user"},
139158
{"content", pending_system},
140159
});
@@ -217,7 +236,7 @@ class chat_template {
217236
}
218237
}
219238
}
220-
actual_messages.push_back(message);
239+
add_message(message);
221240
}
222241
flush_sys();
223242
} else {

‎common/minja.hpp

+26-2
Original file line numberDiff line numberDiff line change
@@ -693,7 +693,7 @@ enum SpaceHandling { Keep, Strip, StripSpaces, StripNewline };
693693

694694
class TemplateToken {
695695
public:
696-
enum class Type { Text, Expression, If, Else, Elif, EndIf, For, EndFor, Set, EndSet, Comment, Macro, EndMacro, Filter, EndFilter };
696+
enum class Type { Text, Expression, If, Else, Elif, EndIf, For, EndFor, Generation, EndGeneration, Set, EndSet, Comment, Macro, EndMacro, Filter, EndFilter };
697697

698698
static std::string typeToString(Type t) {
699699
switch (t) {
@@ -712,6 +712,8 @@ class TemplateToken {
712712
case Type::EndMacro: return "endmacro";
713713
case Type::Filter: return "filter";
714714
case Type::EndFilter: return "endfilter";
715+
case Type::Generation: return "generation";
716+
case Type::EndGeneration: return "endgeneration";
715717
}
716718
return "Unknown";
717719
}
@@ -788,6 +790,14 @@ struct EndForTemplateToken : public TemplateToken {
788790
EndForTemplateToken(const Location & location, SpaceHandling pre, SpaceHandling post) : TemplateToken(Type::EndFor, location, pre, post) {}
789791
};
790792

793+
struct GenerationTemplateToken : public TemplateToken {
794+
GenerationTemplateToken(const Location & location, SpaceHandling pre, SpaceHandling post) : TemplateToken(Type::Generation, location, pre, post) {}
795+
};
796+
797+
struct EndGenerationTemplateToken : public TemplateToken {
798+
EndGenerationTemplateToken(const Location & location, SpaceHandling pre, SpaceHandling post) : TemplateToken(Type::EndGeneration, location, pre, post) {}
799+
};
800+
791801
struct SetTemplateToken : public TemplateToken {
792802
std::string ns;
793803
std::vector<std::string> var_names;
@@ -2149,7 +2159,7 @@ class Parser {
21492159
static std::regex comment_tok(R"(\{#([-~]?)(.*?)([-~]?)#\})");
21502160
static std::regex expr_open_regex(R"(\{\{([-~])?)");
21512161
static std::regex block_open_regex(R"(^\{%([-~])?[\s\n\r]*)");
2152-
static std::regex block_keyword_tok(R"((if|else|elif|endif|for|endfor|set|endset|block|endblock|macro|endmacro|filter|endfilter)\b)");
2162+
static std::regex block_keyword_tok(R"((if|else|elif|endif|for|endfor|generation|endgeneration|set|endset|block|endblock|macro|endmacro|filter|endfilter)\b)");
21532163
static std::regex non_text_open_regex(R"(\{\{|\{%|\{#)");
21542164
static std::regex expr_close_regex(R"([\s\n\r]*([-~])?\}\})");
21552165
static std::regex block_close_regex(R"([\s\n\r]*([-~])?%\})");
@@ -2229,6 +2239,12 @@ class Parser {
22292239
} else if (keyword == "endfor") {
22302240
auto post_space = parseBlockClose();
22312241
tokens.push_back(std::make_unique<EndForTemplateToken>(location, pre_space, post_space));
2242+
} else if (keyword == "generation") {
2243+
auto post_space = parseBlockClose();
2244+
tokens.push_back(std::make_unique<GenerationTemplateToken>(location, pre_space, post_space));
2245+
} else if (keyword == "endgeneration") {
2246+
auto post_space = parseBlockClose();
2247+
tokens.push_back(std::make_unique<EndGenerationTemplateToken>(location, pre_space, post_space));
22322248
} else if (keyword == "set") {
22332249
static std::regex namespaced_var_regex(R"((\w+)[\s\n\r]*\.[\s\n\r]*(\w+))");
22342250

@@ -2330,6 +2346,13 @@ class Parser {
23302346
throw unterminated(**start);
23312347
}
23322348
children.emplace_back(std::make_shared<ForNode>(token->location, std::move(for_token->var_names), std::move(for_token->iterable), std::move(for_token->condition), std::move(body), for_token->recursive, std::move(else_body)));
2349+
} else if (dynamic_cast<GenerationTemplateToken*>(token.get())) {
2350+
auto body = parseTemplate(begin, it, end);
2351+
if (it == end || (*(it++))->type != TemplateToken::Type::EndGeneration) {
2352+
throw unterminated(**start);
2353+
}
2354+
// Treat as a no-op, as our scope is templates for inference, not training (`{% generation %}` wraps generated tokens for masking).
2355+
children.emplace_back(std::move(body));
23332356
} else if (auto text_token = dynamic_cast<TextTemplateToken*>(token.get())) {
23342357
SpaceHandling pre_space = (it - 1) != begin ? (*(it - 2))->post_space : SpaceHandling::Keep;
23352358
SpaceHandling post_space = it != end ? (*it)->pre_space : SpaceHandling::Keep;
@@ -2397,6 +2420,7 @@ class Parser {
23972420
|| dynamic_cast<EndFilterTemplateToken*>(token.get())
23982421
|| dynamic_cast<EndIfTemplateToken*>(token.get())
23992422
|| dynamic_cast<ElseTemplateToken*>(token.get())
2423+
|| dynamic_cast<EndGenerationTemplateToken*>(token.get())
24002424
|| dynamic_cast<ElifTemplateToken*>(token.get())) {
24012425
it--; // unconsume the token
24022426
break; // exit the loop

0 commit comments

Comments
 (0)