red-candle 1.0.0.pre.5 → 1.0.0.pre.7

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 91a4c43a1a12d6d8960f1a1d190c9bfe8ea60db75f687233012d09b8c90b5020
-  data.tar.gz: 3f6ce143cd38856365231baebe25a188e7d5824d930ecdae79c1660b3ad6c787
+  metadata.gz: 6bea0c9f27d5bbbc43e0c2e8fa5456b2b79511f164f7083512025538ca172077
+  data.tar.gz: b591f559c63c1bdb9fa96beaf1079186284b219e2ceba50d8e2891fa4f3096bd
 SHA512:
-  metadata.gz: 273a01c438b085509a433602097b5ad4bcdb3420fc19ebe84c4fd37bf43ae5e6e1040701ecaad9e2864c0cda31d6128d4b2df6c395290994c17f135620282be6
-  data.tar.gz: 9193d01d8bfc704b982c839f9aebac23217ae5d56b5afae6228daae73b60e219b049fe8fc8d5ad13657ceb110476489c68bbb9aa6397688bf70a976a8d62d41e
+  metadata.gz: 380cde32abf995f9766056638d4d0612acef746a84b31451c0aa2fb8cd205fc76b76b8909729a9af267ee6c549f9149253b2bfb70717385b5afe2cd2ad0b3152
+  data.tar.gz: 7acdfc4de263501be51e3f74933c528b9d723b4da93e78e3abae2454dc2746fb7696c09722fe6666084a48dff098501b2d1af503963b5d739e879bd0fcb96e43

data/README.md

@@ -47,9 +47,46 @@ Red-Candle now supports Large Language Models (LLMs) with GPU acceleration!
 
 ### Supported Models
 
+- **Gemma**: Google's Gemma models (e.g., `google/gemma-2b`, `google/gemma-7b`, `google/gemma-2b-it`)
 - **Llama**: Llama 2 and Llama 3 models (e.g., `TinyLlama/TinyLlama-1.1B-Chat-v1.0`, `meta-llama/Llama-2-7b-hf`, `NousResearch/Llama-2-7b-hf`)
 - **Mistral**: All Mistral models (e.g., `mistralai/Mistral-7B-Instruct-v0.1`)
 
+### Quantized Model Support (GGUF)
+
+Red-Candle supports quantized models in GGUF format, offering 4-8x memory reduction:
+
+> **Note on GGUF Support**: Red-Candle now uses a unified GGUF loader that automatically detects the model architecture from the GGUF file. This means all GGUF models (including Mistral models from TheBloke) should now work correctly! The loader automatically selects the appropriate tokenizer based on the model type to ensure proper text generation.
+
+```ruby
+# Load quantized models - always specify the GGUF filename
+llm = Candle::LLM.from_pretrained("TheBloke/Llama-2-7B-Chat-GGUF", 
+                                  device: device, 
+                                  gguf_file: "llama-2-7b-chat.Q4_K_M.gguf")
+
+# Register custom tokenizer mappings for your models
+Candle::LLM.register_tokenizer("my-org/my-model-GGUF", "my-org/my-tokenizer")
+
+# Popular quantized model sources:
+# - TheBloke: Extensive collection of GGUF models
+# - Search HuggingFace for "GGUF" models
+```
+
+**Memory usage comparison (7B models):**
+- Full precision: ~28 GB
+- Q8_0 (8-bit): ~7 GB - Best quality, larger size
+- Q5_K_M (5-bit): ~4.5 GB - Very good quality  
+- Q4_K_M (4-bit): ~4 GB - Recommended default, best balance
+- Q3_K_M (3-bit): ~3 GB - Good for memory-constrained systems
+
+**Quantization levels explained:**
+- **Q8_0**: Almost identical to full model, use when quality is paramount
+- **Q5_K_M**: Excellent quality with good compression
+- **Q4_K_M**: Best balance of quality/size/speed (recommended default)
+- **Q3_K_M**: Noticeable quality reduction but very compact
+- **Q2_K**: ⚠️ **Not recommended** - Can cause inference errors due to extreme quantization
+
+> **Warning**: Q2_K quantization can lead to "weight is negative, too large or not a valid number" errors during inference. Use Q3_K_M or higher for stable operation.
+
 > ### ⚠️ Huggingface login warning
 > 
 > Many models, including the one below, require you to agree to the terms. You'll need to:
