red-candle 1.0.1 → 1.1.0

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

@@ -1,15 +1,18 @@
 use magnus::{function, method, prelude::*, Error, Module, RArray, RHash, RModule, Ruby, TryConvert, Value};
 use std::cell::RefCell;
+use std::sync::Arc;
 
-use crate::llm::{GenerationConfig as RustGenerationConfig, TextGenerator, mistral::Mistral as RustMistral, llama::Llama as RustLlama, gemma::Gemma as RustGemma, QuantizedGGUF as RustQuantizedGGUF};
+use crate::llm::{GenerationConfig as RustGenerationConfig, TextGenerator, mistral::Mistral as RustMistral, llama::Llama as RustLlama, gemma::Gemma as RustGemma, qwen::Qwen as RustQwen, phi::Phi as RustPhi, QuantizedGGUF as RustQuantizedGGUF};
 use crate::ruby::{Result, Device};
+use crate::ruby::structured::StructuredConstraint;
 
 // Use an enum to handle different model types instead of trait objects
-#[derive(Debug)]
 enum ModelType {
     Mistral(RustMistral),
     Llama(RustLlama),
     Gemma(RustGemma),
+    Qwen(RustQwen),
+    Phi(RustPhi),
     QuantizedGGUF(RustQuantizedGGUF),
 }
 
@@ -19,6 +22,8 @@ impl ModelType {
             ModelType::Mistral(m) => m.generate(prompt, config),
             ModelType::Llama(m) => m.generate(prompt, config),
             ModelType::Gemma(m) => m.generate(prompt, config),
+            ModelType::Qwen(m) => m.generate(prompt, config),
+            ModelType::Phi(m) => m.generate(prompt, config),
             ModelType::QuantizedGGUF(m) => m.generate(prompt, config),
         }
     }
@@ -33,6 +38,8 @@ impl ModelType {
             ModelType::Mistral(m) => m.generate_stream(prompt, config, callback),
             ModelType::Llama(m) => m.generate_stream(prompt, config, callback),
             ModelType::Gemma(m) => m.generate_stream(prompt, config, callback),
+            ModelType::Qwen(m) => m.generate_stream(prompt, config, callback),
+            ModelType::Phi(m) => m.generate_stream(prompt, config, callback),
             ModelType::QuantizedGGUF(m) => m.generate_stream(prompt, config, callback),
         }
     }
@@ -42,6 +49,8 @@ impl ModelType {
             ModelType::Mistral(m) => m.clear_cache(),
             ModelType::Llama(m) => m.clear_cache(),
             ModelType::Gemma(m) => m.clear_cache(),
+            ModelType::Qwen(m) => m.clear_cache(),
+            ModelType::Phi(m) => m.clear_cache(),
             ModelType::QuantizedGGUF(m) => m.clear_cache(),
         }
     }
@@ -67,6 +76,8 @@ impl ModelType {
             },
             ModelType::Llama(m) => m.apply_chat_template(messages),
             ModelType::Gemma(m) => m.apply_chat_template(messages),
+            ModelType::Qwen(m) => m.apply_chat_template(messages),
+            ModelType::Phi(m) => m.apply_chat_template(messages),
             ModelType::QuantizedGGUF(m) => m.apply_chat_template(messages),
         }
     }
@@ -146,6 +157,13 @@ impl GenerationConfig {
             }
         }
         
+        // Handle constraint parameter
+        if let Some(value) = kwargs.get(magnus::Symbol::new("constraint")) {
+            if let Ok(constraint) = <&StructuredConstraint as TryConvert>::try_convert(value) {
+                config.constraint = Some(Arc::clone(&constraint.index));
+            }
+        }
+        
         Ok(Self { inner: config })
     }
 
@@ -191,9 +209,14 @@ impl GenerationConfig {
     pub fn debug_tokens(&self) -> bool {
         self.inner.debug_tokens
     }
+    pub fn constraint(&self) -> Option<StructuredConstraint> {
+        self.inner.constraint.as_ref().map(|c| StructuredConstraint {
+            index: Arc::clone(c),
+        })
+    }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone)]
 #[magnus::wrap(class = "Candle::LLM", mark, free_immediately)]
 pub struct LLM {
     model: std::sync::Arc<std::sync::Mutex<RefCell<ModelType>>>,
@@ -251,10 +274,22 @@ impl LLM {
                 })
                 .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to load model: {}", e)))?;
                 ModelType::Gemma(gemma)
