red-candle 1.0.0.pre.1 → 1.0.0.pre.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 34ca4771af8508ace5ee8df5d5111a407c79fea307f9fbc2b72d9e46d6dd2099
-  data.tar.gz: 39a316c2990f8766257e07c5ae8803ba852aebc5c93b14349548b1da82b6ae4a
+  metadata.gz: 9cdd24afacff8070c3011dad25e975cb4bff9420d7a760ec105d840480bd0a4f
+  data.tar.gz: df1d8df277e5c8447c0a6598c2f7ed1389e74b356b42b3bb8f51b2f4c7863884
 SHA512:
-  metadata.gz: 51fcc58936f25818485bacd5dfc86911ca1ecbc0a8b87f6fdcf2d015a91faccd6aeb889df33a931eae6d864984254ab4a81ab7050fb00caf14d090820cf4d86b
-  data.tar.gz: f6e58cc4e723c12e3d290fcb45b1cab5923594bfd011be223ec327c65c7c91cb7d37f9330ed82eadd4b7eb3085ef0ba9b874a6b98c15221b21c60cd01de11418
+  metadata.gz: b1a4c0df852b01a3cdff7f78dbe0a6553cd39ff51027b35cca4f8ff6c63d4d89b411493d644fe4d4da29d516bbfa48c1b58a5a3b317556041848004909b63edc
+  data.tar.gz: facd28b47751a365e5ef1fdf900f55100392eeee4d5528c3f46e340ec8d69c913ea5698aa4c2900c926947b7e444bc9f876cc79644c81fe731d71f0970715ff3

data/ext/candle/build.rs

@@ -0,0 +1,116 @@
+use std::env;
+use std::path::Path;
+
+fn main() {
+    // Register our custom cfg flags with rustc
+    println!("cargo::rustc-check-cfg=cfg(force_cpu)");
+    println!("cargo::rustc-check-cfg=cfg(has_cuda)");
+    println!("cargo::rustc-check-cfg=cfg(has_metal)");
+    println!("cargo::rustc-check-cfg=cfg(has_mkl)");
+    println!("cargo::rustc-check-cfg=cfg(has_accelerate)");
+    
+    println!("cargo:rerun-if-changed=build.rs");
+    println!("cargo:rerun-if-env-changed=CANDLE_FORCE_CPU");
+    println!("cargo:rerun-if-env-changed=CANDLE_CUDA_PATH");
+    println!("cargo:rerun-if-env-changed=CUDA_ROOT");
+    println!("cargo:rerun-if-env-changed=CUDA_PATH");
+    println!("cargo:rerun-if-env-changed=CANDLE_FEATURES");
+    println!("cargo:rerun-if-env-changed=CANDLE_ENABLE_CUDA");
+
+    // Check if we should force CPU only
+    if env::var("CANDLE_FORCE_CPU").is_ok() {
+        println!("cargo:rustc-cfg=force_cpu");
+        println!("cargo:warning=CANDLE_FORCE_CPU is set, disabling all acceleration");
+        return;
+    }
+
+    // Detect CUDA availability
+    let cuda_available = detect_cuda();
+    let cuda_enabled = env::var("CANDLE_ENABLE_CUDA").is_ok();
+    
+    if cuda_available && cuda_enabled {
+        println!("cargo:rustc-cfg=has_cuda");
+        println!("cargo:warning=CUDA detected and enabled via CANDLE_ENABLE_CUDA");
+    } else if cuda_available && !cuda_enabled {
+        println!("cargo:warning=CUDA detected but not enabled. To enable CUDA support (coming soon), set CANDLE_ENABLE_CUDA=1");
+    }
+
+    // Detect Metal availability (macOS only)
+    #[cfg(target_os = "macos")]
+    {
+        println!("cargo:rustc-cfg=has_metal");
+        println!("cargo:warning=Metal detected (macOS), Metal acceleration will be available");
+    }
+
+    // Detect MKL availability
+    if detect_mkl() {
+        println!("cargo:rustc-cfg=has_mkl");
+        println!("cargo:warning=Intel MKL detected, MKL acceleration will be available");
+    }
+
+    // Detect Accelerate framework (macOS)
+    #[cfg(target_os = "macos")]
+    {
+        println!("cargo:rustc-cfg=has_accelerate");
+        println!("cargo:warning=Accelerate framework detected (macOS)");
+    }
+}
+
+fn detect_cuda() -> bool {
+    // Check environment variables first
+    if env::var("CANDLE_CUDA_PATH").is_ok() {
+        return true;
+    }
+
+    if env::var("CUDA_ROOT").is_ok() || env::var("CUDA_PATH").is_ok() {
+        return true;
+    }
+
+    // Check common CUDA installation paths
+    let cuda_paths = [
+        "/usr/local/cuda",
+        "/opt/cuda",
+        "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA",
+        "C:\\CUDA",
+    ];
+
+    for path in &cuda_paths {
+        if Path::new(path).exists() {
+            return true;
+        }
+    }
+
+    // Check if nvcc is in PATH
+    if let Ok(path_var) = env::var("PATH") {
+        for path in env::split_paths(&path_var) {
+            if path.join("nvcc").exists() || path.join("nvcc.exe").exists() {
+                return true;
+            }
+        }
+    }
+
+    false
+}
+
+fn detect_mkl() -> bool {
+    // Check environment variables
+    if env::var("MKLROOT").is_ok() || env::var("MKL_ROOT").is_ok() {
+        return true;
+    }
+
+    // Check common MKL installation paths
+    let mkl_paths = [
+        "/opt/intel/mkl",
+        "/opt/intel/oneapi/mkl/latest",
+        "C:\\Program Files (x86)\\Intel\\oneAPI\\mkl\\latest",
+        "C:\\Program Files\\Intel\\oneAPI\\mkl\\latest",
+    ];
+
+    for path in &mkl_paths {
+        if Path::new(path).exists() {
+            return true;
+        }
+    }
+
+    false
+}

