Skip to content

Commit 10e9780

Browse files
ngxsonggerganov
andauthored
chat: fix int overflow, prevent size calculation in float/double (ggml-org#17357)
* chat: fix int overflow, prevent size calculation in float/double * Update common/chat.cpp Co-authored-by: Georgi Gerganov <[email protected]> --------- Co-authored-by: Georgi Gerganov <[email protected]>
1 parent a045492 commit 10e9780

File tree

1 file changed

+8
-2
lines changed

1 file changed

+8
-2
lines changed

common/chat.cpp

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3359,7 +3359,7 @@ static common_chat_params common_chat_templates_apply_legacy(
33593359
const struct common_chat_templates * tmpls,
33603360
const struct common_chat_templates_inputs & inputs)
33613361
{
3362-
int alloc_size = 0;
3362+
size_t alloc_size = 0;
33633363
std::vector<llama_chat_message> chat;
33643364
std::vector<std::string> contents;
33653365

@@ -3381,7 +3381,8 @@ static common_chat_params common_chat_templates_apply_legacy(
33813381
const auto & msg = inputs.messages[i];
33823382
const auto & content = contents[i];
33833383
chat.push_back({msg.role.c_str(), content.c_str()});
3384-
alloc_size += (msg.role.size() + content.size()) * 1.25;
3384+
size_t msg_size = msg.role.size() + content.size();
3385+
alloc_size += msg_size + (msg_size / 4); // == msg_size * 1.25 but avoiding float ops
33853386
}
33863387

33873388
std::vector<char> buf(alloc_size);
@@ -3403,6 +3404,11 @@ static common_chat_params common_chat_templates_apply_legacy(
34033404
res = llama_chat_apply_template(src.c_str(), chat.data(), chat.size(), inputs.add_generation_prompt, buf.data(), buf.size());
34043405
}
34053406

3407+
// for safety, we check the result again
3408+
if (res < 0 || (size_t) res > buf.size()) {
3409+
throw std::runtime_error("failed to apply chat template, try using --jinja");
3410+
}
3411+
34063412
common_chat_params params;
34073413
params.prompt = std::string(buf.data(), res);
34083414
if (!inputs.json_schema.empty()) {

0 commit comments

Comments
 (0)