red-candle 1.0.0.pre.7 → 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,7 +361,7 @@ impl EmbeddingModel {
         }
     }
 
-    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()?)
     }
 
@@ -391,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")
@@ -405,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

@@ -2,7 +2,7 @@ use magnus::{function, method, prelude::*, Error, Module, RArray, RHash, RModule
 use std::cell::RefCell;
 
 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 as RbResult, Device as RbDevice};
+use crate::ruby::{Result, Device};
 
 // Use an enum to handle different model types instead of trait objects
 #[derive(Debug)]
@@ -79,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
@@ -198,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
@@ -267,7 +267,7 @@ 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();
@@ -283,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();
@@ -315,12 +315,29 @@ 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<()> {
+    pub fn clear_cache(&self) -> Result<()> {
         let model = match self.model.lock() {
             Ok(guard) => guard,
             Err(poisoned) => {
@@ -335,7 +352,7 @@ impl LLM {
     }
     
     /// 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()
@@ -374,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])?;
@@ -382,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(
@@ -392,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))?;
@@ -413,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