red-candle 1.0.0.pre.7 → 1.0.1

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

@@ -0,0 +1,269 @@
+use magnus::{class, function, method, prelude::*, Error, Module, RArray, RHash, RModule, TryConvert};
+use crate::tokenizer::{TokenizerWrapper as InnerTokenizer, loader::TokenizerLoader};
+use crate::ruby::Result;
+
+#[derive(Clone, Debug)]
+#[magnus::wrap(class = "Candle::Tokenizer", free_immediately, size)]
+pub struct Tokenizer(pub InnerTokenizer);
+
+impl Tokenizer {
+    /// Create a new tokenizer from a file path
+    pub fn from_file(path: String) -> Result<Self> {
+        let tokenizer = TokenizerLoader::from_file(&path)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        Ok(Self(InnerTokenizer::new(tokenizer)))
+    }
+
+    /// Create a new tokenizer from HuggingFace model ID
+    pub fn from_pretrained(model_id: String) -> Result<Self> {
+        // Use tokio runtime for async operations
+        let rt = tokio::runtime::Runtime::new()
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to create runtime: {}", e)))?;
+        
+        let tokenizer = rt.block_on(async {
+            TokenizerLoader::from_hf_hub(&model_id, None).await
+        })
+        .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        
+        Ok(Self(InnerTokenizer::new(tokenizer)))
+    }
+
+    /// Encode text into token IDs
+    pub fn encode(&self, text: String, add_special_tokens: Option<bool>) -> Result<RArray> {
+        let add_special = add_special_tokens.unwrap_or(true);
+        let token_ids = self.0.encode(&text, add_special)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        
+        Ok(RArray::from_vec(token_ids.into_iter().map(|id| id as i64).collect()))
+    }
+    
+    /// Encode text into token strings (words/subwords)
+    pub fn encode_to_tokens(&self, text: String, add_special_tokens: Option<bool>) -> Result<RArray> {
+        let add_special = add_special_tokens.unwrap_or(true);
+        let token_ids = self.0.encode(&text, add_special)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        
+        let mut tokens = Vec::new();
+        for id in token_ids {
+            let token = self.0.token_to_piece(id)
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            tokens.push(token);
+        }
+        
+        Ok(RArray::from_vec(tokens))
+    }
+
+    /// Encode multiple texts in batch
+    pub fn encode_batch(&self, texts: RArray, add_special_tokens: Option<bool>) -> Result<RArray> {
+        let texts: Vec<String> = texts.to_vec()?;
+        let add_special = add_special_tokens.unwrap_or(true);
+        
+        let token_ids_batch = self.0.encode_batch(texts, add_special)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        
+        let result = RArray::new();
+        for token_ids in token_ids_batch {
+            result.push(RArray::from_vec(token_ids.into_iter().map(|id| id as i64).collect()))?;
+        }
+        
+        Ok(result)
+    }
+    
+    /// Encode multiple texts in batch, returning token strings
+    pub fn encode_batch_to_tokens(&self, texts: RArray, add_special_tokens: Option<bool>) -> Result<RArray> {
+        let texts: Vec<String> = texts.to_vec()?;
+        let add_special = add_special_tokens.unwrap_or(true);
+        
+        let token_ids_batch = self.0.encode_batch(texts, add_special)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        
+        let result = RArray::new();
+        for token_ids in token_ids_batch {
+            let mut tokens = Vec::new();
+            for id in token_ids {
+                let token = self.0.token_to_piece(id)
+                    .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+                tokens.push(token);
+            }
+            result.push(RArray::from_vec(tokens))?;
+        }
+        
+        Ok(result)
+    }
+
+    /// Encode text and return both token IDs and token strings
+    pub fn encode_with_tokens(&self, text: String, add_special_tokens: Option<bool>) -> Result<RHash> {
+        let add_special = add_special_tokens.unwrap_or(true);
+        let token_ids = self.0.encode(&text, add_special)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        
+        let mut tokens = Vec::new();
+        for &id in &token_ids {
+            let token = self.0.token_to_piece(id)
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+            tokens.push(token);
+        }
+        
+        let hash = RHash::new();
+        hash.aset("ids", RArray::from_vec(token_ids.into_iter().map(|id| id as i64).collect()))?;
+        hash.aset("tokens", RArray::from_vec(tokens))?;
+        
+        Ok(hash)
+    }
+    
+    /// Decode token IDs back to text
+    pub fn decode(&self, token_ids: RArray, skip_special_tokens: Option<bool>) -> Result<String> {
+        let token_ids: Vec<i64> = token_ids.to_vec()?;
+        let token_ids: Vec<u32> = token_ids.into_iter()
+            .map(|id| id as u32)
+            .collect();
+        let skip_special = skip_special_tokens.unwrap_or(true);
+        
+        self.0.decode(&token_ids, skip_special)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))
+    }
+
+    /// Get the string representation of a single token ID
+    pub fn id_to_token(&self, token_id: i64) -> Result<String> {
+        self.0.token_to_piece(token_id as u32)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))
+    }
+
+    /// Get the vocabulary as a hash of token string to ID
+    pub fn get_vocab(&self, with_added_tokens: Option<bool>) -> Result<RHash> {
+        let with_added = with_added_tokens.unwrap_or(true);
+        let vocab = self.0.inner().get_vocab(with_added);
+        
+        let hash = RHash::new();
+        for (token, id) in vocab {
+            hash.aset(token, id as i64)?;
+        }
+        
+        Ok(hash)
+    }
+
+    /// Get vocabulary size
+    pub fn vocab_size(&self, with_added_tokens: Option<bool>) -> usize {
+        let with_added = with_added_tokens.unwrap_or(true);
+        self.0.inner().get_vocab_size(with_added)
+    }
+
+    /// Enable padding - returns a new tokenizer with padding enabled
+    pub fn with_padding(&self, kwargs: RHash) -> Result<Self> {
+        use tokenizers::{PaddingParams, PaddingStrategy, PaddingDirection};
+        
+        let mut params = PaddingParams::default();
+        
+        // Extract parameters from kwargs
+        if let Some(length) = kwargs.get(magnus::Symbol::new("length")) {
+            if let Ok(len) = usize::try_convert(length) {
+                params.strategy = PaddingStrategy::Fixed(len);
+            }
+        }
+        
+        if let Some(max_length) = kwargs.get(magnus::Symbol::new("max_length")) {
+            if let Ok(_) = usize::try_convert(max_length) {
+                params.strategy = PaddingStrategy::BatchLongest;
+            }
+        }
+        
+        if let Some(direction) = kwargs.get(magnus::Symbol::new("direction")) {
+            if let Ok(dir) = String::try_convert(direction) {
+                params.direction = match dir.as_str() {
+                    "right" => PaddingDirection::Right,
+                    "left" => PaddingDirection::Left,
+                    _ => PaddingDirection::Right,
+                };
+            }
+        }
+        
+        if let Some(pad_id) = kwargs.get(magnus::Symbol::new("pad_id")) {
+            if let Ok(id) = u32::try_convert(pad_id) {
+                params.pad_id = id;
+            }
+        }
+        
+        if let Some(pad_token) = kwargs.get(magnus::Symbol::new("pad_token")) {
+            if let Ok(token) = String::try_convert(pad_token) {
+                params.pad_token = token;
+            }
+        }
+        
+        let mut new_tokenizer = self.0.clone();
+        let _ = new_tokenizer.inner_mut().with_padding(Some(params));
+        Ok(Self(new_tokenizer))
+    }
+
+    /// Enable truncation - returns a new tokenizer with truncation enabled
+    pub fn with_truncation(&self, max_length: usize) -> Result<Self> {
+        use tokenizers::{TruncationParams, TruncationStrategy, TruncationDirection};
+        
+        let params = TruncationParams {
+            max_length,
+            strategy: TruncationStrategy::LongestFirst,
+            stride: 0,
+            direction: TruncationDirection::Right,
+        };
+        
+        let mut new_tokenizer = self.0.clone();
+        let _ = new_tokenizer.inner_mut().with_truncation(Some(params));
+        Ok(Self(new_tokenizer))
+    }
+
+    /// Get special tokens information
+    pub fn get_special_tokens(&self) -> Result<RHash> {
+        let hash = RHash::new();
+        
+        // Common special tokens
+        let special_tokens = vec![
+            ("[CLS]", "cls_token"),
+            ("[SEP]", "sep_token"),
+            ("[PAD]", "pad_token"),
+            ("[UNK]", "unk_token"),
+            ("[MASK]", "mask_token"),
+            ("<s>", "bos_token"),
+            ("</s>", "eos_token"),
+        ];
+        
+        let vocab = self.0.inner().get_vocab(true);
+        
+        for (token, name) in special_tokens {
+            if let Some(id) = vocab.get(token) {
+                hash.aset(name, *id as i64)?;
+            }
+        }
+        
+        Ok(hash)
+    }
+
+    /// String representation
+    pub fn inspect(&self) -> String {
+        format!("#<Candle::Tokenizer vocab_size={}>", self.vocab_size(Some(true)))
+    }
+}
+
+pub fn init(rb_candle: RModule) -> Result<()> {
+    let tokenizer_class = rb_candle.define_class("Tokenizer", class::object())?;
+    
+    // Class methods
+    tokenizer_class.define_singleton_method("from_file", function!(Tokenizer::from_file, 1))?;
+    tokenizer_class.define_singleton_method("from_pretrained", function!(Tokenizer::from_pretrained, 1))?;
+    
+    // Instance methods
+    tokenizer_class.define_method("encode", method!(Tokenizer::encode, 2))?;
+    tokenizer_class.define_method("encode_to_tokens", method!(Tokenizer::encode_to_tokens, 2))?;
+    tokenizer_class.define_method("encode_with_tokens", method!(Tokenizer::encode_with_tokens, 2))?;
+    tokenizer_class.define_method("encode_batch", method!(Tokenizer::encode_batch, 2))?;
+    tokenizer_class.define_method("encode_batch_to_tokens", method!(Tokenizer::encode_batch_to_tokens, 2))?;
+    tokenizer_class.define_method("decode", method!(Tokenizer::decode, 2))?;
+    tokenizer_class.define_method("id_to_token", method!(Tokenizer::id_to_token, 1))?;
+    tokenizer_class.define_method("get_vocab", method!(Tokenizer::get_vocab, 1))?;
+    tokenizer_class.define_method("vocab_size", method!(Tokenizer::vocab_size, 1))?;
+    tokenizer_class.define_method("with_padding", method!(Tokenizer::with_padding, 1))?;
+    tokenizer_class.define_method("with_truncation", method!(Tokenizer::with_truncation, 1))?;
+    tokenizer_class.define_method("get_special_tokens", method!(Tokenizer::get_special_tokens, 0))?;
+    tokenizer_class.define_method("inspect", method!(Tokenizer::inspect, 0))?;
+    tokenizer_class.define_method("to_s", method!(Tokenizer::inspect, 0))?;
+    
+    Ok(())
+}

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