+            } else if model_lower.contains("qwen") {
+                let qwen = rt.block_on(async {
+                    RustQwen::from_pretrained(&model_id, candle_device).await
+                })
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to load model: {}", e)))?;
+                ModelType::Qwen(qwen)
+            } else if model_lower.contains("phi") {
+                let phi = rt.block_on(async {
+                    RustPhi::from_pretrained(&model_id, candle_device).await
+                })
+                .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to load model: {}", e)))?;
+                ModelType::Phi(phi)
             } else {
                 return Err(Error::new(
                     magnus::exception::runtime_error(),
-                    format!("Unsupported model type: {}. Currently Mistral, Llama, and Gemma models are supported.", model_id),
+                    format!("Unsupported model type: {}. Currently Mistral, Llama, Gemma, Qwen, and Phi models are supported.", model_id),
                 ));
             }
         };
@@ -332,6 +367,8 @@ impl LLM {
             ModelType::Mistral(m) => Ok(crate::ruby::tokenizer::Tokenizer(m.tokenizer().clone())),
             ModelType::Llama(m) => Ok(crate::ruby::tokenizer::Tokenizer(m.tokenizer().clone())),
             ModelType::Gemma(m) => Ok(crate::ruby::tokenizer::Tokenizer(m.tokenizer().clone())),
+            ModelType::Qwen(m) => Ok(crate::ruby::tokenizer::Tokenizer(m.tokenizer().clone())),
+            ModelType::Phi(m) => Ok(crate::ruby::tokenizer::Tokenizer(m.tokenizer().clone())),
             ModelType::QuantizedGGUF(m) => Ok(crate::ruby::tokenizer::Tokenizer(m.tokenizer().clone())),
         }
     }
