red-candle 1.2.1 → 1.2.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 15a070c7424707802e4e82d00ef5532691d76d7627d28ab0a4b5f0ac7522471f
-  data.tar.gz: 112d038f09eda1a6b1751935057bfefbc2355d7b4962c85461de192e3c667980
+  metadata.gz: d101e47090a800dddabfa1ee41971ca3f214f9a422e9408c0970b24c03aab8c7
+  data.tar.gz: 7136f0922548fd4bead034930b3c6d1553c498c23fc17802a6c94f7f4a0c7ce5
 SHA512:
-  metadata.gz: cf43744df320c1d69773dad4713bb41e6fa9bd0f75359d3651aa53e429ac48615705682a61e793487d4a09a1f0d5a4aa28df18b3375de63916d9d1b91b2c98b2
-  data.tar.gz: 81fbfe62ba6135b22b34cfeb8ea99b39b12b3dbb2e4c56ac76a758903554ff9f1911132f0be86e8ca130077d6e44997cfcf61a96d48452c8e13995acf76c7e88
+  metadata.gz: 594900e14d64fc335e53e2599a6e2d81ba4fdb1a1f3d5680536a47538bdf05d12332e959b33b67bfa04cca0dd321835a737d2ddf9cdd33c42b937370dde7a2a1
+  data.tar.gz: 4005d21ddccce9182d8600af5b3fd3f34ee9fabdf1b8235eb37eb1be7ede74b644c23d4ea322f6c0e05808e3e98e0c2066251111f538e9e8a934e7ec069ae4c0

data/README.md

@@ -363,6 +363,12 @@ require 'candle'
 # Initialize the reranker with a cross-encoder model
 reranker = Candle::Reranker.from_pretrained("cross-encoder/ms-marco-MiniLM-L-12-v2")
 
