make up for the missing last token output of phi2 example (huggingfac…

…e#2299)
spiceai · Jun 29, 2024 · 74e9e41 · 74e9e41
1 parent e27aac0
commit 74e9e41
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/candle-examples/examples/phi/main.rs b/candle-examples/examples/phi/main.rs
@@ -114,6 +114,10 @@ impl TextGeneration {
             tokens.push(next_token);
             generated_tokens += 1;
             if next_token == eos_token {
+                if let Some(t) = self.tokenizer.decode_rest()? {
+                    print!("{t}");
+                    std::io::stdout().flush()?;
+                }
                 break;
             }
             if let Some(t) = self.tokenizer.next_token(next_token)? {