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

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

@@ -0,0 +1,295 @@
+use magnus::{function, method, prelude::*, Error, Module, RArray, RHash, RModule, Ruby, TryConvert, Value};
+use std::cell::RefCell;
+
+use crate::llm::{GenerationConfig as RustGenerationConfig, TextGenerator, mistral::Mistral as RustMistral};
+use crate::ruby::{Result as RbResult, Device as RbDevice};
+
+// Use an enum to handle different model types instead of trait objects
+#[derive(Debug)]
+enum ModelType {
+    Mistral(RustMistral),
+}
+
+impl ModelType {
+    fn generate(&mut self, prompt: &str, config: &RustGenerationConfig) -> candle_core::Result<String> {
+        match self {
+            ModelType::Mistral(m) => m.generate(prompt, config),
+        }
+    }
+
+    fn generate_stream(
+        &mut self,
+        prompt: &str,
+        config: &RustGenerationConfig,
+        callback: impl FnMut(&str),
+    ) -> candle_core::Result<String> {
+        match self {
+            ModelType::Mistral(m) => m.generate_stream(prompt, config, callback),
+        }
+    }
+
+    #[allow(dead_code)]
+    fn model_name(&self) -> &str {
+        match self {
+            ModelType::Mistral(m) => m.model_name(),
+        }
+    }
+    
+    fn clear_cache(&mut self) {
+        match self {
+            ModelType::Mistral(m) => m.clear_cache(),
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+#[magnus::wrap(class = "Candle::GenerationConfig", mark, free_immediately)]
+pub struct GenerationConfig {
+    inner: RustGenerationConfig,
+}
+
+impl GenerationConfig {
+    pub fn new(kwargs: RHash) -> RbResult<Self> {
+        let mut config = RustGenerationConfig::default();
+        
+        // Extract values from kwargs manually
+        if let Some(value) = kwargs.get(magnus::Symbol::new("max_length")) {
+            if let Ok(v) = TryConvert::try_convert(value) {
+                config.max_length = v;
+            }
+        }
+        
+        if let Some(value) = kwargs.get(magnus::Symbol::new("temperature")) {
+            if let Ok(v) = TryConvert::try_convert(value) {
+                config.temperature = v;
+            }
+        }
+        
+        if let Some(value) = kwargs.get(magnus::Symbol::new("top_p")) {
+            if let Ok(v) = TryConvert::try_convert(value) {
+                config.top_p = Some(v);
+            }
+        }
+        
+        if let Some(value) = kwargs.get(magnus::Symbol::new("top_k")) {
+            if let Ok(v) = TryConvert::try_convert(value) {
+                config.top_k = Some(v);
+            }
+        }
+        
+        if let Some(value) = kwargs.get(magnus::Symbol::new("repetition_penalty")) {
+            if let Ok(v) = TryConvert::try_convert(value) {
+                config.repetition_penalty = v;
+            }
+        }
+        
+        if let Some(value) = kwargs.get(magnus::Symbol::new("repetition_penalty_last_n")) {
+            if let Ok(v) = TryConvert::try_convert(value) {
+                config.repetition_penalty_last_n = v;
+            }
+        }
+        
+        if let Some(value) = kwargs.get(magnus::Symbol::new("seed")) {
+            if let Ok(v) = TryConvert::try_convert(value) {
+                config.seed = v;
+            }
+        }
+        
+        if let Some(value) = kwargs.get(magnus::Symbol::new("include_prompt")) {
+            if let Ok(v) = TryConvert::try_convert(value) {
+                config.include_prompt = v;
+            }
+        }
+        
+        if let Some(value) = kwargs.get(magnus::Symbol::new("stop_sequences")) {
+            if let Ok(arr) = <RArray as TryConvert>::try_convert(value) {
+                config.stop_sequences = arr
+                    .into_iter()
+                    .filter_map(|v| <String as TryConvert>::try_convert(v).ok())
+                    .collect();
+            }
+        }
+        
+        Ok(Self { inner: config })
+    }
+
+    pub fn default() -> Self {
+        Self {
+            inner: RustGenerationConfig::default(),
+        }
+    }
+
+    // Getters
+    pub fn max_length(&self) -> usize {
+        self.inner.max_length
+    }
+
+    pub fn temperature(&self) -> f64 {
+        self.inner.temperature
+    }
+
+    pub fn top_p(&self) -> Option<f64> {
+        self.inner.top_p
+    }
+
+    pub fn top_k(&self) -> Option<usize> {
+        self.inner.top_k
+    }
+
+    pub fn repetition_penalty(&self) -> f32 {
+        self.inner.repetition_penalty
+    }
+
+    pub fn seed(&self) -> u64 {
+        self.inner.seed
+    }
+
+    pub fn stop_sequences(&self) -> Vec<String> {
+        self.inner.stop_sequences.clone()
+    }
+
+    pub fn include_prompt(&self) -> bool {
+        self.inner.include_prompt
+    }
+}
+
+#[derive(Clone, Debug)]
+#[magnus::wrap(class = "Candle::LLM", mark, free_immediately)]
+pub struct LLM {
+    model: std::sync::Arc<std::sync::Mutex<RefCell<ModelType>>>,
+    model_id: String,
+    device: RbDevice,
+}
+
+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);
+        let candle_device = device.as_device()?;
+        
+        // For now, we'll use tokio runtime directly
+        // In production, you might want to share a runtime
+        let rt = tokio::runtime::Runtime::new()
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to create runtime: {}", e)))?;
+        
+        // Determine model type from ID and load appropriately
+        let model_lower = model_id.to_lowercase();
+        let model = if model_lower.contains("mistral") {
+            let mistral = rt.block_on(async {
+                RustMistral::from_pretrained(&model_id, candle_device).await
+            })
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to load model: {}", e)))?;
+            ModelType::Mistral(mistral)
+        } else {
+            return Err(Error::new(
+                magnus::exception::runtime_error(),
+                format!("Unsupported model type: {}. Currently only Mistral models are supported.", model_id),
+            ));
+        };
+        
+        Ok(Self {
+            model: std::sync::Arc::new(std::sync::Mutex::new(RefCell::new(model))),
+            model_id,
+            device,
+        })
+    }
+
+    /// Generate text from a prompt
+    pub fn generate(&self, prompt: String, config: Option<&GenerationConfig>) -> RbResult<String> {
+        let config = config
+            .map(|c| c.inner.clone())
+            .unwrap_or_default();
+        
+        let model = self.model.lock().unwrap();
+        let mut model_ref = model.borrow_mut();
+        
+        model_ref.generate(&prompt, &config)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Generation failed: {}", e)))
+    }
+
+    /// Generate text with streaming output
+    pub fn generate_stream(&self, prompt: String, config: Option<&GenerationConfig>) -> RbResult<String> {
+        let config = config
+            .map(|c| c.inner.clone())
+            .unwrap_or_default();
+        
+        let ruby = Ruby::get().unwrap();
+        let block = ruby.block_proc();
+        if let Err(_) = block {
+            return Err(Error::new(magnus::exception::runtime_error(), "No block given"));
+        }
+        let block = block.unwrap();
+        
+        let model = self.model.lock().unwrap();
+        let mut model_ref = model.borrow_mut();
+        
+        let result = model_ref.generate_stream(&prompt, &config, |token| {
+            // Call the Ruby block with each token
+            let _ = block.call::<(String,), Value>((token.to_string(),));
+        });
+        
+        result.map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Generation failed: {}", e)))
+    }
+
+    /// Get the model name
+    pub fn model_name(&self) -> String {
+        self.model_id.clone()
+    }
+
+    /// Get the device the model is running on
+    pub fn device(&self) -> RbDevice {
+        self.device
+    }
+    
+    /// Clear the model's cache (e.g., KV cache for transformers)
+    pub fn clear_cache(&self) -> RbResult<()> {
+        let model = self.model.lock().unwrap();
+        let mut model_ref = model.borrow_mut();
+        model_ref.clear_cache();
+        Ok(())
+    }
+}
+
+// Define a standalone function for from_pretrained that handles variable arguments
+fn from_pretrained_wrapper(args: &[Value]) -> RbResult<LLM> {
+    match args.len() {
+        1 => {
+            let model_id: String = TryConvert::try_convert(args[0])?;
+            LLM::from_pretrained(model_id, None)
+        },
+        2 => {
+            let model_id: String = TryConvert::try_convert(args[0])?;
+            let device: RbDevice = TryConvert::try_convert(args[1])?;
+            LLM::from_pretrained(model_id, Some(device))
+        },
+        _ => Err(Error::new(
+            magnus::exception::arg_error(),
+            "wrong number of arguments (expected 1..2)"
+        ))
+    }
+}
+
+pub fn init_llm(rb_candle: RModule) -> RbResult<()> {
+    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))?;
+    
+    rb_generation_config.define_method("max_length", method!(GenerationConfig::max_length, 0))?;
+    rb_generation_config.define_method("temperature", method!(GenerationConfig::temperature, 0))?;
+    rb_generation_config.define_method("top_p", method!(GenerationConfig::top_p, 0))?;
+    rb_generation_config.define_method("top_k", method!(GenerationConfig::top_k, 0))?;
+    rb_generation_config.define_method("repetition_penalty", method!(GenerationConfig::repetition_penalty, 0))?;
+    rb_generation_config.define_method("seed", method!(GenerationConfig::seed, 0))?;
+    rb_generation_config.define_method("stop_sequences", method!(GenerationConfig::stop_sequences, 0))?;
+    rb_generation_config.define_method("include_prompt", method!(GenerationConfig::include_prompt, 0))?;
+    
+    let rb_llm = rb_candle.define_class("LLM", magnus::class::object())?;
+    rb_llm.define_singleton_method("_from_pretrained", function!(from_pretrained_wrapper, -1))?;
+    rb_llm.define_method("_generate", method!(LLM::generate, 2))?;
+    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("clear_cache", method!(LLM::clear_cache, 0))?;
+    
+    Ok(())
+}

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