@@ -1,11 +1,10 @@
-use magnus::{function, Error, Module, Object};
+use magnus::{function, Module, Object};
 
-use ::candle_core::Tensor;
+use ::candle_core::Tensor as CoreTensor;
 
-use crate::ruby::errors::wrap_candle_err;
-use crate::ruby::{Result as RbResult, Tensor as RbTensor};
+use crate::ruby::Result;
 
-pub fn actual_index(t: &Tensor, dim: usize, index: i64) -> candle_core::Result<usize> {
+pub fn actual_index(t: &CoreTensor, dim: usize, index: i64) -> candle_core::Result<usize> {
     let dim = t.dim(dim)?;
     if 0 <= index {
         let index = index as usize;
@@ -21,7 +20,7 @@ pub fn actual_index(t: &Tensor, dim: usize, index: i64) -> candle_core::Result<u
     }
 }
 
-pub fn actual_dim(t: &Tensor, dim: i64) -> candle_core::Result<usize> {
+pub fn actual_dim(t: &CoreTensor, dim: i64) -> candle_core::Result<usize> {
     let rank = t.rank();
     if 0 <= dim {
         let dim = dim as usize;
@@ -61,7 +60,7 @@ fn get_num_threads() -> usize {
     candle_core::utils::get_num_threads()
 }
 
-pub fn candle_utils(rb_candle: magnus::RModule) -> Result<(), Error> {
+pub fn candle_utils(rb_candle: magnus::RModule) -> Result<()> {
     let rb_utils = rb_candle.define_module("Utils")?;
     rb_utils.define_singleton_method("cuda_is_available", function!(cuda_is_available, 0))?;
     rb_utils.define_singleton_method("get_num_threads", function!(get_num_threads, 0))?;
@@ -69,20 +68,3 @@ pub fn candle_utils(rb_candle: magnus::RModule) -> Result<(), Error> {
     rb_utils.define_singleton_method("has_mkl", function!(has_mkl, 0))?;
     Ok(())
 }
-
-/// Applies the Softmax function to a given tensor.#
-/// &RETURNS&: Tensor
-#[allow(dead_code)]
-fn softmax(tensor: RbTensor, dim: i64) -> RbResult<RbTensor> {
-    let dim = actual_dim(&tensor, dim).map_err(wrap_candle_err)?;
-    let sm = candle_nn::ops::softmax(&tensor.0, dim).map_err(wrap_candle_err)?;
-    Ok(RbTensor(sm))
-}
-
-/// Applies the Sigmoid Linear Unit (SiLU) function to a given tensor.
-/// &RETURNS&: Tensor
-#[allow(dead_code)]
-fn silu(tensor: RbTensor) -> RbResult<RbTensor> {
-    let s = candle_nn::ops::silu(&tensor.0).map_err(wrap_candle_err)?;
-    Ok(RbTensor(s))
-}

data/ext/candle/src/tokenizer/loader.rs

@@ -0,0 +1,108 @@
+use candle_core::Result as CandleResult;
+use hf_hub::api::tokio::{Api, ApiRepo};
+use tokenizers::Tokenizer;
+use tokenizers::{PaddingParams, PaddingStrategy, TruncationParams};
+use std::path::PathBuf;
+
+/// Standard padding configuration for models
+pub fn standard_padding_params() -> PaddingParams {
+    PaddingParams {
+        strategy: PaddingStrategy::BatchLongest,
+        direction: tokenizers::PaddingDirection::Left,
+        pad_to_multiple_of: None,
+        pad_id: 0,
+        pad_type_id: 0,
+        pad_token: "[PAD]".to_string(),
+    }
+}
+
+/// Unified tokenizer loader with common download logic
+pub struct TokenizerLoader;
+
+impl TokenizerLoader {
+    /// Load tokenizer from a local file path
+    pub fn from_file(path: &str) -> CandleResult<Tokenizer> {
+        Tokenizer::from_file(path)
+            .map_err(|e| candle_core::Error::Msg(format!("Failed to load tokenizer from file: {}", e)))
+    }
+
+    /// Download and load tokenizer from HuggingFace
+    pub async fn from_hf_hub(repo_id: &str, filename: Option<&str>) -> CandleResult<Tokenizer> {
+        let api = Api::new()
+            .map_err(|e| candle_core::Error::Msg(format!("Failed to create HF API: {}", e)))?;
+        
+        let repo = api.model(repo_id.to_string());
+        let tokenizer_path = Self::download_tokenizer_file(&repo, filename).await?;
+        
+        Self::from_file(tokenizer_path.to_str()
+            .ok_or_else(|| candle_core::Error::Msg("Invalid tokenizer path".to_string()))?)
+    }
+
+    /// Download tokenizer file from repository
+    async fn download_tokenizer_file(repo: &ApiRepo, filename: Option<&str>) -> CandleResult<PathBuf> {
+        if let Some(file) = filename {
+            // Try specific filename
+            repo.get(file).await
+                .map_err(|e| candle_core::Error::Msg(
+                    format!("Failed to download tokenizer file '{}': {}", file, e)
+                ))
+        } else {
+            // Try common tokenizer filenames in order
+            let filenames = ["tokenizer.json", "tokenizer.model"];
+            
+            for file in filenames {
+                if let Ok(path) = repo.get(file).await {
+                    return Ok(path);
+                }
+            }
+            
+            Err(candle_core::Error::Msg(
+                "No tokenizer file found. Tried: tokenizer.json, tokenizer.model".to_string()
+            ))
+        }
+    }
+
+    /// Configure tokenizer with standard padding for batch processing
+    pub fn with_padding(mut tokenizer: Tokenizer, padding_params: Option<PaddingParams>) -> Tokenizer {
+        let params = padding_params.unwrap_or_else(standard_padding_params);
+        let _ = tokenizer.with_padding(Some(params));
+        tokenizer
+    }
+
+    /// Configure tokenizer with truncation
+    pub fn with_truncation(mut tokenizer: Tokenizer, max_length: usize) -> Tokenizer {
+        let _ = tokenizer.with_truncation(Some(TruncationParams {
+            max_length,
+            strategy: tokenizers::TruncationStrategy::LongestFirst,
+            stride: 0,
+            direction: tokenizers::TruncationDirection::Right,
+        }));
+        tokenizer
+    }
+
+    /// Download tokenizer from a specific source (for GGUF models)
+    pub async fn from_source(api: &Api, source: &str) -> CandleResult<PathBuf> {
+        // Check if it's a local file path
+        if source.ends_with(".json") && std::path::Path::new(source).exists() {
+            return Ok(PathBuf::from(source));
+        }
+        
+        // Otherwise treat it as a HuggingFace repo
+        let repo = api.model(source.to_string());
+        
+        // Try tokenizer.json first
+        if let Ok(path) = repo.get("tokenizer.json").await {
+            return Ok(path);
+        }
+        
+        // Try tokenizer.model (for models that use sentencepiece)
+        if let Ok(path) = repo.get("tokenizer.model").await {
+            return Ok(path);
+        }
+        
+        Err(candle_core::Error::Msg(format!(
+            "Failed to find tokenizer in specified source: {}. Please check network connectivity and that the model repository exists.",
+            source
+        )))
+    }
+}

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

@@ -0,0 +1,103 @@
+use candle_core::Result as CandleResult;
+use tokenizers::Tokenizer;
+
+pub mod loader;
+
+/// Common structure for managing tokenizer
+#[derive(Debug, Clone)]
+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)))
+    }
+    
+    /// 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_len = full_text
+            .chars()
+            .zip(previous_text.chars())
+            .take_while(|(a, b)| a == b)
+            .count();
+        
+        Ok(full_text.chars().skip(common_len).collect())
+    }
+    
+    /// 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 piece = self.token_to_piece(token)?;
+            result.push_str(&format!("[{}:{}]", token, piece));
+        }
+        Ok(result)
+    }
+
+    /// Encode a batch of texts (needed for reranker)
+    pub fn encode_batch(&self, texts: Vec<String>, add_special_tokens: bool) -> CandleResult<Vec<Vec<u32>>> {
+        let encodings = self.tokenizer
+            .encode_batch(texts, add_special_tokens)
+            .map_err(|e| candle_core::Error::Msg(format!("Tokenizer batch error: {}", e)))?;
+        
+        Ok(encodings.into_iter()
+            .map(|encoding| encoding.get_ids().to_vec())
+            .collect())
+    }
+
+    /// Get the underlying tokenizer (for advanced use cases)
+    pub fn inner(&self) -> &Tokenizer {
+        &self.tokenizer
+    }
+
+    /// Get a mutable reference to the underlying tokenizer (for configuration)
+    pub fn inner_mut(&mut self) -> &mut Tokenizer {
+        &mut self.tokenizer
+    }
+}

data/ext/candle/target/release/build/bindgen-0f89ba23b9ca1395/out/host-target.txt

@@ -0,0 +1 @@
+aarch64-apple-darwin