Skip to content

Commit d42deaa

Browse files
committed
Set the secondary EOS for Gemma2
So that we can remove the `<end_of_turn>` filter that was set up specifically for Gemma2.
1 parent 2bad79f commit d42deaa

File tree

2 files changed

+2
-15
lines changed

2 files changed

+2
-15
lines changed

gemma/configs.cc

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,8 @@ static ModelConfig ConfigBaseGemmaV2() {
3535
ModelConfig config = ConfigNoSSM();
3636
config.att_cap = 50.0f;
3737
config.final_cap = 30.0f;
38+
config.eos_id = 1;
39+
config.secondary_eos_id = 107;
3840
return config;
3941
}
4042

gemma/run.cc

Lines changed: 0 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,6 @@ void ReplGemma(Gemma& model, KVCache& kv_cache, const AppArgs& app,
8585
size_t abs_pos = 0; // across turns
8686
size_t tokens_generated_this_turn = 0; // differentiates prefill from reply
8787
size_t prompt_size = 0;
88-
bool end_of_turn_seen = false;
8988

9089
std::mt19937 gen;
9190
InitGenerator(args, gen);
@@ -140,13 +139,6 @@ void ReplGemma(Gemma& model, KVCache& kv_cache, const AppArgs& app,
140139
std::cout << "\n\n";
141140
}
142141
}
143-
if (token_text == "<end_of_turn>") {
144-
// We don't want to show the <end_of_turn> token to the user.
145-
// We also need to remember that we've seen it, so that we can rewind
146-
// abs_pos appropriately. We expect EOS as the next token.
147-
end_of_turn_seen = true;
148-
return true;
149-
}
150142
std::cout << token_text << std::flush;
151143
return true;
152144
};
@@ -232,13 +224,6 @@ void ReplGemma(Gemma& model, KVCache& kv_cache, const AppArgs& app,
232224
HWY_ASSERT(abs_pos > 0);
233225
abs_pos--;
234226
}
235-
if (end_of_turn_seen && abs_pos > 0) {
236-
// If we have seen an end_of_turn token, we need to rewind abs_pos by one
237-
// more, because we will prepend it again to the prompt in
238-
// WrapAndTokenize.
239-
abs_pos--;
240-
}
241-
end_of_turn_seen = false;
242227
}
243228
}
244229

0 commit comments

Comments
 (0)