data/ext/candle/src/lib.rs

@@ -1,7 +1,7 @@
-use magnus::{function, method, prelude::*, Ruby};
+use magnus::{function, prelude::*, Ruby};
 
 use crate::ruby::candle_utils;
-use crate::ruby::{DType, Device, QTensor, Result as RbResult, Tensor};
+use crate::ruby::Result as RbResult;
 
 pub mod llm;
 pub mod reranker;
@@ -42,101 +42,11 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     ruby::init_embedding_model(rb_candle)?;
     ruby::init_llm(rb_candle)?;
     reranker::init(rb_candle)?;
+    ruby::dtype::init(rb_candle)?;
+    ruby::qtensor::init(rb_candle)?;
+    ruby::device::init(rb_candle)?;
+    ruby::tensor::init(rb_candle)?;
     candle_utils(rb_candle)?;
-    let rb_tensor = rb_candle.define_class("Tensor", Ruby::class_object(ruby))?;
-    rb_tensor.define_singleton_method("new", function!(Tensor::new, 3))?;
-    // rb_tensor.define_singleton_method("cat", function!(Tensor::cat, 2))?;
-    // rb_tensor.define_singleton_method("stack", function!(Tensor::stack, 2))?;
-    rb_tensor.define_singleton_method("rand", function!(Tensor::rand, 2))?;
-    rb_tensor.define_singleton_method("randn", function!(Tensor::randn, 2))?;
-    rb_tensor.define_singleton_method("ones", function!(Tensor::ones, 2))?;
-    rb_tensor.define_singleton_method("zeros", function!(Tensor::zeros, 2))?;
-    rb_tensor.define_method("values", method!(Tensor::values, 0))?;
-    rb_tensor.define_method("values_f32", method!(Tensor::values_f32, 0))?;
-    rb_tensor.define_method("item", method!(Tensor::item, 0))?;
-    rb_tensor.define_method("shape", method!(Tensor::shape, 0))?;
-    rb_tensor.define_method("stride", method!(Tensor::stride, 0))?;
-    rb_tensor.define_method("dtype", method!(Tensor::dtype, 0))?;
-    rb_tensor.define_method("device", method!(Tensor::device, 0))?;
-    rb_tensor.define_method("rank", method!(Tensor::rank, 0))?;
-    rb_tensor.define_method("elem_count", method!(Tensor::elem_count, 0))?;
-    rb_tensor.define_method("sin", method!(Tensor::sin, 0))?;
-    rb_tensor.define_method("cos", method!(Tensor::cos, 0))?;
-    rb_tensor.define_method("log", method!(Tensor::log, 0))?;
-    rb_tensor.define_method("sqr", method!(Tensor::sqr, 0))?;
-    rb_tensor.define_method("mean", method!(Tensor::mean, 1))?;
-    rb_tensor.define_method("sum", method!(Tensor::sum, 1))?;
-    rb_tensor.define_method("sqrt", method!(Tensor::sqrt, 0))?;
-    rb_tensor.define_method("/", method!(Tensor::__truediv__, 1))?; // Accepts Tensor, Float, or Integer
-    rb_tensor.define_method("recip", method!(Tensor::recip, 0))?;
-    rb_tensor.define_method("exp", method!(Tensor::exp, 0))?;
-    rb_tensor.define_method("powf", method!(Tensor::powf, 1))?;
-    rb_tensor.define_method("index_select", method!(Tensor::index_select, 2))?;
-    rb_tensor.define_method("matmul", method!(Tensor::matmul, 1))?;
-    rb_tensor.define_method("broadcast_add", method!(Tensor::broadcast_add, 1))?;
-    rb_tensor.define_method("broadcast_sub", method!(Tensor::broadcast_sub, 1))?;
-    rb_tensor.define_method("broadcast_mul", method!(Tensor::broadcast_mul, 1))?;
-    rb_tensor.define_method("broadcast_div", method!(Tensor::broadcast_div, 1))?;
-    rb_tensor.define_method("where_cond", method!(Tensor::where_cond, 2))?;
-    rb_tensor.define_method("+", method!(Tensor::__add__, 1))?;
-    rb_tensor.define_method("*", method!(Tensor::__mul__, 1))?;
-    rb_tensor.define_method("-", method!(Tensor::__sub__, 1))?;
-    rb_tensor.define_method("reshape", method!(Tensor::reshape, 1))?;
-    rb_tensor.define_method("broadcast_as", method!(Tensor::broadcast_as, 1))?;
-    rb_tensor.define_method("broadcast_left", method!(Tensor::broadcast_left, 1))?;
-    rb_tensor.define_method("squeeze", method!(Tensor::squeeze, 1))?;
-    rb_tensor.define_method("unsqueeze", method!(Tensor::unsqueeze, 1))?;
-    rb_tensor.define_method("get", method!(Tensor::get, 1))?;
-    rb_tensor.define_method("[]", method!(Tensor::get, 1))?;
-    rb_tensor.define_method("transpose", method!(Tensor::transpose, 2))?;
-    rb_tensor.define_method("narrow", method!(Tensor::narrow, 3))?;
-    rb_tensor.define_method("argmax_keepdim", method!(Tensor::argmax_keepdim, 1))?;
-    rb_tensor.define_method("argmin_keepdim", method!(Tensor::argmin_keepdim, 1))?;
-    rb_tensor.define_method("max_keepdim", method!(Tensor::max_keepdim, 1))?;
-    rb_tensor.define_method("min_keepdim", method!(Tensor::min_keepdim, 1))?;
-    // rb_tensor.define_method("eq", method!(Tensor::eq, 1))?;
-    // rb_tensor.define_method("ne", method!(Tensor::ne, 1))?;
-    // rb_tensor.define_method("lt", method!(Tensor::lt, 1))?;
-    // rb_tensor.define_method("gt", method!(Tensor::gt, 1))?;
-    // rb_tensor.define_method("ge", method!(Tensor::ge, 1))?;
-    // rb_tensor.define_method("le", method!(Tensor::le, 1))?;
-    rb_tensor.define_method("sum_all", method!(Tensor::sum_all, 0))?;
-    rb_tensor.define_method("mean_all", method!(Tensor::mean_all, 0))?;
-    rb_tensor.define_method("flatten_from", method!(Tensor::flatten_from, 1))?;
-    rb_tensor.define_method("flatten_to", method!(Tensor::flatten_to, 1))?;
-    rb_tensor.define_method("flatten_all", method!(Tensor::flatten_all, 0))?;
-    rb_tensor.define_method("t", method!(Tensor::t, 0))?;
-    rb_tensor.define_method("contiguous", method!(Tensor::contiguous, 0))?;
-    rb_tensor.define_method("is_contiguous", method!(Tensor::is_contiguous, 0))?;
-    rb_tensor.define_method(
-        "is_fortran_contiguous",
-        method!(Tensor::is_fortran_contiguous, 0),
-    )?;
-    rb_tensor.define_method("detach", method!(Tensor::detach, 0))?;
-    rb_tensor.define_method("copy", method!(Tensor::copy, 0))?;
-    rb_tensor.define_method("to_dtype", method!(Tensor::to_dtype, 1))?;
-    rb_tensor.define_method("to_device", method!(Tensor::to_device, 1))?;
-    rb_tensor.define_method("to_s", method!(Tensor::__str__, 0))?;
-    rb_tensor.define_method("inspect", method!(Tensor::__repr__, 0))?;
-
-    let rb_dtype = rb_candle.define_class("DType", Ruby::class_object(ruby))?;
-    rb_dtype.define_method("to_s", method!(DType::__str__, 0))?;
-    rb_dtype.define_method("inspect", method!(DType::__repr__, 0))?;
-
-    let rb_device = rb_candle.define_class("Device", Ruby::class_object(ruby))?;
-    rb_device.define_singleton_method("cpu", function!(Device::cpu, 0))?;
-    rb_device.define_singleton_method("cuda", function!(Device::cuda, 0))?;
-    rb_device.define_singleton_method("metal", function!(Device::metal, 0))?;
-    rb_device.define_singleton_method("available_devices", function!(ruby::device::available_devices, 0))?;
-    rb_device.define_singleton_method("default", function!(ruby::device::default_device, 0))?;
-    rb_device.define_method("to_s", method!(Device::__str__, 0))?;
-    rb_device.define_method("inspect", method!(Device::__repr__, 0))?;
-
-    let rb_qtensor = rb_candle.define_class("QTensor", Ruby::class_object(ruby))?;
-    rb_qtensor.define_method("ggml_dtype", method!(QTensor::ggml_dtype, 0))?;
-    rb_qtensor.define_method("rank", method!(QTensor::rank, 0))?;
-    rb_qtensor.define_method("shape", method!(QTensor::shape, 0))?;
-    rb_qtensor.define_method("dequantize", method!(QTensor::dequantize, 0))?;
 
     Ok(())
 }

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

