From 8d04f70f4d1bd67c42fb7d63e7031d49cf780a61 Mon Sep 17 00:00:00 2001
From: Laurent Mazare <laurent.mazare@gmail.com>
Date: Sat, 24 Feb 2024 11:07:02 +0100
Subject: [PATCH] Fix the eos token for gemma. (#1753)

---
 candle-examples/examples/gemma/main.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/candle-examples/examples/gemma/main.rs b/candle-examples/examples/gemma/main.rs
index a4a1b3de..fb59f0bb 100644
--- 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 {