red-candle 1.0.0.pre.6 → 1.0.0

data/ext/candle/src/ruby/embedding_model.rs

@@ -3,7 +3,8 @@
 use crate::ruby::{
     errors::{wrap_candle_err, wrap_hf_err, wrap_std_err},
 };
-use crate::ruby::{Tensor, Device, Result as RbResult};
+use crate::ruby::{Tensor, Device, Result};
+use crate::tokenizer::{TokenizerWrapper, loader::TokenizerLoader};
 use candle_core::{DType as CoreDType, Device as CoreDevice, Module, Tensor as CoreTensor};
 use safetensors::tensor::SafeTensors;
 use candle_nn::VarBuilder;
@@ -14,7 +15,6 @@ use candle_transformers::models::{
 };
 use magnus::{class, function, method, prelude::*, Error, RModule};
 use std::path::Path;
-use tokenizers::Tokenizer;
 use serde_json;
 
 
@@ -70,12 +70,12 @@ pub struct EmbeddingModelInner {
     model_path: Option<String>,
     embedding_model_type: Option<EmbeddingModelType>,
     model: Option<EmbeddingModelVariant>,
-    tokenizer: Option<Tokenizer>,
+    tokenizer: Option<TokenizerWrapper>,
     embedding_size: Option<usize>,
 }
 
 impl EmbeddingModel {
-    pub fn new(model_path: Option<String>, tokenizer_path: Option<String>, device: Option<Device>, embedding_model_type: Option<String>, embedding_size: Option<usize>) -> RbResult<Self> {
+    pub fn new(model_path: Option<String>, tokenizer_path: Option<String>, device: Option<Device>, embedding_model_type: Option<String>, embedding_size: Option<usize>) -> Result<Self> {
         let device = device.unwrap_or(Device::Cpu).as_device()?;
         let embedding_model_type = embedding_model_type
             .and_then(|mt| EmbeddingModelType::from_string(&mt))
@@ -102,7 +102,7 @@ impl EmbeddingModel {
     /// Generates an embedding vector for the input text using the specified pooling method.
     /// &RETURNS&: Tensor
     /// pooling_method: "pooled", "pooled_normalized", or "cls" (default: "pooled")
-    pub fn embedding(&self, input: String, pooling_method: String) -> RbResult<Tensor> {
+    pub fn embedding(&self, input: String, pooling_method: String) -> Result<Tensor> {
         match &self.0.model {
             Some(model) => {
                 match &self.0.tokenizer {
@@ -116,7 +116,7 @@ impl EmbeddingModel {
 
     /// Returns the unpooled embedding tensor ([1, SEQLENGTH, DIM]) for the input text
     /// &RETURNS&: Tensor
-    pub fn embeddings(&self, input: String) -> RbResult<Tensor> {
+    pub fn embeddings(&self, input: String) -> Result<Tensor> {
         match &self.0.model {
             Some(model) => {
                 match &self.0.tokenizer {
@@ -130,27 +130,27 @@ impl EmbeddingModel {
 
     /// Pools and normalizes a sequence embedding tensor ([1, SEQLENGTH, DIM]) to [1, DIM]
     /// &RETURNS&: Tensor
-    pub fn pool_embedding(&self, tensor: &Tensor) -> RbResult<Tensor> {
+    pub fn pool_embedding(&self, tensor: &Tensor) -> Result<Tensor> {
         let pooled = Self::pooled_embedding(&tensor.0)?;
         Ok(Tensor(pooled))
     }
 
     /// Pools and normalizes a sequence embedding tensor ([1, SEQLENGTH, DIM]) to [1, DIM]
     /// &RETURNS&: Tensor
-    pub fn pool_and_normalize_embedding(&self, tensor: &Tensor) -> RbResult<Tensor> {
+    pub fn pool_and_normalize_embedding(&self, tensor: &Tensor) -> Result<Tensor> {
         let pooled = Self::pooled_normalized_embedding(&tensor.0)?;
         Ok(Tensor(pooled))
     }
 
     /// Pools the embedding tensor by extracting the CLS token ([1, SEQLENGTH, DIM] -> [1, DIM])
     /// &RETURNS&: Tensor
-    pub fn pool_cls_embedding(&self, tensor: &Tensor) -> RbResult<Tensor> {
+    pub fn pool_cls_embedding(&self, tensor: &Tensor) -> Result<Tensor> {
         let pooled = Self::pooled_cls_embedding(&tensor.0)?;
         Ok(Tensor(pooled))
     }
 
     /// Infers and validates the embedding size from a safetensors file
-    fn resolve_embedding_size(model_path: &Path, embedding_size: Option<usize>) -> Result<usize, magnus::Error> {
+    fn resolve_embedding_size(model_path: &Path, embedding_size: Option<usize>) -> std::result::Result<usize, magnus::Error> {
         match embedding_size {
             Some(user_dim) => {
                 Ok(user_dim)
@@ -170,7 +170,7 @@ impl EmbeddingModel {
         }
     }
 
-    fn build_embedding_model(model_path: &Path, device: CoreDevice, embedding_model_type: EmbeddingModelType, embedding_size: Option<usize>) -> RbResult<EmbeddingModelVariant> {
+    fn build_embedding_model(model_path: &Path, device: CoreDevice, embedding_model_type: EmbeddingModelType, embedding_size: Option<usize>) -> Result<EmbeddingModelVariant> {
         use hf_hub::{api::sync::Api, Repo, RepoType};
         let api = Api::new().map_err(wrap_hf_err)?;
         let repo = Repo::new(model_path.to_str().unwrap().to_string(), RepoType::Model);
@@ -257,7 +257,7 @@ impl EmbeddingModel {
         }
     }
 
-    fn build_tokenizer(tokenizer_path: String) -> RbResult<Tokenizer> {
+    fn build_tokenizer(tokenizer_path: String) -> Result<TokenizerWrapper> {
         use hf_hub::{api::sync::Api, Repo, RepoType};
         let tokenizer_path = Api::new()
                 .map_err(wrap_hf_err)?
@@ -267,20 +267,16 @@ impl EmbeddingModel {
                 ))
                 .get("tokenizer.json")
                 .map_err(wrap_hf_err)?;
-        let mut tokenizer = tokenizers::Tokenizer::from_file(tokenizer_path)
+        let tokenizer = tokenizers::Tokenizer::from_file(tokenizer_path)
             .map_err(wrap_std_err)?;
-        let pp = tokenizers::PaddingParams {
-            strategy: tokenizers::PaddingStrategy::BatchLongest,
-            ..Default::default()
-        };
-        tokenizer.with_padding(Some(pp));
-
-        Ok(tokenizer)
+        
+        let tokenizer = TokenizerLoader::with_padding(tokenizer, None);
+        Ok(TokenizerWrapper::new(tokenizer))
     }
 
     /// Pools the embedding tensor by extracting the CLS token ([1, SEQLENGTH, DIM] -> [1, DIM])
     /// &RETURNS&: Tensor
-    fn pooled_cls_embedding(result: &CoreTensor) -> Result<CoreTensor, Error> {
+    fn pooled_cls_embedding(result: &CoreTensor) -> std::result::Result<CoreTensor, Error> {
         // 1) sanity-check that we have a 3D tensor
         let (_batch, _seq_len, _hidden_size) = result.dims3().map_err(wrap_candle_err)?;
     
@@ -298,14 +294,14 @@ impl EmbeddingModel {
         Ok(cls)
     }
 
-    fn pooled_embedding(result: &CoreTensor) -> Result<CoreTensor, Error> {
+    fn pooled_embedding(result: &CoreTensor) -> std::result::Result<CoreTensor, Error> {
         let (_n_sentence, n_tokens, _hidden_size) = result.dims3().map_err(wrap_candle_err)?;
         let sum = result.sum(1).map_err(wrap_candle_err)?;
         let mean = (sum / (n_tokens as f64)).map_err(wrap_candle_err)?;
         Ok(mean)
     }
 
-    fn pooled_normalized_embedding(result: &CoreTensor) -> Result<CoreTensor, Error> {
+    fn pooled_normalized_embedding(result: &CoreTensor) -> std::result::Result<CoreTensor, Error> {
         let mean = Self::pooled_embedding(result)?;
         let norm = Self::normalize_l2(&mean).map_err(wrap_candle_err)?;
         Ok(norm)
@@ -315,13 +311,11 @@ impl EmbeddingModel {
         &self,
         prompt: String,
         model: &EmbeddingModelVariant,
-        tokenizer: &Tokenizer,
-    ) -> Result<CoreTensor, Error> {
+        tokenizer: &TokenizerWrapper,
+    ) -> std::result::Result<CoreTensor, Error> {
         let tokens = tokenizer
-            .encode(prompt, true)
-            .map_err(wrap_std_err)?
-            .get_ids()
-            .to_vec();
+            .encode(&prompt, true)
+            .map_err(wrap_candle_err)?;
         let token_ids = CoreTensor::new(&tokens[..], &self.0.device)
             .map_err(wrap_candle_err)?
             .unsqueeze(0)
@@ -355,9 +349,9 @@ impl EmbeddingModel {
         &self,
         prompt: String,
         model: &EmbeddingModelVariant,
-        tokenizer: &Tokenizer,
+        tokenizer: &TokenizerWrapper,
         pooling_method: &str,
-    ) -> Result<CoreTensor, Error> {
+    ) -> std::result::Result<CoreTensor, Error> {
         let result = self.compute_embeddings(prompt, model, tokenizer)?;
         match pooling_method {
             "pooled" => Self::pooled_embedding(&result),
@@ -367,8 +361,7 @@ impl EmbeddingModel {
         }
     }
 
-    #[allow(dead_code)]
-    fn normalize_l2(v: &CoreTensor) -> Result<CoreTensor, candle_core::Error> {
+    fn normalize_l2(v: &CoreTensor) -> candle_core::Result<CoreTensor> {
         v.broadcast_div(&v.sqr()?.sum_keepdim(1)?.sqrt()?)
     }
 
@@ -392,9 +385,17 @@ impl EmbeddingModel {
     pub fn __str__(&self) -> String {
         self.__repr__()
     }
+
+    /// Get the tokenizer used by this model
+    pub fn tokenizer(&self) -> Result<crate::ruby::tokenizer::Tokenizer> {
+        match &self.0.tokenizer {
+            Some(tokenizer) => Ok(crate::ruby::tokenizer::Tokenizer(tokenizer.clone())),
+            None => Err(magnus::Error::new(magnus::exception::runtime_error(), "No tokenizer loaded for this model"))
+        }
+    }
 }
 
-pub fn init(rb_candle: RModule) -> Result<(), Error> {
+pub fn init(rb_candle: RModule) -> Result<()> {
     let rb_embedding_model = rb_candle.define_class("EmbeddingModel", class::object())?;
     rb_embedding_model.define_singleton_method("_create", function!(EmbeddingModel::new, 5))?;
     // Expose embedding with an optional pooling_method argument (default: "pooled")
@@ -406,5 +407,6 @@ pub fn init(rb_candle: RModule) -> Result<(), Error> {
     rb_embedding_model.define_method("embedding_model_type", method!(EmbeddingModel::embedding_model_type, 0))?;
     rb_embedding_model.define_method("to_s", method!(EmbeddingModel::__str__, 0))?;
     rb_embedding_model.define_method("inspect", method!(EmbeddingModel::__repr__, 0))?;
+    rb_embedding_model.define_method("tokenizer", method!(EmbeddingModel::tokenizer, 0))?;
     Ok(())
 }

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

@@ -1,8 +1,8 @@
 use magnus::{function, method, prelude::*, Error, Module, RArray, RHash, RModule, Ruby, TryConvert, Value};
 use std::cell::RefCell;
 
-use crate::llm::{GenerationConfig as RustGenerationConfig, TextGenerator, mistral::Mistral as RustMistral, llama::Llama as RustLlama, gemma::Gemma as RustGemma};
-use crate::ruby::{Result as RbResult, Device as RbDevice};
+use crate::llm::{GenerationConfig as RustGenerationConfig, TextGenerator, mistral::Mistral as RustMistral, llama::Llama as RustLlama, gemma::Gemma as RustGemma, QuantizedGGUF as RustQuantizedGGUF};
+use crate::ruby::{Result, Device};
 
 // Use an enum to handle different model types instead of trait objects
 #[derive(Debug)]
@@ -10,6 +10,7 @@ enum ModelType {
     Mistral(RustMistral),
     Llama(RustLlama),
     Gemma(RustGemma),
+    QuantizedGGUF(RustQuantizedGGUF),
 }
 
 impl ModelType {
@@ -18,6 +19,7 @@ impl ModelType {
             ModelType::Mistral(m) => m.generate(prompt, config),
             ModelType::Llama(m) => m.generate(prompt, config),
             ModelType::Gemma(m) => m.generate(prompt, config),
+            ModelType::QuantizedGGUF(m) => m.generate(prompt, config),
         }
     }
 
@@ -31,15 +33,7 @@ impl ModelType {
             ModelType::Mistral(m) => m.generate_stream(prompt, config, callback),
             ModelType::Llama(m) => m.generate_stream(prompt, config, callback),
             ModelType::Gemma(m) => m.generate_stream(prompt, config, callback),
-        }
-    }
-
-    #[allow(dead_code)]
-    fn model_name(&self) -> &str {
-        match self {
-            ModelType::Mistral(m) => m.model_name(),
-            ModelType::Llama(m) => m.model_name(),
-            ModelType::Gemma(m) => m.model_name(),
+            ModelType::QuantizedGGUF(m) => m.generate_stream(prompt, config, callback),
         }
     }
     
@@ -48,6 +42,7 @@ impl ModelType {
             ModelType::Mistral(m) => m.clear_cache(),
             ModelType::Llama(m) => m.clear_cache(),
             ModelType::Gemma(m) => m.clear_cache(),
+            ModelType::QuantizedGGUF(m) => m.clear_cache(),
         }
     }
     
@@ -72,6 +67,7 @@ impl ModelType {
             },
             ModelType::Llama(m) => m.apply_chat_template(messages),
             ModelType::Gemma(m) => m.apply_chat_template(messages),
+            ModelType::QuantizedGGUF(m) => m.apply_chat_template(messages),
         }
     }
 }
@@ -83,7 +79,7 @@ pub struct GenerationConfig {
 }
 
 impl GenerationConfig {
-    pub fn new(kwargs: RHash) -> RbResult<Self> {
+    pub fn new(kwargs: RHash) -> Result<Self> {
         let mut config = RustGenerationConfig::default();
         
         // Extract values from kwargs manually
@@ -144,6 +140,12 @@ impl GenerationConfig {
             }
         }
         
+        if let Some(value) = kwargs.get(magnus::Symbol::new("debug_tokens")) {
+            if let Ok(v) = TryConvert::try_convert(value) {
+                config.debug_tokens = v;
+            }
+        }
+        
         Ok(Self { inner: config })
     }
 
@@ -185,6 +187,10 @@ impl GenerationConfig {
     pub fn include_prompt(&self) -> bool {
         self.inner.include_prompt
     }
+
+    pub fn debug_tokens(&self) -> bool {
+        self.inner.debug_tokens
+    }
 }
 
 #[derive(Clone, Debug)]
@@ -192,13 +198,13 @@ impl GenerationConfig {
 pub struct LLM {
     model: std::sync::Arc<std::sync::Mutex<RefCell<ModelType>>>,
     model_id: String,
-    device: RbDevice,
+    device: Device,
 }
 
 impl LLM {
     /// Create a new LLM from a pretrained model
-    pub fn from_pretrained(model_id: String, device: Option<RbDevice>) -> RbResult<Self> {
-        let device = device.unwrap_or(RbDevice::Cpu);
+    pub fn from_pretrained(model_id: String, device: Option<Device>) -> Result<Self> {
+        let device = device.unwrap_or(Device::Cpu);
         let candle_device = device.as_device()?;
         
         // For now, we'll use tokio runtime directly
@@ -206,31 +212,51 @@ impl LLM {
         let rt = tokio::runtime::Runtime::new()
             .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to create runtime: {}", e)))?;
         
-        // Determine model type from ID and load appropriately
+        // Determine model type from ID and whether it's quantized
         let model_lower = model_id.to_lowercase();
-        let model = if model_lower.contains("mistral") {
-            let mistral = rt.block_on(async {
-                RustMistral::from_pretrained(&model_id, candle_device).await
-            })
-            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to load model: {}", e)))?;
-            ModelType::Mistral(mistral)
-        } else if model_lower.contains("llama") || model_lower.contains("meta-llama") || model_lower.contains("tinyllama") {
-            let llama = rt.block_on(async {
-                RustLlama::from_pretrained(&model_id, candle_device).await
-            })
-            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to load model: {}", e)))?;
-            ModelType::Llama(llama)
-        } else if model_lower.contains("gemma") || model_lower.contains("google/gemma") {
-            let gemma = rt.block_on(async {
-                RustGemma::from_pretrained(&model_id, candle_device).await
+        let is_quantized = model_lower.contains("gguf") || model_lower.contains("-q4") || model_lower.contains("-q5") || model_lower.contains("-q8");
+        
+        let model = if is_quantized {
+            // Extract tokenizer source if provided in model_id
+            let (model_id_clean, tokenizer_source) = if let Some(pos) = model_id.find("@@") {
+                let (id, _tok) = model_id.split_at(pos);
+                (id.to_string(), Some(&model_id[pos+2..]))
+            } else {
+                (model_id.clone(), None)
+            };
+            
+            // Use unified GGUF loader for all quantized models
+            let gguf_model = rt.block_on(async {
+                RustQuantizedGGUF::from_pretrained(&model_id_clean, candle_device, tokenizer_source).await
             })
-            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to load model: {}", e)))?;
-            ModelType::Gemma(gemma)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to load GGUF model: {}", e)))?;
+            ModelType::QuantizedGGUF(gguf_model)
         } else {
-            return Err(Error::new(
-                magnus::exception::runtime_error(),
-                format!("Unsupported model type: {}. Currently Mistral, Llama, and Gemma models are supported.", model_id),
-            ));
+            // Load non-quantized models
+            if model_lower.contains("mistral") {
+                let mistral = rt.block_on(async {
+                    RustMistral::from_pretrained(&model_id, candle_device).await
+                })
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to load model: {}", e)))?;
+                ModelType::Mistral(mistral)
+            } else if model_lower.contains("llama") || model_lower.contains("meta-llama") || model_lower.contains("tinyllama") {
+                let llama = rt.block_on(async {
+                    RustLlama::from_pretrained(&model_id, candle_device).await
+                })
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to load model: {}", e)))?;
+                ModelType::Llama(llama)
+            } else if model_lower.contains("gemma") || model_lower.contains("google/gemma") {
+                let gemma = rt.block_on(async {
+                    RustGemma::from_pretrained(&model_id, candle_device).await
+                })
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to load model: {}", e)))?;
+                ModelType::Gemma(gemma)
+            } else {
+                return Err(Error::new(
+                    magnus::exception::runtime_error(),
+                    format!("Unsupported model type: {}. Currently Mistral, Llama, and Gemma models are supported.", model_id),
+                ));
+            }
         };
         
         Ok(Self {
@@ -241,12 +267,15 @@ impl LLM {
     }
 
     /// Generate text from a prompt
-    pub fn generate(&self, prompt: String, config: Option<&GenerationConfig>) -> RbResult<String> {
+    pub fn generate(&self, prompt: String, config: Option<&GenerationConfig>) -> Result<String> {
         let config = config
             .map(|c| c.inner.clone())
             .unwrap_or_default();
         
-        let model = self.model.lock().unwrap();
+        let model = match self.model.lock() {
+            Ok(guard) => guard,
+            Err(poisoned) => poisoned.into_inner(),
+        };
         let mut model_ref = model.borrow_mut();
         
         model_ref.generate(&prompt, &config)
@@ -254,7 +283,7 @@ impl LLM {
     }
 
     /// Generate text with streaming output
-    pub fn generate_stream(&self, prompt: String, config: Option<&GenerationConfig>) -> RbResult<String> {
+    pub fn generate_stream(&self, prompt: String, config: Option<&GenerationConfig>) -> Result<String> {
         let config = config
             .map(|c| c.inner.clone())
             .unwrap_or_default();
@@ -266,7 +295,10 @@ impl LLM {
         }
         let block = block.unwrap();
         
-        let model = self.model.lock().unwrap();
+        let model = match self.model.lock() {
+            Ok(guard) => guard,
+            Err(poisoned) => poisoned.into_inner(),
+        };
         let mut model_ref = model.borrow_mut();
         
         let result = model_ref.generate_stream(&prompt, &config, |token| {
@@ -283,20 +315,44 @@ impl LLM {
     }
 
     /// Get the device the model is running on
-    pub fn device(&self) -> RbDevice {
+    pub fn device(&self) -> Device {
         self.device
     }
+
+    /// Get the tokenizer used by this model
+    pub fn tokenizer(&self) -> Result<crate::ruby::tokenizer::Tokenizer> {
+        let model = match self.model.lock() {
+            Ok(guard) => guard,
+            Err(poisoned) => poisoned.into_inner(),
+        };
+        let model_ref = model.borrow();
+        
+        // Clone the tokenizer from the model
+        match &*model_ref {
+            ModelType::Mistral(m) => Ok(crate::ruby::tokenizer::Tokenizer(m.tokenizer().clone())),
+            ModelType::Llama(m) => Ok(crate::ruby::tokenizer::Tokenizer(m.tokenizer().clone())),
+            ModelType::Gemma(m) => Ok(crate::ruby::tokenizer::Tokenizer(m.tokenizer().clone())),
+            ModelType::QuantizedGGUF(m) => Ok(crate::ruby::tokenizer::Tokenizer(m.tokenizer().clone())),
+        }
+    }
     
     /// Clear the model's cache (e.g., KV cache for transformers)
-    pub fn clear_cache(&self) -> RbResult<()> {
-        let model = self.model.lock().unwrap();
+    pub fn clear_cache(&self) -> Result<()> {
+        let model = match self.model.lock() {
+            Ok(guard) => guard,
+            Err(poisoned) => {
+                // If the mutex is poisoned, we can still recover the data
+                // This happens when another thread panicked while holding the lock
+                poisoned.into_inner()
+            }
+        };
         let mut model_ref = model.borrow_mut();
         model_ref.clear_cache();
         Ok(())
     }
     
     /// Apply chat template to messages
-    pub fn apply_chat_template(&self, messages: RArray) -> RbResult<String> {
+    pub fn apply_chat_template(&self, messages: RArray) -> Result<String> {
         // Convert Ruby array to JSON values
         let json_messages: Vec<serde_json::Value> = messages
             .into_iter()
@@ -323,7 +379,10 @@ impl LLM {
             })
             .collect();
         
-        let model = self.model.lock().unwrap();
+        let model = match self.model.lock() {
+            Ok(guard) => guard,
+            Err(poisoned) => poisoned.into_inner(),
+        };
         let model_ref = model.borrow();
         
         model_ref.apply_chat_template(&json_messages)
@@ -332,7 +391,7 @@ impl LLM {
 }
 
 // Define a standalone function for from_pretrained that handles variable arguments
-fn from_pretrained_wrapper(args: &[Value]) -> RbResult<LLM> {
+fn from_pretrained_wrapper(args: &[Value]) -> Result<LLM> {
     match args.len() {
         1 => {
             let model_id: String = TryConvert::try_convert(args[0])?;
@@ -340,7 +399,7 @@ fn from_pretrained_wrapper(args: &[Value]) -> RbResult<LLM> {
         },
         2 => {
             let model_id: String = TryConvert::try_convert(args[0])?;
-            let device: RbDevice = TryConvert::try_convert(args[1])?;
+            let device: Device = TryConvert::try_convert(args[1])?;
             LLM::from_pretrained(model_id, Some(device))
         },
         _ => Err(Error::new(
@@ -350,7 +409,7 @@ fn from_pretrained_wrapper(args: &[Value]) -> RbResult<LLM> {
     }
 }
 
-pub fn init_llm(rb_candle: RModule) -> RbResult<()> {
+pub fn init_llm(rb_candle: RModule) -> Result<()> {
     let rb_generation_config = rb_candle.define_class("GenerationConfig", magnus::class::object())?;
     rb_generation_config.define_singleton_method("new", function!(GenerationConfig::new, 1))?;
     rb_generation_config.define_singleton_method("default", function!(GenerationConfig::default, 0))?;
@@ -363,6 +422,7 @@ pub fn init_llm(rb_candle: RModule) -> RbResult<()> {
     rb_generation_config.define_method("seed", method!(GenerationConfig::seed, 0))?;
     rb_generation_config.define_method("stop_sequences", method!(GenerationConfig::stop_sequences, 0))?;
     rb_generation_config.define_method("include_prompt", method!(GenerationConfig::include_prompt, 0))?;
+    rb_generation_config.define_method("debug_tokens", method!(GenerationConfig::debug_tokens, 0))?;
     
     let rb_llm = rb_candle.define_class("LLM", magnus::class::object())?;
     rb_llm.define_singleton_method("_from_pretrained", function!(from_pretrained_wrapper, -1))?;
@@ -370,6 +430,7 @@ pub fn init_llm(rb_candle: RModule) -> RbResult<()> {
     rb_llm.define_method("_generate_stream", method!(LLM::generate_stream, 2))?;
     rb_llm.define_method("model_name", method!(LLM::model_name, 0))?;
     rb_llm.define_method("device", method!(LLM::device, 0))?;
+    rb_llm.define_method("tokenizer", method!(LLM::tokenizer, 0))?;
     rb_llm.define_method("clear_cache", method!(LLM::clear_cache, 0))?;
     rb_llm.define_method("apply_chat_template", method!(LLM::apply_chat_template, 1))?;
     

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

@@ -2,17 +2,16 @@ pub mod embedding_model;
 pub mod tensor;
 pub mod device;
 pub mod dtype;
-pub mod qtensor;
 pub mod result;
 pub mod errors;
 pub mod utils;
 pub mod llm;
+pub mod tokenizer;
 
 pub use embedding_model::{EmbeddingModel, EmbeddingModelInner};
 pub use tensor::Tensor;
 pub use device::Device;
 pub use dtype::DType;
-pub use qtensor::QTensor;
 pub use result::Result;
 
 // Re-export for convenience