Skip to content

Commit 379a22d

Browse files
committed
Fix EOS
1 parent 93aa3c4 commit 379a22d

1 file changed

Lines changed: 4 additions & 0 deletions

File tree

examples/models/gemma4_31b/main.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,10 @@ int main(int argc, char** argv) {
224224

225225
auto eos_ids = llm::get_eos_ids(tokenizer.get(), module.get());
226226
eos_ids.insert(static_cast<uint64_t>(FLAGS_eos_id));
227+
auto turn_ids = tokenizer->encode("<turn|>", /*bos=*/0, /*eos=*/0);
228+
if (turn_ids.ok() && turn_ids->size() == 1) {
229+
eos_ids.insert(turn_ids.get()[0]);
230+
}
227231

228232
// Read prompt
229233
std::string prompt_text = FLAGS_prompt;

0 commit comments

Comments
 (0)