@@ -423,6 +460,7 @@ pub fn init_llm(rb_candle: RModule) -> Result<()> {
     rb_generation_config.define_method("stop_sequences", method!(GenerationConfig::stop_sequences, 0))?;
     rb_generation_config.define_method("include_prompt", method!(GenerationConfig::include_prompt, 0))?;
     rb_generation_config.define_method("debug_tokens", method!(GenerationConfig::debug_tokens, 0))?;
+    rb_generation_config.define_method("constraint", method!(GenerationConfig::constraint, 0))?;
     
     let rb_llm = rb_candle.define_class("LLM", magnus::class::object())?;
     rb_llm.define_singleton_method("_from_pretrained", function!(from_pretrained_wrapper, -1))?;

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

@@ -7,6 +7,7 @@ pub mod errors;
 pub mod utils;
 pub mod llm;
 pub mod tokenizer;
+pub mod structured;
 
 pub use embedding_model::{EmbeddingModel, EmbeddingModelInner};
 pub use tensor::Tensor;

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

@@ -0,0 +1,47 @@
+use magnus::{Error, Module, RModule, function, Object};
+use std::sync::Arc;
+
+use crate::structured::{SchemaProcessor, VocabularyAdapter, Index};
+use crate::ruby::{Result, tokenizer::Tokenizer};
+
+/// Ruby wrapper for structured generation constraints
+#[derive(Clone, Debug)]
+#[magnus::wrap(class = "Candle::StructuredConstraint", mark, free_immediately)]
+pub struct StructuredConstraint {
+    pub(crate) index: Arc<Index>,
+}
+
+impl StructuredConstraint {
+    /// Create a constraint from a JSON schema
+    pub fn from_schema(schema: String, tokenizer: &Tokenizer) -> Result<Self> {
+        let vocabulary = VocabularyAdapter::from_tokenizer(&tokenizer.0)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to create vocabulary: {}", e)))?;
+        
+        let processor = SchemaProcessor::new();
+        let index = processor.process_schema(&schema, &vocabulary)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to process schema: {}", e)))?;
+        
+        Ok(Self { index })
+    }
+    
+    /// Create a constraint from a regex pattern
+    pub fn from_regex(pattern: String, tokenizer: &Tokenizer) -> Result<Self> {
+        let vocabulary = VocabularyAdapter::from_tokenizer(&tokenizer.0)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to create vocabulary: {}", e)))?;
+        
+        let processor = SchemaProcessor::new();
+        let index = processor.process_regex(&pattern, &vocabulary)
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to process regex: {}", e)))?;
+        
+        Ok(Self { index })
+    }
+}
+
+pub fn init_structured(rb_candle: RModule) -> Result<()> {
+    let class = rb_candle.define_class("StructuredConstraint", magnus::class::object())?;
+    
+    class.define_singleton_method("from_schema", function!(StructuredConstraint::from_schema, 2))?;
+    class.define_singleton_method("from_regex", function!(StructuredConstraint::from_regex, 2))?;
+    
+    Ok(())
+}

data/ext/candle/src/structured/integration_test.rs

@@ -0,0 +1,130 @@
+#[cfg(test)]
+mod integration_tests {
+    use super::super::*;
+    use crate::tokenizer::{TokenizerWrapper, loader::TokenizerLoader};
+    use std::sync::Arc;
+    
+    #[tokio::test]
+    async fn test_schema_processor_with_vocabulary() {
+        // This test requires a tokenizer to create a vocabulary
+        let tokenizer_result = TokenizerLoader::from_hf_hub("bert-base-uncased", None).await;
+        
+        if let Ok(tokenizer) = tokenizer_result {
+            let wrapper = TokenizerWrapper::new(tokenizer);
+            
+            // Create vocabulary from tokenizer
+            let vocabulary = VocabularyAdapter::from_tokenizer(&wrapper)
+                .expect("Should create vocabulary");
+            
+            // Create schema processor
+            let processor = SchemaProcessor::new();
+            
+            // Test with a simple JSON schema
+            let schema = r#"{
+                "type": "object",
+                "properties": {
+                    "name": {"type": "string"},
+                    "age": {"type": "integer"}
+                },
+                "required": ["name", "age"]
+            }"#;
+            
+            // Process schema into Index
+            let index_result = processor.process_schema(schema, &vocabulary);
+            assert!(index_result.is_ok(), "Should process schema successfully");
+            
+            // Test caching - second call should use cache
+            let index2_result = processor.process_schema(schema, &vocabulary);
+            assert!(index2_result.is_ok(), "Should retrieve from cache");
+            
+            // Both should be the same Arc
+            let index1 = index_result.unwrap();
+            let index2 = index2_result.unwrap();
+            assert!(Arc::ptr_eq(&index1, &index2), "Should return cached Index");
+            
+            // Check cache stats
+            let (size, _) = processor.cache_stats();
+            assert_eq!(size, 1, "Cache should have one entry");
+        } else {
+            eprintln!("Skipping integration test - couldn't load tokenizer");
+        }
+    }
+    
+    #[tokio::test]
+    async fn test_regex_processing() {
+        let tokenizer_result = TokenizerLoader::from_hf_hub("bert-base-uncased", None).await;
+        
+        if let Ok(tokenizer) = tokenizer_result {
+            let wrapper = TokenizerWrapper::new(tokenizer);
+            let vocabulary = VocabularyAdapter::from_tokenizer(&wrapper)
+                .expect("Should create vocabulary");
+            
+            let processor = SchemaProcessor::new();
+            
+            // Test with a simple regex pattern
+            let email_regex = r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}";
+            
+            let index_result = processor.process_regex(email_regex, &vocabulary);
+            assert!(index_result.is_ok(), "Should process regex successfully");
+            
+            // Test different regex
+            let phone_regex = r"\d{3}-\d{3}-\d{4}";
+            let phone_index_result = processor.process_regex(phone_regex, &vocabulary);
+            assert!(phone_index_result.is_ok(), "Should process phone regex");
+            
+            // Cache should have both
+            let (size, _) = processor.cache_stats();
+            assert_eq!(size, 2, "Cache should have two entries");
+            
+            // Clear cache
+            processor.clear_cache();
+            let (size, _) = processor.cache_stats();
+            assert_eq!(size, 0, "Cache should be empty after clear");
+        }
+    }
+    
+    #[test]
+    fn test_various_json_schemas() {
+        let _processor = SchemaProcessor::new();
+        
+        // Array schema
+        let array_schema = serde_json::json!({
+            "type": "array",
+            "items": {"type": "string"}
+        });
+        
+        // Process as a full schema instead of testing private method
+        // This would need a mock vocabulary in a real test
+        // For now, just verify the schema is valid JSON
+        let json_str = serde_json::to_string(&array_schema).unwrap();
+        assert!(!json_str.is_empty(), "Should serialize array schema");
+        
+        // Nested object schema
+        let nested_schema = serde_json::json!({
+            "type": "object",
+            "properties": {
+                "user": {
+                    "type": "object",
+                    "properties": {
+                        "id": {"type": "integer"},
+                        "email": {"type": "string", "format": "email"}
+                    }
+                }
+            }
+        });
+        
+        // Verify nested schema is valid
+        let json_str = serde_json::to_string(&nested_schema).unwrap();
+        assert!(json_str.contains("properties"), "Should have nested properties");
+        
+        // Schema with enum
+        let enum_schema = serde_json::json!({
+            "type": "string",
+            "enum": ["red", "green", "blue"]
+        });
+        
+        // Verify enum schema is valid
+        let json_str = serde_json::to_string(&enum_schema).unwrap();
+        assert!(json_str.contains("enum"), "Should have enum values");
+    }
+}

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

