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

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

@@ -180,87 +180,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 +197,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 +219,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 +234,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

@@ -3,11 +3,14 @@ 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;
 
 /// Trait for text generation models
 pub trait TextGenerator: Send + Sync {
@@ -66,4 +69,60 @@ impl TokenizerWrapper {
             .map(|s| s.to_string())
             .ok_or_else(|| candle_core::Error::Msg(format!("Unknown token id: {}", token)))
     }
+    
+    /// Decode a single token for streaming output
+    pub fn decode_token(&self, token: u32) -> CandleResult<String> {
+        // Decode the single token properly
+        self.decode(&[token], true)
+    }
+    
+    /// Decode tokens incrementally for streaming
+    /// This is more efficient than decoding single tokens
+    pub fn decode_incremental(&self, all_tokens: &[u32], new_tokens_start: usize) -> CandleResult<String> {
+        if new_tokens_start >= all_tokens.len() {
+            return Ok(String::new());
+        }
+        
+        // Decode all tokens up to this point
+        let full_text = self.decode(all_tokens, true)?;
+        
+        // If we're at the start, return everything
+        if new_tokens_start == 0 {
+            return Ok(full_text);
+        }
+        
+        // Otherwise, decode up to the previous token and return the difference
+        let previous_text = self.decode(&all_tokens[..new_tokens_start], true)?;
+        
+        // Find the common prefix between the two strings to handle cases where
+        // the tokenizer might produce slightly different text when decoding
+        // different token sequences
+        let common_prefix_len = full_text
+            .char_indices()
+            .zip(previous_text.chars())
+            .take_while(|((_, c1), c2)| c1 == c2)
+            .count();
+        
+        // Find the byte position of the character boundary
+        let byte_pos = full_text
+            .char_indices()
+            .nth(common_prefix_len)
+            .map(|(pos, _)| pos)
+            .unwrap_or(full_text.len());
+        
+        // Return only the new portion
+        Ok(full_text[byte_pos..].to_string())
+    }
+    
+    /// Format tokens with debug information
+    pub fn format_tokens_with_debug(&self, tokens: &[u32]) -> CandleResult<String> {
+        let mut result = String::new();
+        
+        for &token in tokens {
+            let token_piece = self.token_to_piece(token)?;
+            result.push_str(&format!("[{}:{}]", token, token_piece));
+        }
+        
+        Ok(result)
+    }
 }