red-candle 1.0.0.pre.7 → 1.0.1

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

@@ -1,5 +1,4 @@
 use candle_core::{Device, Result as CandleResult};
-use tokenizers::Tokenizer;
 
 pub mod mistral;
 pub mod llama;
@@ -11,6 +10,7 @@ pub mod quantized_gguf;
 pub use generation_config::GenerationConfig;
 pub use text_generation::TextGeneration;
 pub use quantized_gguf::QuantizedGGUF;
+pub use crate::tokenizer::TokenizerWrapper;
 
 /// Trait for text generation models
 pub trait TextGenerator: Send + Sync {
@@ -37,92 +37,4 @@ pub trait TextGenerator: Send + Sync {
     
     /// 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)))
-    }
-    
-    /// 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)
-    }
 }

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

@@ -28,6 +28,11 @@ enum ModelType {
 }
 
 impl QuantizedGGUF {
+    /// Get the tokenizer
+    pub fn tokenizer(&self) -> &TokenizerWrapper {
+        &self.tokenizer
+    }
+    
     /// Load a quantized model from a GGUF file
     pub async fn from_pretrained(model_id: &str, device: Device, tokenizer_source: Option<&str>) -> CandleResult<Self> {
         // Check if user specified an exact GGUF filename

data/ext/candle/src/ner.rs

@@ -0,0 +1,423 @@
+use magnus::{class, function, method, prelude::*, Error, RModule, RArray, RHash};
+use candle_transformers::models::bert::{BertModel, Config};
+use candle_core::{Device as CoreDevice, Tensor, DType, Module as CanModule};
+use candle_nn::{VarBuilder, Linear};
+use hf_hub::{api::sync::Api, Repo, RepoType};
+use std::collections::HashMap;
+use serde::{Deserialize, Serialize};
+use crate::ruby::{Device, Result};
+use crate::tokenizer::{TokenizerWrapper, loader::TokenizerLoader};
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct NERConfig {
+    pub id2label: HashMap<i64, String>,
+    pub label2id: HashMap<String, i64>,
+}
+
+#[derive(Debug, Clone)]
+pub struct EntitySpan {
+    pub text: String,
+    pub label: String,
+    pub start: usize,
+    pub end: usize,
+    pub token_start: usize,
+    pub token_end: usize,
+    pub confidence: f32,
+}
+
+#[magnus::wrap(class = "Candle::NER", free_immediately, size)]
+pub struct NER {
+    model: BertModel,
+    tokenizer: TokenizerWrapper,
+    classifier: Linear,
+    config: NERConfig,
+    device: CoreDevice,
+    model_id: String,
+}
+
+impl NER {
+    pub fn new(model_id: String, device: Option<Device>, tokenizer_id: Option<String>) -> Result<Self> {
+        let device = device.unwrap_or(Device::Cpu).as_device()?;
+        
+        // Load model in a separate thread to avoid blocking
+        let device_clone = device.clone();
+        let model_id_clone = model_id.clone();
+        
+        let handle = std::thread::spawn(move || -> std::result::Result<(BertModel, TokenizerWrapper, Linear, NERConfig), Box<dyn std::error::Error + Send + Sync>> {
+            let api = Api::new()?;
+            let repo = api.repo(Repo::new(model_id_clone.clone(), RepoType::Model));
+            
+            // Download model files
+            let config_filename = repo.get("config.json")?;
+            
+            // Handle tokenizer loading with optional tokenizer_id
+            let tokenizer = if let Some(tok_id) = tokenizer_id {
+                // Use the specified tokenizer
+                let tok_repo = api.repo(Repo::new(tok_id, RepoType::Model));
+                let tokenizer_filename = tok_repo.get("tokenizer.json")?;
+                let tokenizer = tokenizers::Tokenizer::from_file(tokenizer_filename)?;
+                TokenizerLoader::with_padding(tokenizer, None)
+            } else {
+                // Try to load tokenizer from model repo
+                let tokenizer_filename = repo.get("tokenizer.json")?;
+                let tokenizer = tokenizers::Tokenizer::from_file(tokenizer_filename)?;
+                TokenizerLoader::with_padding(tokenizer, None)
+            };
+            let weights_filename = repo.get("pytorch_model.safetensors")
+                .or_else(|_| repo.get("model.safetensors"))?;
+            
+            // Load BERT config
+            let config_str = std::fs::read_to_string(&config_filename)?;
+            let config_json: serde_json::Value = serde_json::from_str(&config_str)?;
+            let bert_config: Config = serde_json::from_value(config_json.clone())?;
+            
+            // Extract NER label configuration
+            let id2label = config_json["id2label"]
+                .as_object()
+                .ok_or("Missing id2label in config")?
+                .iter()
+                .map(|(k, v)| {
+                    let id = k.parse::<i64>().unwrap_or(0);
+                    let label = v.as_str().unwrap_or("O").to_string();
+                    (id, label)
+                })
+                .collect::<HashMap<_, _>>();
+            
+            let label2id = id2label.iter()
+                .map(|(id, label)| (label.clone(), *id))
+                .collect::<HashMap<_, _>>();
+            
+            let num_labels = id2label.len();
+            let ner_config = NERConfig { id2label, label2id };
+            
+            // Load model weights
+            let vb = unsafe {
+                VarBuilder::from_mmaped_safetensors(&[weights_filename], DType::F32, &device_clone)?
+            };
+            
+            // Load BERT model
+            let model = BertModel::load(vb.pp("bert"), &bert_config)?;
+            
+            // Load classification head for token classification
+            let classifier = candle_nn::linear(
+                bert_config.hidden_size,
+                num_labels,
+                vb.pp("classifier")
+            )?;
+            
+            Ok((model, TokenizerWrapper::new(tokenizer), classifier, ner_config))
+        });
+        
+        match handle.join() {
+            Ok(Ok((model, tokenizer, classifier, config))) => {
+                Ok(Self {
+                    model,
+                    tokenizer,
+                    classifier,
+                    config,
+                    device,
+                    model_id,
+                })
+            }
+            Ok(Err(e)) => Err(Error::new(
+                magnus::exception::runtime_error(),
+                format!("Failed to load NER model: {}", e)
+            )),
+            Err(_) => Err(Error::new(
+                magnus::exception::runtime_error(),
+                "Thread panicked while loading NER model"
+            )),
+        }
+    }
+    
+    /// Extract entities from text with confidence scores
+    pub fn extract_entities(&self, text: String, confidence_threshold: Option<f64>) -> Result<RArray> {
+        let threshold = confidence_threshold.unwrap_or(0.9) as f32;
+        
+        // Tokenize the text
+        let encoding = self.tokenizer.inner().encode(text.as_str(), true)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Tokenization failed: {}", e)))?;
+        
+        let token_ids = encoding.get_ids();
+        let tokens = encoding.get_tokens();
+        let offsets = encoding.get_offsets();
+        
+        // Convert to tensors
+        let input_ids = Tensor::new(token_ids, &self.device)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?
+            .unsqueeze(0)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?; // Add batch dimension
+        
+        let attention_mask = Tensor::ones_like(&input_ids)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        let token_type_ids = Tensor::zeros_like(&input_ids)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        
+        // Forward pass through BERT
+        let output = self.model.forward(&input_ids, &token_type_ids, Some(&attention_mask))
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        
+        // Apply classifier to get logits for each token
+        let logits = self.classifier.forward(&output)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        
+        // Apply softmax to get probabilities
+        let probs = candle_nn::ops::softmax(&logits, 2)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        
+        // Get predictions and confidence scores
+        let probs_vec: Vec<Vec<f32>> = probs.squeeze(0)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?
+            .to_vec2()
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        
+        // Extract entities with BIO decoding
+        let entities = self.decode_entities(
+            &text,
+            &tokens.iter().map(|s| s.as_str()).collect::<Vec<_>>(),
+            offsets,
+            &probs_vec,
+            threshold
+        )?;
+        
+        // Convert to Ruby array
+        let result = RArray::new();
+        for entity in entities {
+            let hash = RHash::new();
+            hash.aset("text", entity.text)?;
+            hash.aset("label", entity.label)?;
+            hash.aset("start", entity.start)?;
+            hash.aset("end", entity.end)?;
+            hash.aset("confidence", entity.confidence)?;
+            hash.aset("token_start", entity.token_start)?;
+            hash.aset("token_end", entity.token_end)?;
+            result.push(hash)?;
+        }
+        
+        Ok(result)
+    }
+    
+    /// Get token-level predictions with labels and confidence scores
+    pub fn predict_tokens(&self, text: String) -> Result<RArray> {
+        // Tokenize the text
+        let encoding = self.tokenizer.inner().encode(text.as_str(), true)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Tokenization failed: {}", e)))?;
+        
+        let token_ids = encoding.get_ids();
+        let tokens = encoding.get_tokens();
+        
+        // Convert to tensors
+        let input_ids = Tensor::new(token_ids, &self.device)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?
+            .unsqueeze(0)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        
+        let attention_mask = Tensor::ones_like(&input_ids)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        let token_type_ids = Tensor::zeros_like(&input_ids)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        
+        // Forward pass
+        let output = self.model.forward(&input_ids, &token_type_ids, Some(&attention_mask))
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        let logits = self.classifier.forward(&output)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        let probs = candle_nn::ops::softmax(&logits, 2)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        
+        // Get predictions
+        let probs_vec: Vec<Vec<f32>> = probs.squeeze(0)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?
+            .to_vec2()
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), e.to_string()))?;
+        
+        // Build result array
+        let result = RArray::new();
+        for (i, (token, probs)) in tokens.iter().zip(probs_vec.iter()).enumerate() {
+            // Find best label
+            let (label_id, confidence) = probs.iter()
+                .enumerate()
+                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
+                .map(|(idx, conf)| (idx as i64, *conf))
+                .unwrap_or((0, 0.0));
+            
+            let label = self.config.id2label.get(&label_id)
+                .unwrap_or(&"O".to_string())
+                .clone();
+            
+            let token_info = RHash::new();
+            token_info.aset("token", token.to_string())?;
+            token_info.aset("label", label)?;
+            token_info.aset("confidence", confidence)?;
+            token_info.aset("index", i)?;
+            
+            // Add probability distribution if needed
+            let probs_hash = RHash::new();
+            for (id, label) in &self.config.id2label {
+                if let Some(prob) = probs.get(*id as usize) {
+                    probs_hash.aset(label.as_str(), *prob)?;
+                }
+            }
+            token_info.aset("probabilities", probs_hash)?;
+            
+            result.push(token_info)?;
+        }
+        
+        Ok(result)
+    }
+    
+    /// Decode BIO-tagged sequences into entity spans
+    fn decode_entities(
+        &self,
+        text: &str,
+        tokens: &[&str],
+        offsets: &[(usize, usize)],
+        probs: &[Vec<f32>],
+        threshold: f32,
+    ) -> Result<Vec<EntitySpan>> {
+        let mut entities = Vec::new();
+        let mut current_entity: Option<(String, usize, usize, Vec<f32>)> = None;
+        
+        for (i, (token, probs_vec)) in tokens.iter().zip(probs).enumerate() {
+            // Skip special tokens
+            if token.starts_with("[") && token.ends_with("]") {
+                continue;
+            }
+            
+            // Get predicted label
+            let (label_id, confidence) = probs_vec.iter()
+                .enumerate()
+                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
+                .map(|(idx, conf)| (idx as i64, *conf))
+                .unwrap_or((0, 0.0));
+            
+            let label = self.config.id2label.get(&label_id)
+                .unwrap_or(&"O".to_string())
+                .clone();
+            
+            // BIO decoding logic
+            if label == "O" || confidence < threshold {
+                // End current entity if exists
+                if let Some((entity_type, start_idx, end_idx, confidences)) = current_entity.take() {
+                    if let (Some(start_offset), Some(end_offset)) = 
+                        (offsets.get(start_idx), offsets.get(end_idx - 1)) {
+                        let entity_text = text[start_offset.0..end_offset.1].to_string();
+                        let avg_confidence = confidences.iter().sum::<f32>() / confidences.len() as f32;
+                        
+                        entities.push(EntitySpan {
+                            text: entity_text,
+                            label: entity_type,
+                            start: start_offset.0,
+                            end: end_offset.1,
+                            token_start: start_idx,
+                            token_end: end_idx,
+                            confidence: avg_confidence,
+                        });
+                    }
+                }
+            } else if label.starts_with("B-") {
+                // Begin new entity
+                if let Some((entity_type, start_idx, end_idx, confidences)) = current_entity.take() {
+                    if let (Some(start_offset), Some(end_offset)) = 
+                        (offsets.get(start_idx), offsets.get(end_idx - 1)) {
+                        let entity_text = text[start_offset.0..end_offset.1].to_string();
+                        let avg_confidence = confidences.iter().sum::<f32>() / confidences.len() as f32;
+                        
+                        entities.push(EntitySpan {
+                            text: entity_text,
+                            label: entity_type,
+                            start: start_offset.0,
+                            end: end_offset.1,
+                            token_start: start_idx,
+                            token_end: end_idx,
+                            confidence: avg_confidence,
+                        });
+                    }
+                }
+                
+                let entity_type = label[2..].to_string();
+                current_entity = Some((entity_type, i, i + 1, vec![confidence]));
+            } else if label.starts_with("I-") {
+                // Continue entity
+                if let Some((ref mut entity_type, _, ref mut end_idx, ref mut confidences)) = current_entity {
+                    let new_type = label[2..].to_string();
+                    if *entity_type == new_type {
+                        *end_idx = i + 1;
+                        confidences.push(confidence);
+                    } else {
+                        // Type mismatch, start new entity
+                        current_entity = Some((new_type, i, i + 1, vec![confidence]));
+                    }
+                } else {
+                    // I- tag without B- tag, treat as beginning
+                    let entity_type = label[2..].to_string();
+                    current_entity = Some((entity_type, i, i + 1, vec![confidence]));
+                }
+            }
+        }
+        
+        // Handle final entity
+        if let Some((entity_type, start_idx, end_idx, confidences)) = current_entity {
+            if let (Some(start_offset), Some(end_offset)) = 
+                (offsets.get(start_idx), offsets.get(end_idx - 1)) {
+                let entity_text = text[start_offset.0..end_offset.1].to_string();
+                let avg_confidence = confidences.iter().sum::<f32>() / confidences.len() as f32;
+                
+                entities.push(EntitySpan {
+                    text: entity_text,
+                    label: entity_type,
+                    start: start_offset.0,
+                    end: end_offset.1,
+                    token_start: start_idx,
+                    token_end: end_idx,
+                    confidence: avg_confidence,
+                });
+            }
+        }
+        
+        Ok(entities)
+    }
+    
+    /// Get the label configuration
+    pub fn labels(&self) -> Result<RHash> {
+        let hash = RHash::new();
+        
+        let id2label = RHash::new();
+        for (id, label) in &self.config.id2label {
+            id2label.aset(*id, label.as_str())?;
+        }
+        
+        let label2id = RHash::new();
+        for (label, id) in &self.config.label2id {
+            label2id.aset(label.as_str(), *id)?;
+        }
+        
+        hash.aset("id2label", id2label)?;
+        hash.aset("label2id", label2id)?;
+        hash.aset("num_labels", self.config.id2label.len())?;
+        
+        Ok(hash)
+    }
+    
+    /// Get the tokenizer
+    pub fn tokenizer(&self) -> Result<crate::ruby::tokenizer::Tokenizer> {
+        Ok(crate::ruby::tokenizer::Tokenizer(self.tokenizer.clone()))
+    }
+    
+    /// Get model info
+    pub fn model_info(&self) -> String {
+        format!("NER model: {}, labels: {}", self.model_id, self.config.id2label.len())
+    }
+}
+
+pub fn init(rb_candle: RModule) -> Result<()> {
+    let ner_class = rb_candle.define_class("NER", class::object())?;
+    ner_class.define_singleton_method("new", function!(NER::new, 3))?;
+    ner_class.define_method("extract_entities", method!(NER::extract_entities, 2))?;
+    ner_class.define_method("predict_tokens", method!(NER::predict_tokens, 1))?;
+    ner_class.define_method("labels", method!(NER::labels, 0))?;
+    ner_class.define_method("tokenizer", method!(NER::tokenizer, 0))?;
+    ner_class.define_method("model_info", method!(NER::model_info, 0))?;
+    
+    Ok(())
+}

