Fix the eos token for gemma. (huggingface#1753)

smartyhouses · Feb 24, 2024 · 8d04f70 · 8d04f70
1 parent eeb7e2b
commit 8d04f70
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/candle-examples/examples/gemma/main.rs b/candle-examples/examples/gemma/main.rs
@@ -66,9 +66,9 @@ impl TextGeneration {
         std::io::stdout().flush()?;
 
         let mut generated_tokens = 0usize;
-        let eos_token = match self.tokenizer.get_token("</s>") {
+        let eos_token = match self.tokenizer.get_token("<eos>") {
             Some(token) => token,
-            None => anyhow::bail!("cannot find the </s> token"),
+            None => anyhow::bail!("cannot find the <eos> token"),
         };
         let start_gen = std::time::Instant::now();
         for index in 0..sample_len {