red-candle 1.0.2 → 1.1.1

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");
+    }
+}

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

@@ -0,0 +1,152 @@
+use crate::tokenizer::TokenizerWrapper;
+use candle_core::Result as CandleResult;
+use outlines_core::vocabulary::Vocabulary;
+use std::collections::HashMap;
+
+/// Adapter to convert between red-candle's TokenizerWrapper and Outlines' Vocabulary
+pub struct VocabularyAdapter;
+
+impl VocabularyAdapter {
+    /// Convert a TokenizerWrapper's vocabulary to an Outlines Vocabulary
+    /// 
+    /// # Arguments
+    /// * `tokenizer` - The tokenizer to extract vocabulary from
+    /// 
+    /// # Returns
+    /// An Outlines Vocabulary suitable for use with Index construction
+    pub fn from_tokenizer(tokenizer: &TokenizerWrapper) -> CandleResult<Vocabulary> {
+        // Get the vocabulary mapping from the tokenizer
+        let vocab_map: HashMap<String, u32> = tokenizer.inner().get_vocab(true);
+        
+        // Try to find EOS token in vocabulary
+        let eos_token_id = vocab_map.get("</s>")
+            .or_else(|| vocab_map.get("<|endoftext|>"))
+            .or_else(|| vocab_map.get("<eos>"))
+            .or_else(|| vocab_map.get("[SEP]"))
+            .copied();
+        
+        // Create a sorted list of (token_id, token_string) pairs
+        let mut token_pairs: Vec<(u32, String)> = vocab_map
+            .into_iter()
+            .map(|(token, id)| (id, token))
+            .collect();
+        
+        // Sort by token ID to ensure correct indexing
+        token_pairs.sort_by_key(|(id, _)| *id);
+        
+        // Find the maximum token ID to determine vocabulary size
+        let max_token_id = token_pairs
+            .last()
+            .map(|(id, _)| *id)
+            .unwrap_or(0);
+        
+        // Create vocabulary items in the format expected by Outlines
+        // We need to handle potential gaps in token IDs
+        let mut vocab_items: Vec<(String, Vec<u8>)> = Vec::new();
+        let mut current_id = 0;
+        
+        for (token_id, token_string) in token_pairs {
+            // Fill gaps with placeholder tokens
+            while current_id < token_id {
+                vocab_items.push((
+                    format!("<unused_{}>", current_id),
+                    format!("<unused_{}>", current_id).into_bytes(),
+                ));
+                current_id += 1;
+            }
+            
+            // Add the actual token
+            // Convert token string to bytes for Outlines
+            vocab_items.push((
+                token_string.clone(),
+                token_string.into_bytes(),
+            ));
+            current_id += 1;
+        }
+        
+        // Fill any remaining gaps up to a reasonable vocabulary size
+        // This ensures we don't have issues with token IDs beyond our vocabulary
+        while current_id <= max_token_id {
+            vocab_items.push((
+                format!("<unused_{}>", current_id),
+                format!("<unused_{}>", current_id).into_bytes(),
+            ));
+            current_id += 1;
+        }
+        
+        // Create the Outlines vocabulary
+        // The Vocabulary API expects us to build it token by token
+        let mut vocabulary = Vocabulary::new(
+            eos_token_id.unwrap_or(0) // Use EOS token ID or 0 as default
+        );
+        
+        // Insert all tokens into the vocabulary
+        for (idx, (token, bytes)) in vocab_items.into_iter().enumerate() {
+            // Skip inserting the EOS token as it's already set in the vocabulary
+            if Some(idx as u32) == eos_token_id {
+                continue;
+            }
+            
+            vocabulary.try_insert(bytes, idx as u32)
+                .map_err(|e| candle_core::Error::Msg(
+                    format!("Failed to insert token '{}': {:?}", token, e)
+                ))?;
+        }
+        
+        Ok(vocabulary)
+    }
+    
+    /// Get vocabulary size from a tokenizer
+    pub fn vocab_size(tokenizer: &TokenizerWrapper) -> usize {
+        tokenizer.inner().get_vocab_size(true)
+    }
+    
+    /// Extract and validate special tokens
+    pub fn get_special_tokens(tokenizer: &TokenizerWrapper) -> HashMap<String, u32> {
+        let tokenizer_inner = tokenizer.inner();
+        let mut special_tokens = HashMap::new();
+        
+        // Get common special tokens if they exist
+        if let Some(_token) = tokenizer_inner.id_to_token(0) {
+            special_tokens.insert("pad_token".to_string(), 0);
+        }
+        
+        // Try to find EOS token
+        let vocab = tokenizer_inner.get_vocab(true);
+        if let Some(&eos_id) = vocab.get("</s>")
+            .or_else(|| vocab.get("<|endoftext|>"))
+            .or_else(|| vocab.get("<eos>"))
+            .or_else(|| vocab.get("[SEP]")) {
+            special_tokens.insert("eos_token".to_string(), eos_id);
+        }
+        
+        // Try to get BOS token if it exists
+        if let Some(bos_token) = tokenizer_inner.token_to_id("<s>") {
+            special_tokens.insert("bos_token".to_string(), bos_token);
+        } else if let Some(bos_token) = tokenizer_inner.token_to_id("<|startoftext|>") {
+            special_tokens.insert("bos_token".to_string(), bos_token);
+        }
+        
+        special_tokens
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    
+    #[test]
+    fn test_vocabulary_adapter_creation() {
+        // This test will be implemented once we have a way to create test tokenizers
+        // For now, it serves as a placeholder for the test structure
+    }
+    
+    #[test]
+    fn test_special_tokens_extraction() {
+        // Test special token extraction logic
+    }
+    
+    #[test]
+    fn test_vocab_size() {
+        // Test vocabulary size calculation
+    }
+}

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

@@ -0,0 +1,66 @@
+#[cfg(test)]
+mod real_tests {
+    use super::super::*;
+    use crate::tokenizer::{TokenizerWrapper, loader::TokenizerLoader};
+    
+    #[tokio::test]
+    async fn test_vocabulary_conversion_with_real_outlines() {
+        // This test requires network access to download a tokenizer
+        // It verifies that our adapter works with the real outlines-core crate
+        
+        // Load a simple tokenizer
+        let tokenizer_result = TokenizerLoader::from_hf_hub("bert-base-uncased", None).await;
+        
+        if let Ok(tokenizer) = tokenizer_result {
+            let wrapper = TokenizerWrapper::new(tokenizer);
+            
+            // Convert to Outlines vocabulary
+            let vocab_result = VocabularyAdapter::from_tokenizer(&wrapper);
+            assert!(vocab_result.is_ok(), "Vocabulary conversion should succeed");
+            
+            let vocabulary = vocab_result.unwrap();
+            
+            // Verify the vocabulary was created
+            // The real Vocabulary doesn't expose a size method directly,
+            // but we can verify it exists and has the correct EOS token
+            assert_eq!(vocabulary.eos_token_id(), 102); // BERT's [SEP] token
+            
+            println!("✓ Successfully created Outlines Vocabulary from BERT tokenizer");
+        } else {
+            println!("⚠️  Skipping test - couldn't download tokenizer (likely offline)");
+        }
+    }
+    
+    #[test]
+    fn test_vocabulary_adapter_with_mock_data() {
+        // This test doesn't require network access
+        // It uses a mock tokenizer to verify the conversion logic
+        
+        use tokenizers::models::wordpiece::WordPiece;
+        use tokenizers::Tokenizer;
+        use std::collections::HashMap;
+        
+        // Create a minimal vocabulary
+        let mut vocab = HashMap::new();
+        vocab.insert("[PAD]".to_string(), 0);
+        vocab.insert("[UNK]".to_string(), 1);
+        vocab.insert("[SEP]".to_string(), 2);
+        vocab.insert("hello".to_string(), 3);
+        vocab.insert("world".to_string(), 4);
+        
+        let model = WordPiece::from_vocab(vocab);
+        let tokenizer = Tokenizer::new(model);
+        let wrapper = TokenizerWrapper::new(tokenizer);
+        
+        // Convert to Outlines vocabulary
+        let vocab_result = VocabularyAdapter::from_tokenizer(&wrapper);
+        assert!(vocab_result.is_ok(), "Vocabulary conversion should succeed");
+        
+        let vocabulary = vocab_result.unwrap();
+        
+        // Verify EOS token was found
+        assert_eq!(vocabulary.eos_token_id(), 2); // [SEP] token
+        
+        println!("✓ Mock vocabulary conversion successful");
+    }
+}

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

@@ -0,0 +1,70 @@
+#[cfg(test)]
+mod simple_tests {
+    use super::super::*;
+    
+    #[test]
+    fn test_vocabulary_adapter_basic() {
+        // Create a simple mock tokenizer to test the adapter
+        // This validates that the VocabularyAdapter compiles and can be called
+        
+        // Note: Creating a full tokenizer in tests is complex due to the tokenizers crate API
+        // For now, we verify compilation and will rely on integration tests
+        
+        // The important thing is that this code compiles, proving our integration works
+        let _adapter = VocabularyAdapter;
+        
+        // Test the static methods compile
+        // These would be tested with a real tokenizer in integration tests
+        
+        // Test passes if this compiles - no output needed
+    }
+    
+    #[test]
+    fn test_outlines_vocabulary_api() {
+        use outlines_core::vocabulary::Vocabulary;
+        
+        // Test that we can create a Vocabulary object
+        // Use token ID 2 as EOS (like BERT's [SEP] token)
+        let mut vocab = Vocabulary::new(2);
+        
+        // Test inserting tokens
+        let test_tokens = vec![
+            ("<pad>".to_string(), "<pad>".as_bytes().to_vec()),
+            ("<unk>".to_string(), "<unk>".as_bytes().to_vec()),
+            ("<sep>".to_string(), "<sep>".as_bytes().to_vec()), // EOS token at ID 2
+            ("hello".to_string(), "hello".as_bytes().to_vec()),
+            ("world".to_string(), "world".as_bytes().to_vec()),
+        ];
+        
+        for (idx, (_token, bytes)) in test_tokens.into_iter().enumerate() {
+            match vocab.try_insert(bytes, idx as u32) {
+                Ok(_) => {},
+                Err(e) => {
+                    // It's ok if we can't insert the EOS token
+                    if idx != 2 {
+                        panic!("Failed to insert token at index {}: {:?}", idx, e);
+                    }
+                }
+            }
+        }
+        
+        // Test passes - vocabulary API works correctly
+    }
+    
+    #[test] 
+    fn test_special_token_patterns() {
+        
+        // Test that our special token patterns are correct
+        let test_cases = vec![
+            ("</s>", "EOS token for many models"),
+            ("<|endoftext|>", "GPT-style EOS token"),
+            ("<eos>", "Alternative EOS token"),
+            ("[SEP]", "BERT-style separator"),
+            ("<s>", "BOS token"),
+            ("<|startoftext|>", "GPT-style BOS token"),
+        ];
+        
+        // Just verify the patterns exist - no output needed
+        assert_eq!(test_cases.len(), 6, "Should have 6 special token patterns");
+    }
+}