data/ext/candle/src/reranker.rs

@@ -3,28 +3,29 @@ use candle_transformers::models::bert::{BertModel, Config};
 use candle_core::{Device as CoreDevice, Tensor, IndexOp, DType};
 use candle_nn::{VarBuilder, Linear, Module, ops::sigmoid};
 use hf_hub::{api::sync::Api, Repo, RepoType};
-use tokenizers::{PaddingParams, Tokenizer, EncodeInput};
+use tokenizers::{EncodeInput, Tokenizer};
 use std::thread;
-use crate::ruby::{Device as RbDevice, Result as RbResult};
+use crate::ruby::{Device, Result};
+use crate::tokenizer::{TokenizerWrapper, loader::TokenizerLoader};
 
 #[magnus::wrap(class = "Candle::Reranker", free_immediately, size)]
 pub struct Reranker {
     model: BertModel,
-    tokenizer: Tokenizer,
+    tokenizer: TokenizerWrapper,
     pooler: Linear,
     classifier: Linear,
     device: CoreDevice,
 }
 
 impl Reranker {
-    pub fn new(model_id: String, device: Option<RbDevice>) -> RbResult<Self> {
-        let device = device.unwrap_or(RbDevice::Cpu).as_device()?;
+    pub fn new(model_id: String, device: Option<Device>) -> Result<Self> {
+        let device = device.unwrap_or(Device::Cpu).as_device()?;
         Self::new_with_core_device(model_id, device)
     }
         
-    fn new_with_core_device(model_id: String, device: CoreDevice) -> Result<Self, Error> {
+    fn new_with_core_device(model_id: String, device: CoreDevice) -> std::result::Result<Self, Error> {
         let device_clone = device.clone();
-        let handle = thread::spawn(move || -> Result<(BertModel, Tokenizer, Linear, Linear), Box<dyn std::error::Error + Send + Sync>> {
+        let handle = thread::spawn(move || -> std::result::Result<(BertModel, TokenizerWrapper, Linear, Linear), Box<dyn std::error::Error + Send + Sync>> {
             let api = Api::new()?;
             let repo = api.repo(Repo::new(model_id.clone(), RepoType::Model));
             
@@ -38,12 +39,8 @@ impl Reranker {
             let config: Config = serde_json::from_str(&config)?;
 
             // Setup tokenizer with padding
-            let mut tokenizer = Tokenizer::from_file(tokenizer_filename)?;
-            let pp = PaddingParams {
-                strategy: tokenizers::PaddingStrategy::BatchLongest,
-                ..Default::default()
-            };
-            tokenizer.with_padding(Some(pp));
+            let tokenizer = Tokenizer::from_file(tokenizer_filename)?;
+            let tokenizer = TokenizerLoader::with_padding(tokenizer, None);
             
             // Load model weights
             let vb = unsafe {
@@ -59,7 +56,7 @@ impl Reranker {
             // Load classifier layer for cross-encoder (single output score)
             let classifier = candle_nn::linear(config.hidden_size, 1, vb.pp("classifier"))?;
             
-            Ok((model, tokenizer, pooler, classifier))
+            Ok((model, TokenizerWrapper::new(tokenizer), pooler, classifier))
         });
         
         match handle.join() {
@@ -71,12 +68,12 @@ impl Reranker {
         }
     }
     
-    pub fn debug_tokenization(&self, query: String, document: String) -> Result<magnus::RHash, Error> {
+    pub fn debug_tokenization(&self, query: String, document: String) -> std::result::Result<magnus::RHash, Error> {
         // Create query-document pair for cross-encoder
         let query_doc_pair: EncodeInput = (query.clone(), document.clone()).into();
         
-        // Tokenize
-        let encoding = self.tokenizer.encode(query_doc_pair, true)
+        // Tokenize using the inner tokenizer for detailed info
+        let encoding = self.tokenizer.inner().encode(query_doc_pair, true)
             .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Tokenization failed: {}", e)))?;
         
         // Get token information
@@ -95,7 +92,7 @@ impl Reranker {
         Ok(result)
     }
     
-    pub fn rerank_with_options(&self, query: String, documents: RArray, pooling_method: String, apply_sigmoid: bool) -> Result<RArray, Error> {
+    pub fn rerank_with_options(&self, query: String, documents: RArray, pooling_method: String, apply_sigmoid: bool) -> std::result::Result<RArray, Error> {
         let documents: Vec<String> = documents.to_vec()?;
         
         // Create query-document pairs for cross-encoder
@@ -104,8 +101,8 @@ impl Reranker {
             .map(|d| (query.clone(), d.clone()).into())
             .collect();
 
-        // Tokenize batch
-        let encodings = self.tokenizer.encode_batch(query_and_docs, true)
+        // Tokenize batch using inner tokenizer for access to token type IDs
+        let encodings = self.tokenizer.inner().encode_batch(query_and_docs, true)
             .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Tokenization failed: {}", e)))?;
         
         // Convert to tensors
@@ -256,12 +253,18 @@ impl Reranker {
         }
         Ok(result_array)
     }
+
+    /// Get the tokenizer used by this model
+    pub fn tokenizer(&self) -> std::result::Result<crate::ruby::tokenizer::Tokenizer, Error> {
+        Ok(crate::ruby::tokenizer::Tokenizer(self.tokenizer.clone()))
+    }
 }
 
-pub fn init(rb_candle: RModule) -> Result<(), Error> {
+pub fn init(rb_candle: RModule) -> std::result::Result<(), Error> {
     let c_reranker = rb_candle.define_class("Reranker", class::object())?;
     c_reranker.define_singleton_method("_create", function!(Reranker::new, 2))?;
     c_reranker.define_method("rerank_with_options", method!(Reranker::rerank_with_options, 4))?;
     c_reranker.define_method("debug_tokenization", method!(Reranker::debug_tokenization, 2))?;
+    c_reranker.define_method("tokenizer", method!(Reranker::tokenizer, 0))?;
     Ok(())
 }