From ad7c53953b6aadf7d741b2263aa2fd32223b023b Mon Sep 17 00:00:00 2001
From: Laurent Mazare <laurent.mazare@gmail.com>
Date: Thu, 17 Aug 2023 15:26:44 +0100
Subject: [PATCH] Add a verbose-prompt mode, similar to llama.cpp. (#489)

---
 candle-examples/examples/quantized/main.rs | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/candle-examples/examples/quantized/main.rs b/candle-examples/examples/quantized/main.rs
index 94efb03f..b9b2ec9a 100644
--- a/candle-examples/examples/quantized/main.rs
+++ b/candle-examples/examples/quantized/main.rs
@@ -303,6 +303,10 @@ struct Args {
     /// Enable tracing (generates a trace-timestamp.json file).
     #[arg(long)]
     tracing: bool,
+
+    /// Display the token for the specified prompt.
+    #[arg(long)]
+    verbose_prompt: bool,
 }
 
 impl Args {
@@ -375,11 +379,15 @@ fn main() -> anyhow::Result<()> {
 
     let tokenizer = args.tokenizer()?;
     let prompt = args.prompt.as_ref().map_or(DEFAULT_PROMPT, |p| p.as_str());
-    let mut tokens = tokenizer
-        .encode(prompt, true)
-        .map_err(anyhow::Error::msg)?
-        .get_ids()
-        .to_vec();
+    let tokens = tokenizer.encode(prompt, true).map_err(anyhow::Error::msg)?;
+    if args.verbose_prompt {
+        for (token, id) in tokens.get_tokens().iter().zip(tokens.get_ids().iter()) {
+            let token = token.replace('▁', " ").replace("<0x0A>", "\n");
+            println!("{id:7} -> '{token}'");
+        }
+    }
+
+    let mut tokens = tokens.get_ids().to_vec();
     let mut index_pos = 0;
     let mut logits_processor = LogitsProcessor::new(args.seed, args.temperature);
     let start_gen = std::time::Instant::now();