@@ -0,0 +1,21 @@
+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 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
+pub use embedding_model::init as init_embedding_model;
+pub use utils::candle_utils;
+pub use llm::init_llm;

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

@@ -0,0 +1,69 @@
+use std::sync::Arc;
+use magnus::{method, class, RModule, Error, Module};
+
+use crate::ruby::errors::wrap_candle_err;
+use crate::ruby::{Tensor, Result as RbResult};
+use ::candle_core::{quantized::QTensor as CoreQTensor, Device as CoreDevice};
+
+#[derive(Debug)]
+#[magnus::wrap(class = "Candle::QTensor", free_immediately, size)]
+/// A quantized tensor.
+pub struct QTensor(Arc<CoreQTensor>);
+
+impl std::ops::Deref for QTensor {
+    type Target = CoreQTensor;
+
+    fn deref(&self) -> &Self::Target {
+        self.0.as_ref()
+    }
+}
+
+impl QTensor {
+    ///Gets the tensors quantized dtype.
+    /// &RETURNS&: str
+    pub fn ggml_dtype(&self) -> String {
+        format!("{:?}", self.0.dtype())
+    }
+
+    ///Gets the rank of the tensor.
+    /// &RETURNS&: int
+    pub fn rank(&self) -> usize {
+        self.0.rank()
+    }
+
+    ///Gets the shape of the tensor.
+    /// &RETURNS&: Tuple[int]
+    pub fn shape(&self) -> Vec<usize> {
+        self.0.shape().dims().to_vec()
+    }
+
+    pub fn __repr__(&self) -> String {
+        format!("{:?}", self.0)
+    }
+
+    pub fn __str__(&self) -> String {
+        self.__repr__()
+    }
+
+    /// Dequantizes the tensor.
+    /// &RETURNS&: Tensor
+    pub fn dequantize(&self) -> RbResult<Tensor> {
+        let tensor = self.0.dequantize(&CoreDevice::Cpu).map_err(wrap_candle_err)?;
+        Ok(Tensor(tensor))
+    }
+
+    // fn matmul_t(&self, lhs: &Tensor) -> RbResult<Tensor> {
+    //     let qmatmul = ::candle_core::quantized::QMatMul::from_arc(self.0.clone());
+    //     let res = qmatmul.forward(lhs).map_err(wrap_candle_err)?;
+    //     Ok(Tensor(res))
+    // }
+}
+
+pub fn init(rb_candle: RModule) -> Result<(), Error> {
+    let rb_qtensor = rb_candle.define_class("QTensor", class::object())?;
+    rb_qtensor.define_method("ggml_dtype", method!(QTensor::ggml_dtype, 0))?;
+    rb_qtensor.define_method("rank", method!(QTensor::rank, 0))?;
+    rb_qtensor.define_method("shape", method!(QTensor::shape, 0))?;
+    rb_qtensor.define_method("dequantize", method!(QTensor::dequantize, 0))?;
+    Ok(())
+}

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

@@ -0,0 +1,3 @@
+use magnus::Error;
+
+pub type Result<T> = std::result::Result<T, Error>;