@@ -0,0 +1,49 @@
+use std::time::{SystemTime, UNIX_EPOCH};
+
+/// Configuration for text generation
+#[derive(Debug, Clone)]
+pub struct GenerationConfig {
+    /// The maximum number of tokens to generate
+    pub max_length: usize,
+    /// The temperature for sampling
+    pub temperature: f64,
+    /// The top-p value for nucleus sampling
+    pub top_p: Option<f64>,
+    /// The top-k value for top-k sampling
+    pub top_k: Option<usize>,
+    /// The repetition penalty
+    pub repetition_penalty: f32,
+    /// The repetition penalty range
+    pub repetition_penalty_last_n: usize,
+    /// Random seed for sampling
+    pub seed: u64,
+    /// Stop sequences
+    pub stop_sequences: Vec<String>,
+    /// Whether to return the prompt in the output
+    pub include_prompt: bool,
+}
+
+/// Generate a random seed based on current time
+fn random_seed() -> u64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|d| d.as_nanos() as u64)
+        .unwrap_or(42)
+}
+
+impl Default for GenerationConfig {
+    fn default() -> Self {
+        Self {
+            max_length: 512,
+            temperature: 0.7,
+            top_p: None,
+            top_k: None,
+            repetition_penalty: 1.1,
+            repetition_penalty_last_n: 64,
+            seed: random_seed(),
+            stop_sequences: vec![],
+            include_prompt: false,
+        }
+    }
+}
+

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

