red-candle 1.0.0.pre.6 → 1.0.0

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

@@ -21,6 +21,11 @@ impl Gemma {
         self.model.clear_kv_cache();
     }
     
+    /// Get the tokenizer
+    pub fn tokenizer(&self) -> &TokenizerWrapper {
+        &self.tokenizer
+    }
+    
     /// Load a Gemma model from HuggingFace Hub
     pub async fn from_pretrained(model_id: &str, device: Device) -> CandleResult<Self> {
         let api = Api::new()
@@ -169,77 +174,14 @@ impl Gemma {
             
             // 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)?;
-            
-            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);
                 }
             }
             
@@ -249,12 +191,7 @@ impl Gemma {
             }
             
             // 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;
             }
@@ -312,7 +249,12 @@ impl TextGenerator for Gemma {
     ) -> 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(
@@ -322,7 +264,7 @@ impl TextGenerator for Gemma {
         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)
     }
 

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

@@ -28,6 +28,11 @@ impl Llama {
         }
     }
     
+    /// Get the tokenizer
+    pub fn tokenizer(&self) -> &TokenizerWrapper {
+        &self.tokenizer
+    }
+    
     /// Load a Llama model from HuggingFace Hub
     pub async fn from_pretrained(model_id: &str, device: Device) -> CandleResult<Self> {
         let api = Api::new()
@@ -205,77 +210,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 +227,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 +311,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 +326,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)
     }
 

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

@@ -21,6 +21,11 @@ impl Mistral {
         self.model.clear_kv_cache();
     }
     
+    /// Get the tokenizer
+    pub fn tokenizer(&self) -> &TokenizerWrapper {
+        &self.tokenizer
+    }
+    
     /// Load a Mistral model from HuggingFace Hub
     pub async fn from_pretrained(model_id: &str, device: Device) -> CandleResult<Self> {
         let api = Api::new()
@@ -180,87 +185,14 @@ impl Mistral {
             
             // 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();
-        
-        // For incremental decoding
-        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)?;
-            // Ensure input tensor is contiguous for Metal backend
-            let input = input.contiguous()?;
-            let logits = self.model.forward(&input, start_pos)?;
-            
-            // The model returns logits of shape [batch_size, seq_len, vocab_size]
-            // We need to get the logits for the last token only
-            let logits = logits.squeeze(0)?; // Remove batch dimension
-            let logits = if logits.dims().len() == 2 {
-                // If we still have [seq_len, vocab_size], take the last token
-                let seq_len = logits.dim(0)?;
-                logits.narrow(0, seq_len - 1, 1)?.squeeze(0)?
-            } else {
-                // Already [vocab_size]
-                logits
-            };
-            
-            // Convert to F32 for sampling if needed
-            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 {
-                // Decode all generated tokens so far
-                let current_decoded = self.tokenizer.decode(&all_tokens[start_gen..], true)?;
-                
-                // Only emit the new text since last callback
-                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);
                 }
             }
             
@@ -270,12 +202,7 @@ impl Mistral {
             }
             
             // 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;
             }
@@ -297,7 +224,12 @@ impl TextGenerator for Mistral {
     ) -> 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(
@@ -307,7 +239,7 @@ impl TextGenerator for Mistral {
         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)
     }
 

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

@@ -1,14 +1,16 @@
 use candle_core::{Device, Result as CandleResult};
-use tokenizers::Tokenizer;
 
 pub mod mistral;
 pub mod llama;
 pub mod gemma;
 pub mod generation_config;
 pub mod text_generation;
+pub mod quantized_gguf;
 
 pub use generation_config::GenerationConfig;
 pub use text_generation::TextGeneration;
+pub use quantized_gguf::QuantizedGGUF;
+pub use crate::tokenizer::TokenizerWrapper;
 
 /// Trait for text generation models
 pub trait TextGenerator: Send + Sync {
@@ -35,36 +37,4 @@ pub trait TextGenerator: Send + Sync {
     
     /// Clear any cached state (like KV cache)
     fn clear_cache(&mut self);
-}
-
-/// Common structure for managing tokenizer
-#[derive(Debug)]
-pub struct TokenizerWrapper {
-    tokenizer: Tokenizer,
-}
-
-impl TokenizerWrapper {
-    pub fn new(tokenizer: Tokenizer) -> Self {
-        Self { tokenizer }
-    }
-
-    pub fn encode(&self, text: &str, add_special_tokens: bool) -> CandleResult<Vec<u32>> {
-        let encoding = self.tokenizer
-            .encode(text, add_special_tokens)
-            .map_err(|e| candle_core::Error::Msg(format!("Tokenizer error: {}", e)))?;
-        Ok(encoding.get_ids().to_vec())
-    }
-
-    pub fn decode(&self, tokens: &[u32], skip_special_tokens: bool) -> CandleResult<String> {
-        self.tokenizer
-            .decode(tokens, skip_special_tokens)
-            .map_err(|e| candle_core::Error::Msg(format!("Tokenizer decode error: {}", e)))
-    }
-
-    pub fn token_to_piece(&self, token: u32) -> CandleResult<String> {
-        self.tokenizer
-            .id_to_token(token)
-            .map(|s| s.to_string())
-            .ok_or_else(|| candle_core::Error::Msg(format!("Unknown token id: {}", token)))
-    }
 }