@@ -0,0 +1,31 @@
+/// Structured generation support using Outlines
+/// 
+/// This module provides functionality to constrain language model generation
+/// to follow specific patterns, such as JSON schemas or regular expressions.
+
+pub mod vocabulary_adapter;
+pub mod schema_processor;
+
+pub use vocabulary_adapter::VocabularyAdapter;
+pub use schema_processor::SchemaProcessor;
+
+// Re-export commonly used types from outlines-core
+pub use outlines_core::prelude::Index;
+pub use outlines_core::vocabulary::Vocabulary;
+
+#[cfg(test)]
+mod vocabulary_adapter_simple_test;
+
+#[cfg(test)]
+mod integration_test;
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    
+    #[test]
+    fn test_module_imports() {
+        // Ensure all exports are available
+        let _ = VocabularyAdapter;
+    }
+}

data/ext/candle/src/structured/schema_processor.rs

@@ -0,0 +1,215 @@
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
+use candle_core::Result as CandleResult;
+use candle_core::Error as CandleError;
+use outlines_core::prelude::Index;
+use outlines_core::vocabulary::Vocabulary;
+use serde_json::Value as JsonValue;
+use outlines_core::json_schema;
+
+/// Processes JSON schemas into compiled Index objects for structured generation
+pub struct SchemaProcessor {
+    /// Cache of compiled Index objects keyed by schema hash
+    cache: Arc<Mutex<HashMap<u64, Arc<Index>>>>,
+}
+
+impl SchemaProcessor {
+    /// Create a new schema processor with an empty cache
+    pub fn new() -> Self {
+        Self {
+            cache: Arc::new(Mutex::new(HashMap::new())),
+        }
+    }
+    
+    /// Process a JSON schema into a compiled Index
+    /// 
+    /// # Arguments
+    /// * `schema` - JSON schema as a string
+    /// * `vocabulary` - The tokenizer's vocabulary
+    /// 
+    /// # Returns
+    /// A compiled Index ready for constrained generation
+    pub fn process_schema(
+        &self,
+        schema: &str,
+        vocabulary: &Vocabulary,
+    ) -> CandleResult<Arc<Index>> {
+        // Calculate hash of the schema for caching
+        let schema_hash = self.calculate_hash(schema);
+        
+        // Check cache first
+        if let Ok(cache) = self.cache.lock() {
+            if let Some(cached_index) = cache.get(&schema_hash) {
+                return Ok(Arc::clone(cached_index));
+            }
+        }
+        
+        // Parse the JSON schema
+        let schema_value: JsonValue = serde_json::from_str(schema)
+            .map_err(|e| CandleError::Msg(format!("Invalid JSON schema: {}", e)))?;
+        
+        // Convert schema to regex using Outlines
+        let regex = self.schema_to_regex(&schema_value)?;
+        
+        // Compile regex into Index
+        let index = self.compile_regex(&regex, vocabulary)?;
+        let index_arc = Arc::new(index);
+        
+        // Cache the compiled Index
+        if let Ok(mut cache) = self.cache.lock() {
+            cache.insert(schema_hash, Arc::clone(&index_arc));
+        }
+        
+        Ok(index_arc)
+    }
+    
+    /// Process a regex pattern directly into an Index
+    /// 
+    /// # Arguments
+    /// * `regex` - Regular expression pattern
+    /// * `vocabulary` - The tokenizer's vocabulary
+    /// 
+    /// # Returns
+    /// A compiled Index for the regex pattern
+    pub fn process_regex(
+        &self,
+        regex: &str,
+        vocabulary: &Vocabulary,
+    ) -> CandleResult<Arc<Index>> {
+        // Calculate hash for caching
+        let regex_hash = self.calculate_hash(regex);
+        
+        // Check cache
+        if let Ok(cache) = self.cache.lock() {
+            if let Some(cached_index) = cache.get(&regex_hash) {
+                return Ok(Arc::clone(cached_index));
+            }
+        }
+        
+        // Compile the regex
+        let index = self.compile_regex(regex, vocabulary)?;
+        let index_arc = Arc::new(index);
+        
+        // Cache it
+        if let Ok(mut cache) = self.cache.lock() {
+            cache.insert(regex_hash, Arc::clone(&index_arc));
+        }
+        
+        Ok(index_arc)
+    }
+    
+    /// Convert a JSON schema to a regex pattern
+    fn schema_to_regex(&self, schema: &JsonValue) -> CandleResult<String> {
+        // Use Outlines' built-in JSON schema to regex conversion
+        json_schema::regex_from_value(schema, None)
+            .map_err(|e| CandleError::Msg(format!("Failed to convert schema to regex: {:?}", e)))
+    }
+    
+    /// Compile a regex pattern into an Index
+    fn compile_regex(&self, regex: &str, vocabulary: &Vocabulary) -> CandleResult<Index> {
+        // Use Outlines to build the Index from regex
+        Index::new(regex, vocabulary)
+            .map_err(|e| CandleError::Msg(format!("Failed to build index from regex: {:?}", e)))
+    }
+    
+    /// Calculate a hash for caching
+    fn calculate_hash(&self, input: &str) -> u64 {
+        use std::collections::hash_map::DefaultHasher;
+        use std::hash::{Hash, Hasher};
+        
+        let mut hasher = DefaultHasher::new();
+        input.hash(&mut hasher);
+        hasher.finish()
+    }
+    
+    /// Clear the cache
+    pub fn clear_cache(&self) {
+        if let Ok(mut cache) = self.cache.lock() {
+            cache.clear();
+        }
+    }
+    
+    /// Get cache statistics
+    pub fn cache_stats(&self) -> (usize, usize) {
+        if let Ok(cache) = self.cache.lock() {
+            let size = cache.len();
+            let capacity = cache.capacity();
+            (size, capacity)
+        } else {
+            (0, 0)
+        }
+    }
+}
+
+impl Default for SchemaProcessor {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    
+    #[test]
+    fn test_schema_processor_creation() {
+        let processor = SchemaProcessor::new();
+        let (size, _) = processor.cache_stats();
+        assert_eq!(size, 0, "Cache should start empty");
+    }
+    
+    #[test]
+    fn test_cache_operations() {
+        let processor = SchemaProcessor::new();
+        
+        // Initially empty
+        let (size, _) = processor.cache_stats();
+        assert_eq!(size, 0);
+        
+        // After clear (should still be empty)
+        processor.clear_cache();
+        let (size, _) = processor.cache_stats();
+        assert_eq!(size, 0);
+    }
+    
+    #[test]
+    fn test_schema_to_regex_basic_types() {
+        let processor = SchemaProcessor::new();
+        
+        // Test string type
+        let string_schema = serde_json::json!({
+            "type": "string"
+        });
+        let regex = processor.schema_to_regex(&string_schema).unwrap();
+        // Just verify it produces a regex, exact format depends on Outlines
+        assert!(!regex.is_empty(), "String schema should produce a regex");
+        
+        // Test number type
+        let number_schema = serde_json::json!({
+            "type": "number"
+        });
+        let regex = processor.schema_to_regex(&number_schema).unwrap();
+        assert!(!regex.is_empty(), "Number schema should produce a regex");
+        
+        // Test boolean type
+        let bool_schema = serde_json::json!({
+            "type": "boolean"
+        });
+        let regex = processor.schema_to_regex(&bool_schema).unwrap();
+        assert!(regex.contains("true") && regex.contains("false"), "Boolean regex should contain true/false");
+    }
+    
+    #[test]
+    fn test_schema_with_pattern() {
+        let processor = SchemaProcessor::new();
+        
+        let schema = serde_json::json!({
+            "type": "string",
+            "pattern": r"^\d{3}-\d{3}-\d{4}$"
+        });
+        
+        let regex = processor.schema_to_regex(&schema).unwrap();
+        // Pattern should be included in the generated regex
+        assert!(regex.contains("\\d{3}"), "Should contain digit pattern");
+    }
+}