@@ -0,0 +1,325 @@
+use candle_core::{DType, Device, Result as CandleResult, Tensor};
+use candle_nn::VarBuilder;
+use candle_transformers::models::mistral::{Config, Model as MistralModel};
+use hf_hub::{api::tokio::Api, Repo};
+use tokenizers::Tokenizer;
+
+use super::{GenerationConfig, TextGeneration, TextGenerator, TokenizerWrapper};
+
+#[derive(Debug)]
+pub struct Mistral {
+    model: MistralModel,
+    tokenizer: TokenizerWrapper,
+    device: Device,
+    model_id: String,
+    eos_token_id: u32,
+}
+
+impl Mistral {
+    /// Clear the KV cache between generations
+    pub fn clear_kv_cache(&mut self) {
+        self.model.clear_kv_cache();
+    }
+    
+    /// Load a Mistral 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 if let Ok(consolidated_file) = repo.get("consolidated.safetensors").await {
+            // Some Mistral models use consolidated.safetensors
+            vec![consolidated_file]
+        } else {
+            // Try to find sharded model files
+            let mut sharded_files = Vec::new();
+            let mut index = 1;
+            loop {
+                // Try common shard counts
+                let mut found = false;
+                for total in [2, 3, 4, 5, 6, 7, 8] {
+                    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() {
+                // Try single pytorch_model.bin as last resort (though we prefer safetensors)
+                if let Ok(_pytorch_file) = repo.get("pytorch_model.bin").await {
+                    return Err(candle_core::Error::Msg(
+                        "Only safetensors format is supported. This model uses pytorch_model.bin format.".to_string()
+                    ));
+                } else {
+                    return Err(candle_core::Error::Msg(
+                        "Could not find model weights. Tried: model.safetensors, consolidated.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)))?;
+        
+        let eos_token_id = tokenizer
+            .get_vocab(true)
+            .get("</s>")
+            .copied()
+            .unwrap_or(2);
+        
+        // Load model weights
+        let vb = unsafe {
+            VarBuilder::from_mmaped_safetensors(&weights_filenames, DType::F32, &device)?
+        };
+        
+        let model = MistralModel::new(&config, vb)?;
+        
+        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: MistralModel,
+        tokenizer: Tokenizer,
+        device: Device,
+        model_id: String,
+    ) -> Self {
+        let eos_token_id = tokenizer
+            .get_vocab(true)
+            .get("</s>")
+            .copied()
+            .unwrap_or(2);
+        
+        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)?;
+            // 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
+            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;
+                }
+            }
+            
+            // Check stop conditions
+            if text_gen.should_stop(next_token, config.max_length) {
+                break;
+            }
+            
+            // Check stop sequences
+            let generated_text = if callback.is_some() {
+                previously_decoded.clone()
+            } else {
+                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()
+        })
+    }
+}
+
+impl TextGenerator for Mistral {
+    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)>)?;
+        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_decoded(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/mod.rs

@@ -0,0 +1,68 @@
+use candle_core::{Device, Result as CandleResult};
+use tokenizers::Tokenizer;
+
+pub mod mistral;
+pub mod generation_config;
+pub mod text_generation;
+
+pub use generation_config::GenerationConfig;
+pub use text_generation::TextGeneration;
+
+/// Trait for text generation models
+pub trait TextGenerator: Send + Sync {
+    /// Generate text from a prompt
+    fn generate(
+        &mut self,
+        prompt: &str,
+        config: &GenerationConfig,
+    ) -> CandleResult<String>;
+
+    /// Generate text with streaming callback
+    fn generate_stream(
+        &mut self,
+        prompt: &str,
+        config: &GenerationConfig,
+        callback: impl FnMut(&str),
+    ) -> CandleResult<String>;
+
+    /// Get the model's name
+    fn model_name(&self) -> &str;
+
+    /// Get the device the model is running on
+    fn device(&self) -> &Device;
+    
+    /// 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)))
+    }
+}