red-candle 1.0.0.pre.1 → 1.0.0.pre.3

data/ext/candle/src/reranker.rs

@@ -0,0 +1,267 @@
+use magnus::{class, function, method, prelude::*, Error, RModule, Float, RArray};
+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 std::thread;
+use crate::ruby::{Device as RbDevice, Result as RbResult};
+
+#[magnus::wrap(class = "Candle::Reranker", free_immediately, size)]
+pub struct Reranker {
+    model: BertModel,
+    tokenizer: Tokenizer,
+    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()?;
+        Self::new_with_core_device(model_id, device)
+    }
+        
+    fn new_with_core_device(model_id: String, device: CoreDevice) -> 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 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
+            let mut tokenizer = Tokenizer::from_file(tokenizer_filename)?;
+            let pp = PaddingParams {
+                strategy: tokenizers::PaddingStrategy::BatchLongest,
+                ..Default::default()
+            };
+            tokenizer.with_padding(Some(pp));
+            
+            // 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"), &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, tokenizer, pooler, classifier))
+        });
+        
+        match handle.join() {
+            Ok(Ok((model, tokenizer, pooler, classifier))) => {
+                Ok(Self { model, tokenizer, pooler, classifier, device })
+            }
+            Ok(Err(e)) => Err(Error::new(magnus::exception::runtime_error(), format!("Failed to load model: {}", e))),
+            Err(_) => Err(Error::new(magnus::exception::runtime_error(), "Thread panicked while loading model")),
+        }
+    }
+    
+    pub fn debug_tokenization(&self, query: String, document: String) -> 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)
+            .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) -> 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()
+            .map(|d| (query.clone(), d.clone()).into())
+            .collect();
+
+        // Tokenize batch
+        let encodings = self.tokenizer.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())
+            .collect::<Vec<_>>();
+
+        let token_ids = Tensor::new(token_ids, &self.device)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to create tensor: {}", e)))?;
+        let token_type_ids = Tensor::new(token_type_ids, &self.device)
+            .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" => {
+                // Extract [CLS] token and apply pooler (dense + tanh)
+                // Work around Metal indexing issue by using narrow instead of i((.., 0))
+                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 {
+                        let start_idx = i * _seq_len;
+                        let cls_vec = reshaped.narrow(0, start_idx, 1)
+                            .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)))?
+                        .contiguous()
+                        .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to make contiguous: {}", e)))?
+                } else {
+                    embeddings.i((.., 0))
+                        .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to extract CLS token: {}", e)))?
+                };
+                // Ensure tensor is contiguous before linear layer
+                let cls_embeddings = cls_embeddings.contiguous()
+                    .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to make cls_embeddings contiguous: {}", e)))?;
+                let pooled = self.pooler.forward(&cls_embeddings)
+                    .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Pooler forward failed: {}", e)))?;
+                pooled.tanh()
+                    .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Tanh activation failed: {}", e)))?
+            },
+            "cls" => {
+                // Just use the [CLS] token embeddings directly (no pooler layer)
+                // Work around Metal indexing issue
+                let cls_embeddings = if self.device.is_metal() {
+                    // Use same approach as pooler method
+                    let (batch_size, _seq_len, hidden_size) = embeddings.dims3()
+                        .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to get dims: {}", e)))?;
+                    
+                    let reshaped = embeddings.reshape((batch_size * _seq_len, hidden_size))
+                        .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to reshape: {}", e)))?;
+                    
+                    let mut cls_vecs = Vec::new();
+                    for i in 0..batch_size {
+                        let start_idx = i * _seq_len;
+                        let cls_vec = reshaped.narrow(0, start_idx, 1)
+                            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to extract CLS: {}", e)))?;
+                        cls_vecs.push(cls_vec);
+                    }
+                    
+                    Tensor::cat(&cls_vecs, 0)
+                        .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to cat CLS tokens: {}", e)))?
+                        .contiguous()
+                        .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to make contiguous: {}", e)))?
+                } else {
+                    embeddings.i((.., 0))
+                        .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to extract CLS token: {}", e)))?
+                };
+                // Ensure contiguous for classifier
+                cls_embeddings.contiguous()
+                    .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to make CLS embeddings contiguous: {}", e)))?
+            },
+            "mean" => {
+                // Mean pooling across all tokens
+                let (_batch, seq_len, _hidden) = embeddings.dims3()
+                    .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to get tensor dimensions: {}", e)))?;
+                let sum = embeddings.sum(1)
+                    .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to sum embeddings: {}", e)))?;
+                (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(), 
+                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()
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to make pooled_embeddings contiguous: {}", e)))?;
+        let logits = self.classifier.forward(&pooled_embeddings)
+            .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)
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Sigmoid failed: {}", e)))?
+        } 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)))?;
+
+        // Create tuples with document, score, and original index
+        let mut ranked_docs: Vec<(String, f32, usize)> = documents
+            .into_iter()
+            .zip(scores_vec)
+            .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));
+
+        // Build result array with [doc, score, doc_id]
+        let result_array = RArray::new();
+        for (doc, score, doc_id) in ranked_docs {
+            let tuple = RArray::new();
+            tuple.push(doc)?;
+            tuple.push(Float::from_f64(score as f64))?;
+            tuple.push(doc_id)?;
+            result_array.push(tuple)?;
+        }
+        Ok(result_array)
+    }
+}
+
+pub fn init(rb_candle: RModule) -> 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))?;
+    Ok(())
+}

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