@@ -67,10 +104,10 @@ device = Candle::Device.cpu     # CPU (default)
 device = Candle::Device.metal   # Apple GPU (Metal)
 device = Candle::Device.cuda    # NVIDIA GPU (CUDA)
 
-# Load a Llama model
-llm = Candle::LLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", device: device)
-# Or a Mistral model
-llm = Candle::LLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", device: device)
+# Load a model
+llm = Candle::LLM.from_pretrained("google/gemma-2b-it", device: device)  # Gemma
+# llm = Candle::LLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", device: device)  # Llama
+# llm = Candle::LLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", device: device)  # Mistral
 
 # Generate text
 response = llm.generate("What is Ruby?", config: Candle::GenerationConfig.balanced)
@@ -102,6 +139,33 @@ device = Candle::Device.metal
 device = Candle::Device.cuda   # Linux/Windows with NVIDIA GPU
 ```
 
+### Debugging Token Generation
+
+For debugging purposes, you can enable raw token output to see both token IDs and their raw representations:
+
+```ruby
+# Enable debug mode to see raw tokens during generation
+config = Candle::GenerationConfig.balanced(debug_tokens: true)
+
+# Non-streaming generation with debug tokens
+result = llm.generate("Hello, world!", config: config)
+puts result
+# Output: [15043:Hello][11:,][1917:world][0:!]
+
+# Streaming generation with debug tokens
+llm.generate_stream("Hello, world!", config: config) do |text|
+  print text  # Will show each token as it's generated: [15043:Hello][11:,][1917:world][0:!]
+end
+
+# Works with all models (Llama, Mistral, Gemma, and quantized GGUF models)
+```
+
+This is particularly useful for:
+- Debugging tokenization issues
+- Understanding how the model processes text
+- Troubleshooting generation problems
+- Analyzing model behavior
+
 ## ⚠️ Model Format Requirement: Safetensors Only
 
 Red-Candle **only supports embedding models that provide their weights in the [safetensors](https://github.com/huggingface/safetensors) format** (i.e., the model repo must contain a `model.safetensors` file). If the model repo does not provide the required file, loading will fail with a clear error. Most official BERT and DistilBERT models do **not** provide safetensors; many Sentence Transformers and JinaBERT models do.
@@ -287,6 +351,102 @@ The reranker uses a BERT-based architecture that:
 
 This joint processing allows cross-encoders to capture subtle semantic relationships between queries and documents, making them more accurate for reranking tasks, though at the cost of higher computational requirements.
 
+## Common Runtime Errors
+
+### 1. Weight is negative, too large or not a valid number
+
+**Error:**
+```
+/Users/cpetersen/src/scientist/red-candle/lib/candle/llm.rb:25:in `_generate_stream': Generation failed: A weight is negative, too large or not a valid number (RuntimeError)
+    from /Users/cpetersen/src/scientist/red-candle/lib/candle/llm.rb:25:in `generate_stream'
+    ...
+```
+
+**Cause:** This error occurs when using overly aggressive quantization levels (particularly Q2_K) that result in numerical instability during inference. The 2-bit quantization can cause weights to become corrupted or produce NaN/Inf values.
+
+**Solution:** Use a higher quantization level. Recommended options:
+- Q4_K_M (4-bit) - Best balance of quality and size
+- Q5_K_M (5-bit) - Higher quality with slightly larger size
+- Q3_K_M (3-bit) - Minimum recommended quantization
+
+```ruby
+# Instead of Q2_K:
+llm = Candle::LLM.from_pretrained("TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF", 
+                                  device: device, 
+                                  gguf_file: "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf")
+```
+
+### 2. Cannot find tensor model.embed_tokens.weight
+
+**Error:**
+```
+Failed to load quantized model: cannot find tensor model.embed_tokens.weight (RuntimeError)
+```
+
+**Cause:** This error was common in earlier versions when loading GGUF files with incompatible tensor naming conventions. The unified GGUF loader in version 1.0.0+ should handle most GGUF files correctly.
+
+**If you still encounter this error:**
+1. Ensure you're using the latest version of red-candle (1.0.0 or higher)
+2. Make sure to specify the exact GGUF filename:
+   ```ruby
+   llm = Candle::LLM.from_pretrained("TheBloke/Mistral-7B-Instruct-v0.2-GGUF", 
+                                     device: device,
+                                     gguf_file: "mistral-7b-instruct-v0.2.Q4_K_M.gguf")
+   ```
+3. If the error persists, the GGUF file may use an unsupported architecture or format
+
+### 3. No GGUF file found in repository
+
+**Error:**
+```
+Failed to load quantized model: No GGUF file found in repository TheBloke/model-name-GGUF. Try specifying a quantization level like Q4_K_M, Q5_K_M, or Q8_0. (RuntimeError)
+```
+
+**Cause:** The automatic GGUF file detection couldn't find a matching file, often due to naming variations.
+
+**Solution:** Specify the exact GGUF filename:
+```ruby
+# Visit the HuggingFace repository to find the exact filename
+llm = Candle::LLM.from_pretrained("TheBloke/Llama-2-7B-Chat-GGUF", 
+                                  device: device, 
+                                  gguf_file: "llama-2-7b-chat.Q4_K_M.gguf")
+```
+
+### 4. Failed to download tokenizer
+
+**Error:**
+```
+Failed to load quantized model: Failed to download tokenizer: request error: HTTP status client error (404 Not Found)
+```
+
+**Cause:** GGUF repositories often don't include separate tokenizer files since they're embedded in the GGUF format.
+
+**Solution:** The code now includes fallback tokenizer loading. If you still encounter this error, ensure you're using the latest version of red-candle.
+
+### 5. Missing metadata in GGUF file
+
+**Error:**
+```
+Failed to load GGUF model: cannot find gemma3.attention.head_count in metadata (RuntimeError)
+```
+or
+```
+Failed to load GGUF model: cannot find llama.attention.head_count in metadata (RuntimeError)
+```
+
+**Cause:** Some GGUF files may have been created with older conversion tools that don't include all required metadata fields.
+
+**Solution:** 
+- Try a different GGUF file from the same model
+- Look for GGUF files from TheBloke or other reputable sources
+- Check if a newer version of the GGUF file is available
+- Some Gemma GGUF files may not be compatible with the current loader
+
+**Known compatibility issues:**
+- `lmstudio-ai/gemma-2b-it-GGUF` - Missing required metadata fields
+- Gemma 3 GGUF files may require specific tokenizers that are not publicly available
+- For best compatibility, use Llama or Mistral GGUF files from TheBloke
+
 ## Development
 
 FORK IT!

data/Rakefile

@@ -8,7 +8,7 @@ task default: :test
 Rake::TestTask.new do |t|
   t.deps << :compile
   t.libs << "test"
-  t.test_files = FileList["test/**/*_test.rb"]
+  t.test_files = FileList["test/**/*_test.rb"].exclude("test/benchmarks/**/*_test.rb")
 end
 
 spec = Bundler.load_gemspec("candle.gemspec")
@@ -36,7 +36,6 @@ end
 
 desc "Run benchmark tests"
 Rake::TestTask.new("test:benchmark") do |t|
-  ENV['CANDLE_RUN_BENCHMARKS'] = 'true'
   t.deps << :compile
   t.libs << "test"
   t.test_files = FileList["test/benchmarks/**/*_test.rb"]
@@ -59,7 +58,6 @@ end
 
 desc "Run benchmarks with device tests"
 task "test:device:benchmark" => :compile do
-  ENV['CANDLE_RUN_BENCHMARKS'] = 'true'
   ENV['CANDLE_TEST_VERBOSE'] = 'true'
   Rake::Task["test:device"].invoke
   Rake::Task["test:benchmark"].invoke

data/ext/candle/src/llm/gemma.rs

@@ -0,0 +1,277 @@
+use candle_core::{DType, Device, Result as CandleResult, Tensor};
+use candle_nn::VarBuilder;
+use candle_transformers::models::gemma::{Config, Model as GemmaModel};
+use hf_hub::{api::tokio::Api, Repo};
+use tokenizers::Tokenizer;
+
+use super::{GenerationConfig, TextGeneration, TextGenerator, TokenizerWrapper};
+
+#[derive(Debug)]
+pub struct Gemma {
+    model: GemmaModel,
+    tokenizer: TokenizerWrapper,
+    device: Device,
+    model_id: String,
+    eos_token_id: u32,
+}
+
+impl Gemma {
+    /// Clear the KV cache between generations
+    pub fn clear_kv_cache(&mut self) {
+        self.model.clear_kv_cache();
+    }
+    
+    /// Load a Gemma model from HuggingFace Hub
+    pub async fn from_pretrained(model_id: &str, device: Device) -> CandleResult<Self> {
+        let api = Api::new()
+            .map_err(|e| candle_core::Error::Msg(format!("Failed to create HF API: {}", e)))?;
+        
+        let repo = api.repo(Repo::model(model_id.to_string()));
+        
+        // Download model files
+        let config_filename = repo
+            .get("config.json")
+            .await
+            .map_err(|e| candle_core::Error::Msg(format!("Failed to download config: {}", e)))?;
+        
+        let tokenizer_filename = repo
+            .get("tokenizer.json")
+            .await
+            .map_err(|e| candle_core::Error::Msg(format!("Failed to download tokenizer: {}", e)))?;
+        
+        // Try different file patterns for model weights
+        let weights_filenames = if let Ok(single_file) = repo.get("model.safetensors").await {
+            vec![single_file]
+        } else {
+            // Try to find sharded model files
+            let mut sharded_files = Vec::new();
+            let mut index = 1;
+            loop {
+                // Try common shard counts for Gemma models
+                let mut found = false;
+                for total in [2, 3, 4, 5, 6, 7, 8, 10, 12] {
+                    let filename = format!("model-{:05}-of-{:05}.safetensors", index, total);
+                    if let Ok(file) = repo.get(&filename).await {
+                        sharded_files.push(file);
+                        found = true;
+                        break;
+                    }
+                }
+                if !found {
+                    break;
+                }
+                index += 1;
+            }
+            
+            if sharded_files.is_empty() {
+                return Err(candle_core::Error::Msg(
+                    "Could not find model weights. Tried: model.safetensors, model-*-of-*.safetensors".to_string()
+                ));
+            }
+            sharded_files
+        };
+        
+        // Load config
+        let config: Config = serde_json::from_reader(std::fs::File::open(config_filename)?)
+            .map_err(|e| candle_core::Error::Msg(format!("Failed to parse config: {}", e)))?;
+        
+        // Load tokenizer
+        let tokenizer = Tokenizer::from_file(tokenizer_filename)
+            .map_err(|e| candle_core::Error::Msg(format!("Failed to load tokenizer: {}", e)))?;
+        
+        // Gemma uses specific tokens
+        let eos_token_id = {
+            let vocab = tokenizer.get_vocab(true);
+            vocab.get("<eos>")
+                .or_else(|| vocab.get("<end_of_turn>"))
+                .copied()
+                .unwrap_or(1) // Default Gemma EOS
+        };
+        
+        // Load model weights
+        let vb = unsafe {
+            VarBuilder::from_mmaped_safetensors(&weights_filenames, DType::F32, &device)?
+        };
+        
+        let model = GemmaModel::new(false, &config, vb)?; // Don't use flash attention for now
+        
+        Ok(Self {
+            model,
+            tokenizer: TokenizerWrapper::new(tokenizer),
+            device,
+            model_id: model_id.to_string(),
+            eos_token_id,
+        })
+    }
+
+    /// Create from existing components (useful for testing)
+    pub fn new(
+        model: GemmaModel,
+        tokenizer: Tokenizer,
+        device: Device,
+        model_id: String,
+    ) -> Self {
+        let eos_token_id = {
+            let vocab = tokenizer.get_vocab(true);
+            vocab.get("<eos>")
+                .or_else(|| vocab.get("<end_of_turn>"))
+                .copied()
+                .unwrap_or(1)
+        };
+        
+        Self {
+            model,
+            tokenizer: TokenizerWrapper::new(tokenizer),
+            device,
+            model_id,
+            eos_token_id,
+        }
+    }
+
+    fn generate_tokens(
+        &mut self,
+        prompt_tokens: Vec<u32>,
+        config: &GenerationConfig,
+        mut callback: Option<impl FnMut(&str)>,
+    ) -> CandleResult<Vec<u32>> {
+        let mut text_gen = TextGeneration::from_config(config);
+        text_gen.set_eos_token_id(self.eos_token_id);
+        text_gen.set_tokens(prompt_tokens.clone());
+        
+        let mut all_tokens = prompt_tokens.clone();
+        let start_gen = all_tokens.len();
+        
+        for index in 0..config.max_length {
+            let context_size = if index > 0 { 1 } else { all_tokens.len() };
+            let start_pos = all_tokens.len().saturating_sub(context_size);
+            let ctxt = &all_tokens[start_pos..];
+            
+            let input = Tensor::new(ctxt, &self.device)?.unsqueeze(0)?;
+            let input = input.contiguous()?;
+            let logits = self.model.forward(&input, start_pos)?;
+            
+            let logits = logits.squeeze(0)?;
+            let logits = if logits.dims().len() == 2 {
+                let seq_len = logits.dim(0)?;
+                logits.narrow(0, seq_len - 1, 1)?.squeeze(0)?
+            } else {
+                logits
+            };
+            
+            let logits = logits.to_dtype(DType::F32)?;
+            
+            let next_token = text_gen.sample_next_token(
+                &logits,
+                Some((config.repetition_penalty, config.repetition_penalty_last_n)),
+            )?;
+            
+            all_tokens.push(next_token);
+            
+            // Stream callback
+            if let Some(ref mut cb) = callback {
+                if config.debug_tokens {
+                    // In debug mode, only show debug tokens
+                    let token_piece = self.tokenizer.token_to_piece(next_token)?;
+                    cb(&format!("[{}:{}]", next_token, token_piece));
+                } else {
+                    // Normal mode: use incremental decoding for proper text
+                    let decoded_text = self.tokenizer.decode_incremental(&all_tokens, all_tokens.len() - 1)?;
+                    cb(&decoded_text);
+                }
+            }
+            
+            // Check stop conditions
+            if text_gen.should_stop(next_token, config.max_length) {
+                break;
+            }
+            
+            // Check stop sequences
+            let generated_text = self.tokenizer.decode(&all_tokens[start_gen..], true)?;
+            if text_gen.check_stop_sequences(&generated_text, &config.stop_sequences) {
+                break;
+            }
+        }
+        
+        Ok(if config.include_prompt {
+            all_tokens
+        } else {
+            all_tokens[start_gen..].to_vec()
+        })
+    }
+    
+    /// Apply Gemma chat template
+    pub fn apply_chat_template(&self, messages: &[serde_json::Value]) -> CandleResult<String> {
+        let mut prompt = String::new();
+        
+        // Gemma uses a specific format:
+        // <start_of_turn>user\n{user_message}<end_of_turn>
+        // <start_of_turn>model\n{model_message}<end_of_turn>
+        
+        for message in messages {
+            let role = message["role"].as_str().unwrap_or("");
+            let content = message["content"].as_str().unwrap_or("");
+            
+            match role {
+                "system" => {
+                    // Gemma doesn't have explicit system messages, prepend to first user message
+                    prompt.push_str(&format!("<start_of_turn>user\nSystem: {}\n", content));
+                }
+                "user" => {
+                    if !prompt.contains("<start_of_turn>user") || prompt.ends_with("<end_of_turn>\n") {
+                        prompt.push_str("<start_of_turn>user\n");
+                    }
+                    prompt.push_str(&format!("{}<end_of_turn>\n", content));
+                }
+                "assistant" | "model" => {
+                    prompt.push_str(&format!("<start_of_turn>model\n{}<end_of_turn>\n", content));
+                }
+                _ => {}
+            }
+        }
+        
+        // Add the model prompt
+        prompt.push_str("<start_of_turn>model\n");
+        
+        Ok(prompt)
+    }
+}
+
+impl TextGenerator for Gemma {
+    fn generate(
+        &mut self,
+        prompt: &str,
+        config: &GenerationConfig,
+    ) -> CandleResult<String> {
+        let prompt_tokens = self.tokenizer.encode(prompt, true)?;
+        let output_tokens = self.generate_tokens(prompt_tokens, config, None::<fn(&str)>)?;
+        
+        if config.debug_tokens {
+            self.tokenizer.format_tokens_with_debug(&output_tokens)
+        } else {
+            self.tokenizer.decode(&output_tokens, true)
+        }
+    }
+
+    fn generate_stream(
+        &mut self,
+        prompt: &str,
+        config: &GenerationConfig,
+        mut callback: impl FnMut(&str),
+    ) -> CandleResult<String> {
+        let prompt_tokens = self.tokenizer.encode(prompt, true)?;
+        let output_tokens = self.generate_tokens(prompt_tokens, config, Some(&mut callback))?;
+        self.tokenizer.decode(&output_tokens, true)
+    }
+
+    fn model_name(&self) -> &str {
+        &self.model_id
+    }
+
+    fn device(&self) -> &Device {
+        &self.device
+    }
+    
+    fn clear_cache(&mut self) {
+        self.clear_kv_cache();
+    }
+}

data/ext/candle/src/llm/generation_config.rs

@@ -21,6 +21,8 @@ pub struct GenerationConfig {
     pub stop_sequences: Vec<String>,
     /// Whether to return the prompt in the output
     pub include_prompt: bool,
+    /// Whether to show raw tokens during generation (for debugging)
+    pub debug_tokens: bool,
 }
 
 /// Generate a random seed based on current time
@@ -43,6 +45,7 @@ impl Default for GenerationConfig {
             seed: random_seed(),
             stop_sequences: vec![],
             include_prompt: false,
+            debug_tokens: false,
         }
     }
 }

data/ext/candle/src/llm/llama.rs

@@ -205,77 +205,14 @@ impl Llama {
             
             // Stream callback
             if let Some(ref mut cb) = callback {
-                let token_text = self.tokenizer.token_to_piece(next_token)?;
-                cb(&token_text);
-            }
-            
-            // Check stop conditions
-            if text_gen.should_stop(next_token, config.max_length) {
-                break;
-            }
-            
-            // Check stop sequences
-            let generated_text = self.tokenizer.decode(&all_tokens[start_gen..], true)?;
-            if text_gen.check_stop_sequences(&generated_text, &config.stop_sequences) {
-                break;
-            }
-        }
-        
-        Ok(if config.include_prompt {
-            all_tokens
-        } else {
-            all_tokens[start_gen..].to_vec()
-        })
-    }
-    
-    fn generate_tokens_decoded(
-        &mut self,
-        prompt_tokens: Vec<u32>,
-        config: &GenerationConfig,
-        mut callback: Option<impl FnMut(&str)>,
-    ) -> CandleResult<Vec<u32>> {
-        let mut text_gen = TextGeneration::from_config(config);
-        text_gen.set_eos_token_id(self.eos_token_id);
-        text_gen.set_tokens(prompt_tokens.clone());
-        
-        let mut all_tokens = prompt_tokens.clone();
-        let start_gen = all_tokens.len();
-        let mut previously_decoded = String::new();
-        
-        for index in 0..config.max_length {
-            let context_size = if index > 0 { 1 } else { all_tokens.len() };
-            let start_pos = all_tokens.len().saturating_sub(context_size);
-            let ctxt = &all_tokens[start_pos..];
-            
-            let input = Tensor::new(ctxt, &self.device)?.unsqueeze(0)?;
-            let input = input.contiguous()?;
-            let logits = self.model.forward(&input, start_pos, &mut self.cache)?;
-            
-            let logits = logits.squeeze(0)?;
-            let logits = if logits.dims().len() == 2 {
-                let seq_len = logits.dim(0)?;
-                logits.narrow(0, seq_len - 1, 1)?.squeeze(0)?
-            } else {
-                logits
-            };
-            
-            let logits = logits.to_dtype(DType::F32)?;
-            
-            let next_token = text_gen.sample_next_token(
-                &logits,
-                Some((config.repetition_penalty, config.repetition_penalty_last_n)),
-            )?;
-            
-            all_tokens.push(next_token);
-            
-            // Stream callback with incremental decoding
-            if let Some(ref mut cb) = callback {
-                let current_decoded = self.tokenizer.decode(&all_tokens[start_gen..], true)?;
-                
-                if current_decoded.len() > previously_decoded.len() {
-                    let new_text = &current_decoded[previously_decoded.len()..];
-                    cb(new_text);
-                    previously_decoded = current_decoded;
+                if config.debug_tokens {
+                    // In debug mode, only show debug tokens
+                    let token_piece = self.tokenizer.token_to_piece(next_token)?;
+                    cb(&format!("[{}:{}]", next_token, token_piece));
+                } else {
+                    // Normal mode: use incremental decoding for proper text
+                    let decoded_text = self.tokenizer.decode_incremental(&all_tokens, all_tokens.len() - 1)?;
+                    cb(&decoded_text);
                 }
             }
             
@@ -285,12 +222,7 @@ impl Llama {
             }
             
             // Check stop sequences
-            let generated_text = if callback.is_some() {
-                previously_decoded.clone()
-            } else {
-                self.tokenizer.decode(&all_tokens[start_gen..], true)?
-            };
-            
+            let generated_text = self.tokenizer.decode(&all_tokens[start_gen..], true)?;
             if text_gen.check_stop_sequences(&generated_text, &config.stop_sequences) {
                 break;
             }
@@ -374,7 +306,12 @@ impl TextGenerator for Llama {
     ) -> CandleResult<String> {
         let prompt_tokens = self.tokenizer.encode(prompt, true)?;
         let output_tokens = self.generate_tokens(prompt_tokens, config, None::<fn(&str)>)?;
-        self.tokenizer.decode(&output_tokens, true)
+        
+        if config.debug_tokens {
+            self.tokenizer.format_tokens_with_debug(&output_tokens)
+        } else {
+            self.tokenizer.decode(&output_tokens, true)
+        }
     }
 
     fn generate_stream(
@@ -384,7 +321,7 @@ impl TextGenerator for Llama {
         mut callback: impl FnMut(&str),
     ) -> CandleResult<String> {
         let prompt_tokens = self.tokenizer.encode(prompt, true)?;
-        let output_tokens = self.generate_tokens_decoded(prompt_tokens, config, Some(&mut callback))?;
+        let output_tokens = self.generate_tokens(prompt_tokens, config, Some(&mut callback))?;
         self.tokenizer.decode(&output_tokens, true)
     }