+# Or with custom max_length for truncation (default is 512)
+reranker = Candle::Reranker.from_pretrained(
+  "cross-encoder/ms-marco-MiniLM-L-12-v2",
+  max_length: 256  # Faster processing with less context
+)
+
 # Define your query and candidate documents
 query = "How many people live in London?"
 documents = [
@@ -469,6 +475,75 @@ The reranker uses a BERT-based architecture that:
 
 This joint processing allows cross-encoders to capture subtle semantic relationships between queries and documents, making them more accurate for reranking tasks, though at the cost of higher computational requirements.
 
+### Performance Considerations
+
+**Important**: The Reranker automatically truncates documents to ensure stable performance. The default maximum is 512 tokens, but this is configurable.
+
+#### Configurable Truncation
+
+You can adjust the `max_length` parameter to balance performance and context:
+
+```ruby
+# Default: 512 tokens (maximum context, ~300ms per doc on CPU)
+reranker = Candle::Reranker.from_pretrained(model_id)
+
+# Faster: 256 tokens (~60% faster, ~120ms per doc on CPU)
+reranker = Candle::Reranker.from_pretrained(model_id, max_length: 256)
+
+# Fastest: 128 tokens (~80% faster, ~60ms per doc on CPU)
+reranker = Candle::Reranker.from_pretrained(model_id, max_length: 128)
+```
+
+Choose based on your needs:
+- **512 tokens**: Maximum context for complex queries (default)
+- **256 tokens**: Good balance of speed and context
+- **128 tokens**: Fast processing for simple matching
+
+#### Performance Guidelines
+
+1. **Document Length**: Documents longer than ~400 words will be truncated
+   - The first 512 tokens (roughly 300-400 words) are used
+   - Consider splitting very long documents into chunks if full coverage is needed
+
+2. **Batch Size**: Process multiple documents in one call for efficiency
+   ```ruby
+   # Good: Single call with multiple documents
+   results = reranker.rerank(query, documents)
+   
+   # Less efficient: Multiple calls
+   documents.map { |doc| reranker.rerank(query, [doc]) }
+   ```
+
+3. **Expected Performance**:
+   - **CPU**: ~0.3-0.5s per query-document pair
+   - **GPU (Metal/CUDA)**: ~0.05-0.1s per query-document pair
+   - Performance is consistent regardless of document length due to truncation
+
+4. **Chunking Strategy** for long documents:
+   ```ruby
+   def rerank_long_document(query, long_text, chunk_size: 300)
+     # Split into overlapping chunks
+     words = long_text.split
+     chunks = []
+     
+     (0...words.length).step(chunk_size - 50) do |i|
+       chunk = words[i...(i + chunk_size)].join(" ")
+       chunks << chunk
+     end
+     
+     # Rerank chunks
+     results = reranker.rerank(query, chunks)
+     
+     # Return best chunk
+     results.max_by { |r| r[:score] }
+   end
+   ```
+
+5. **Memory Usage**: 
+   - Model size: ~125MB
+   - Each batch processes all documents simultaneously
+   - Consider batching if you have many documents
+
 ## Tokenizer
 
 Red-Candle provides direct access to tokenizers for text preprocessing and analysis. This is useful for understanding how models process text, debugging issues, and building custom NLP pipelines.

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

@@ -18,46 +18,48 @@ pub struct Reranker {
 }
 
 impl Reranker {
-    pub fn new(model_id: String, device: Option<Device>) -> Result<Self> {
+    pub fn new(model_id: String, device: Option<Device>, max_length: Option<usize>) -> Result<Self> {
         let device = device.unwrap_or(Device::best()).as_device()?;
-        Self::new_with_core_device(model_id, device)
+        let max_length = max_length.unwrap_or(512);  // Default to 512
+        Self::new_with_core_device(model_id, device, max_length)
     }
-        
-    fn new_with_core_device(model_id: String, device: CoreDevice) -> std::result::Result<Self, Error> {
+
+    fn new_with_core_device(model_id: String, device: CoreDevice, max_length: usize) -> std::result::Result<Self, Error> {
         let result = (|| -> 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));
-            
+
             // Download model files
             let config_filename = repo.get("config.json")?;
             let tokenizer_filename = repo.get("tokenizer.json")?;
             let weights_filename = repo.get("model.safetensors")?;
-            
+
             // Load config
             let config = std::fs::read_to_string(config_filename)?;
             let config: Config = serde_json::from_str(&config)?;
 
-            // Setup tokenizer with padding
+            // Setup tokenizer with padding AND truncation
             let tokenizer = Tokenizer::from_file(tokenizer_filename)?;
             let tokenizer = TokenizerLoader::with_padding(tokenizer, None);
-            
+            let tokenizer = TokenizerLoader::with_truncation(tokenizer, max_length);
+
             // Load model weights
             let vb = unsafe {
                 VarBuilder::from_mmaped_safetensors(&[weights_filename], DType::F32, &device)?
             };
-            
+
             // Load BERT model
             let model = BertModel::load(vb.pp("bert"), &config)?;
-            
+
             // Load pooler layer (dense + tanh activation)
             let pooler = candle_nn::linear(config.hidden_size, config.hidden_size, vb.pp("bert.pooler.dense"))?;
-            
+
             // Load classifier layer for cross-encoder (single output score)
             let classifier = candle_nn::linear(config.hidden_size, 1, vb.pp("classifier"))?;
-            
+
             Ok((model, TokenizerWrapper::new(tokenizer), pooler, classifier))
         })();
-        
+
         match result {
             Ok((model, tokenizer, pooler, classifier)) => {
                 Ok(Self { model, tokenizer, pooler, classifier, device, model_id })
@@ -65,18 +67,18 @@ impl Reranker {
             Err(e) => Err(Error::new(magnus::exception::runtime_error(), format!("Failed to load model: {}", e))),
         }
     }
-    
+
     /// Extract CLS embeddings from the model output, handling Metal device workarounds
     fn extract_cls_embeddings(&self, embeddings: &Tensor) -> std::result::Result<Tensor, Error> {
         let cls_embeddings = if self.device.is_metal() {
             // Metal has issues with tensor indexing, use a different approach
             let (batch_size, seq_len, hidden_size) = embeddings.dims3()
                 .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to get dims: {}", e)))?;
-            
+
             // Reshape to [batch * seq_len, hidden] then take first hidden vectors for each batch
             let reshaped = embeddings.reshape((batch_size * seq_len, hidden_size))
                 .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to reshape: {}", e)))?;
-            
+
             // Extract CLS tokens (first token of each sequence)
             let mut cls_vecs = Vec::new();
             for i in 0..batch_size {
@@ -85,7 +87,7 @@ impl Reranker {
                     .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to extract CLS: {}", e)))?;
                 cls_vecs.push(cls_vec);
             }
-            
+
             // Stack the CLS vectors
             Tensor::cat(&cls_vecs, 0)
                 .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to cat CLS tokens: {}", e)))?
@@ -93,39 +95,39 @@ impl Reranker {
             embeddings.i((.., 0))
                 .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to extract CLS token: {}", e)))?
         };
-        
+
         // Ensure tensor is contiguous for downstream operations
         cls_embeddings.contiguous()
             .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to make CLS embeddings contiguous: {}", e)))
     }
-    
+
     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 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
         let token_ids = encoding.get_ids().to_vec();
         let token_type_ids = encoding.get_type_ids().to_vec();
         let attention_mask = encoding.get_attention_mask().to_vec();
         let tokens = encoding.get_tokens().iter().map(|t| t.to_string()).collect::<Vec<_>>();
-        
+
         // Create result hash
         let result = magnus::RHash::new();
         result.aset("token_ids", RArray::from_vec(token_ids.iter().map(|&id| id as i64).collect::<Vec<_>>()))?;
         result.aset("token_type_ids", RArray::from_vec(token_type_ids.iter().map(|&id| id as i64).collect::<Vec<_>>()))?;
         result.aset("attention_mask", RArray::from_vec(attention_mask.iter().map(|&mask| mask as i64).collect::<Vec<_>>()))?;
         result.aset("tokens", RArray::from_vec(tokens))?;
-        
+
         Ok(result)
     }
-    
+
     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
         let query_and_docs: Vec<EncodeInput> = documents
             .iter()
@@ -135,13 +137,13 @@ impl Reranker {
         // 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
         let token_ids = encodings
             .iter()
             .map(|e| e.get_ids().to_vec())
             .collect::<Vec<_>>();
-            
+
         let token_type_ids = encodings
             .iter()
             .map(|e| e.get_type_ids().to_vec())
@@ -153,11 +155,11 @@ impl Reranker {
             .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to create token type ids tensor: {}", e)))?;
         let attention_mask = token_ids.ne(0u32)
             .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to create attention mask: {}", e)))?;
-        
+
         // Forward pass through BERT
         let embeddings = self.model.forward(&token_ids, &token_type_ids, Some(&attention_mask))
             .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Model forward pass failed: {}", e)))?;
-        
+
         // Apply pooling based on the specified method
         let pooled_embeddings = match pooling_method.as_str() {
             "pooler" => {
@@ -181,10 +183,10 @@ impl Reranker {
                 (sum / (seq_len as f64))
                     .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to compute mean: {}", e)))?
             },
-            _ => return Err(Error::new(magnus::exception::runtime_error(), 
+            _ => return Err(Error::new(magnus::exception::runtime_error(),
                 format!("Unknown pooling method: {}. Use 'pooler', 'cls', or 'mean'", pooling_method)))
         };
-        
+
         // Apply classifier to get relevance scores (raw logits)
         // Ensure tensor is contiguous before linear layer
         let pooled_embeddings = pooled_embeddings.contiguous()
@@ -193,7 +195,7 @@ impl Reranker {
             .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Classifier forward failed: {}", e)))?;
         let scores = logits.squeeze(1)
             .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to squeeze tensor: {}", e)))?;
-        
+
         // Optionally apply sigmoid activation
         let scores = if apply_sigmoid {
             sigmoid(&scores)
@@ -201,7 +203,7 @@ impl Reranker {
         } else {
             scores
         };
-        
+
         let scores_vec: Vec<f32> = scores.to_vec1()
             .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to convert scores to vec: {}", e)))?;
 
@@ -212,7 +214,7 @@ impl Reranker {
             .enumerate()
             .map(|(idx, (doc, score))| (doc, score, idx))
             .collect();
-        
+
         // Sort documents by relevance score (descending)
         ranked_docs.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
 
@@ -232,17 +234,17 @@ impl Reranker {
     pub fn tokenizer(&self) -> std::result::Result<crate::ruby::tokenizer::Tokenizer, Error> {
         Ok(crate::ruby::tokenizer::Tokenizer(self.tokenizer.clone()))
     }
-    
+
     /// Get the model_id
     pub fn model_id(&self) -> String {
         self.model_id.clone()
     }
-    
+
     /// Get the device
     pub fn device(&self) -> Device {
         Device::from_device(&self.device)
     }
-    
+
     /// Get all options as a hash
     pub fn options(&self) -> std::result::Result<magnus::RHash, Error> {
         let hash = magnus::RHash::new();
@@ -254,7 +256,7 @@ impl Reranker {
 
 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_singleton_method("_create", function!(Reranker::new, 3))?;
     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))?;

data/lib/candle/reranker.rb

@@ -6,18 +6,20 @@ module Candle
     # Load a pre-trained reranker model from HuggingFace
     # @param model_id [String] HuggingFace model ID (defaults to cross-encoder/ms-marco-MiniLM-L-12-v2)
     # @param device [Candle::Device] The device to use for computation (defaults to best available)
+    # @param max_length [Integer] Maximum sequence length for truncation (defaults to 512)
     # @return [Reranker] A new Reranker instance
-    def self.from_pretrained(model_id = DEFAULT_MODEL_PATH, device: Candle::Device.best)
-      _create(model_id, device)
+    def self.from_pretrained(model_id = DEFAULT_MODEL_PATH, device: Candle::Device.best, max_length: 512)
+      _create(model_id, device, max_length)
     end
     
     # Constructor for creating a new Reranker with optional parameters
     # @deprecated Use {.from_pretrained} instead
     # @param model_path [String, nil] The path to the model on Hugging Face
     # @param device [Candle::Device, Candle::Device.cpu] The device to use for computation
-    def self.new(model_path: DEFAULT_MODEL_PATH, device: Candle::Device.best)
+    # @param max_length [Integer] Maximum sequence length for truncation (defaults to 512)
+    def self.new(model_path: DEFAULT_MODEL_PATH, device: Candle::Device.best, max_length: 512)
       $stderr.puts "[DEPRECATION] `Reranker.new` is deprecated. Please use `Reranker.from_pretrained` instead."
-      _create(model_path, device)
+      _create(model_path, device, max_length)
     end
 
     # Returns documents ranked by relevance using the specified pooling method.

data/lib/candle/version.rb

@@ -1,5 +1,5 @@
 # :nocov:
 module Candle
-  VERSION = "1.2.1"
+  VERSION = "1.2.2"
 end
 # :nocov:

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: red-candle
 version: !ruby/object:Gem::Version
-  version: 1.2.1
+  version: 1.2.2
 platform: ruby
 authors:
 - Christopher Petersen
@@ -9,7 +9,7 @@ authors:
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2025-08-11 00:00:00.000000000 Z
+date: 2025-08-13 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rb_sys