@@ -0,0 +1,197 @@
+use magnus::Error;
+use magnus::{function, method, class, RModule, Module, Object};
+
+use ::candle_core::Device as CoreDevice;
+use crate::ruby::Result as RbResult;
+
+#[cfg(any(feature = "cuda", feature = "metal"))]
+use crate::ruby::errors::wrap_candle_err;
+
+#[cfg(feature = "cuda")]
+static CUDA_DEVICE: std::sync::Mutex<Option<CoreDevice>> = std::sync::Mutex::new(None);
+
+#[cfg(feature = "metal")]
+static METAL_DEVICE: std::sync::Mutex<Option<CoreDevice>> = std::sync::Mutex::new(None);
+
+/// Get list of available devices based on compile-time features
+pub fn available_devices() -> Vec<String> {
+    let devices = vec!["cpu".to_string()];
+    
+    #[cfg(all(feature = "cuda", not(force_cpu)))]
+    let devices = {
+        let mut devices = devices;
+        devices.push("cuda".to_string());
+        devices
+    };
+    
+    #[cfg(all(feature = "metal", not(force_cpu)))]
+    let devices = {
+        let mut devices = devices;
+        devices.push("metal".to_string());
+        devices
+    };
+    
+    devices
+}
+
+/// Get the default device based on what's available
+pub fn default_device() -> Device {
+    // Return based on compiled features, not detection
+    #[cfg(all(feature = "metal", not(force_cpu)))]
+    {
+        Device::Metal
+    }
+    
+    #[cfg(all(feature = "cuda", not(feature = "metal"), not(force_cpu)))]
+    {
+        Device::Cuda
+    }
+    
+    #[cfg(not(any(all(feature = "metal", not(force_cpu)), all(feature = "cuda", not(feature = "metal"), not(force_cpu)))))]
+    {
+        Device::Cpu
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[magnus::wrap(class = "Candle::Device")]
+pub enum Device {
+    Cpu,
+    Cuda,
+    Metal,
+}
+
+impl Device {
+    /// Create a CPU device
+    pub fn cpu() -> Self {
+        Self::Cpu
+    }
+
+    /// Create a CUDA device (GPU)
+    pub fn cuda() -> RbResult<Self> {
+        #[cfg(not(feature = "cuda"))]
+        {
+            return Err(Error::new(
+                magnus::exception::runtime_error(),
+                "CUDA support not compiled in. Rebuild with CUDA available.",
+            ));
+        }
+        
+        #[cfg(feature = "cuda")]
+        Ok(Self::Cuda)
+    }
+
+    /// Create a Metal device (Apple GPU)
+    pub fn metal() -> RbResult<Self> {
+        #[cfg(not(feature = "metal"))]
+        {
+            return Err(Error::new(
+                magnus::exception::runtime_error(),
+                "Metal support not compiled in. Rebuild on macOS.",
+            ));
+        }
+        
+        #[cfg(feature = "metal")]
+        Ok(Self::Metal)
+    }
+
+    pub fn from_device(device: &CoreDevice) -> Self {
+        match device {
+            CoreDevice::Cpu => Self::Cpu,
+            CoreDevice::Cuda(_) => Self::Cuda,
+            CoreDevice::Metal(_) => Self::Metal,
+        }
+    }
+
+    pub fn as_device(&self) -> RbResult<CoreDevice> {
+        match self {
+            Self::Cpu => Ok(CoreDevice::Cpu),
+            Self::Cuda => {
+                #[cfg(not(feature = "cuda"))]
+                {
+                    return Err(Error::new(
+                        magnus::exception::runtime_error(),
+                        "CUDA support not compiled in. Rebuild with CUDA available.",
+                    ));
+                }
+                
+                #[cfg(feature = "cuda")]
+                {
+                    let mut device = CUDA_DEVICE.lock().unwrap();
+                    if let Some(device) = device.as_ref() {
+                        return Ok(device.clone());
+                    };
+                    // Note: new_cuda() is used here (not cuda_if_available) because
+                    // we want to fail if CUDA isn't available at runtime, not fall back to CPU
+                    let d = CoreDevice::new_cuda(0).map_err(wrap_candle_err)?;
+                    *device = Some(d.clone());
+                    Ok(d)
+                }
+            }
+            Self::Metal => {
+                #[cfg(not(feature = "metal"))]
+                {
+                    return Err(Error::new(
+                        magnus::exception::runtime_error(),
+                        "Metal support not compiled in. Rebuild on macOS.",
+                    ));
+                }
+                
+                #[cfg(feature = "metal")]
+                {
+                    let mut device = METAL_DEVICE.lock().unwrap();
+                    if let Some(device) = device.as_ref() {
+                        return Ok(device.clone());
+                    };
+                    let d = CoreDevice::new_metal(0).map_err(wrap_candle_err)?;
+                    *device = Some(d.clone());
+                    Ok(d)
+                }
+            }
+        }
+    }
+
+    pub fn __repr__(&self) -> String {
+        match self {
+            Self::Cpu => "cpu".to_string(),
+            Self::Cuda => "cuda".to_string(),
+            Self::Metal => "metal".to_string(),
+        }
+    }
+
+    pub fn __str__(&self) -> String {
+        self.__repr__()
+    }
+}
+
+impl magnus::TryConvert for Device {
+    fn try_convert(val: magnus::Value) -> RbResult<Self> {
+        // First check if it's already a wrapped Device object
+        if let Ok(device) = <magnus::typed_data::Obj<Device> as magnus::TryConvert>::try_convert(val) {
+            return Ok(*device);
+        }
+
+        // Otherwise try to convert from string
+        let device = magnus::RString::try_convert(val)?;
+        let device = unsafe { device.as_str() }.unwrap();
+        let device = match device {
+            "cpu" => Device::Cpu,
+            "cuda" => Device::Cuda,
+            "metal" => Device::Metal,
+            _ => return Err(Error::new(magnus::exception::arg_error(), "invalid device")),
+        };
+        Ok(device)
+    }
+}
+
+pub fn init(rb_candle: RModule) -> Result<(), Error> {
+    let rb_device = rb_candle.define_class("Device", class::object())?;
+    rb_device.define_singleton_method("cpu", function!(Device::cpu, 0))?;
+    rb_device.define_singleton_method("cuda", function!(Device::cuda, 0))?;
+    rb_device.define_singleton_method("metal", function!(Device::metal, 0))?;
+    rb_device.define_singleton_method("available_devices", function!(available_devices, 0))?;
+    rb_device.define_singleton_method("default", function!(default_device, 0))?;
+    rb_device.define_method("to_s", method!(Device::__str__, 0))?;
+    rb_device.define_method("inspect", method!(Device::__repr__, 0))?;
+    Ok(())
+}

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

@@ -0,0 +1,37 @@
+use magnus::value::ReprValue;
+use magnus::{method, class, RModule, Error, Module};
+
+use ::candle_core::DType as CoreDType;
+use crate::ruby::Result as RbResult;
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[magnus::wrap(class = "Candle::DType", free_immediately, size)]
+
+/// A `candle` dtype.
+pub struct DType(pub CoreDType);
+
+impl DType {
+    pub fn __repr__(&self) -> String {
+        format!("{:?}", self.0)
+    }
+
+    pub fn __str__(&self) -> String {
+        self.__repr__()
+    }
+}
+
+impl DType {
+    pub fn from_rbobject(dtype: magnus::Symbol) -> RbResult<Self> {
+        let dtype = unsafe { dtype.to_s() }.unwrap().into_owned();
+        use std::str::FromStr;
+        let dtype = CoreDType::from_str(&dtype).unwrap();
+        Ok(Self(dtype))
+    }
+}
+
+pub fn init(rb_candle: RModule) -> Result<(), Error> {
+    let rb_dtype = rb_candle.define_class("DType", class::object())?;
+    rb_dtype.define_method("to_s", method!(DType::__str__, 0))?;
+    rb_dtype.define_method("inspect", method!(DType::__repr__, 0))?;
+    Ok(())
+}