red-candle 1.0.0.pre.5 → 1.0.0.pre.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 91a4c43a1a12d6d8960f1a1d190c9bfe8ea60db75f687233012d09b8c90b5020
-  data.tar.gz: 3f6ce143cd38856365231baebe25a188e7d5824d930ecdae79c1660b3ad6c787
+  metadata.gz: 07ca4e6eb0b65eac5b62f4b3622ed3189f203279265b7174936ccfd5ff3e5099
+  data.tar.gz: f4970f5c4376453cde1ee18b93155f69ca634ccc3e4a359a45b49d7f20379f64
 SHA512:
-  metadata.gz: 273a01c438b085509a433602097b5ad4bcdb3420fc19ebe84c4fd37bf43ae5e6e1040701ecaad9e2864c0cda31d6128d4b2df6c395290994c17f135620282be6
-  data.tar.gz: 9193d01d8bfc704b982c839f9aebac23217ae5d56b5afae6228daae73b60e219b049fe8fc8d5ad13657ceb110476489c68bbb9aa6397688bf70a976a8d62d41e
+  metadata.gz: 10ed0881ec2f67ab1e798401e857eac638049b254b20460bcb5565cee822b24ce2abe23d0ce00275dcb1d1ddebfd926d47eac7e6d54924937da4356a36211224
+  data.tar.gz: d24fa67f74cd62c87ea1666e9488f12e8773d15e2d62b806bd38ca7cb20215d819b0502d352e4d310d1377b9ab64debbdd664148c70fc1ac70f1f2e23e9b516c

data/README.md

@@ -47,6 +47,7 @@ Red-Candle now supports Large Language Models (LLMs) with GPU acceleration!
 
 ### Supported Models
 
+- **Gemma**: Google's Gemma models (e.g., `google/gemma-2b`, `google/gemma-7b`, `google/gemma-2b-it`)
 - **Llama**: Llama 2 and Llama 3 models (e.g., `TinyLlama/TinyLlama-1.1B-Chat-v1.0`, `meta-llama/Llama-2-7b-hf`, `NousResearch/Llama-2-7b-hf`)
 - **Mistral**: All Mistral models (e.g., `mistralai/Mistral-7B-Instruct-v0.1`)
 
@@ -67,10 +68,10 @@ device = Candle::Device.cpu     # CPU (default)
 device = Candle::Device.metal   # Apple GPU (Metal)
 device = Candle::Device.cuda    # NVIDIA GPU (CUDA)
 
-# Load a Llama model
-llm = Candle::LLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", device: device)
-# Or a Mistral model
-llm = Candle::LLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", device: device)
+# Load a model
+llm = Candle::LLM.from_pretrained("google/gemma-2b-it", device: device)  # Gemma
+# llm = Candle::LLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", device: device)  # Llama
+# llm = Candle::LLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", device: device)  # Mistral
 
 # Generate text
 response = llm.generate("What is Ruby?", config: Candle::GenerationConfig.balanced)

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

@@ -0,0 +1,340 @@
+use candle_core::{DType, Device, Result as CandleResult, Tensor};
+use candle_nn::VarBuilder;
+use candle_transformers::models::gemma::{Config, Model as GemmaModel};
+use hf_hub::{api::tokio::Api, Repo};
+use tokenizers::Tokenizer;
+
+use super::{GenerationConfig, TextGeneration, TextGenerator, TokenizerWrapper};
+
+#[derive(Debug)]
+pub struct Gemma {
+    model: GemmaModel,
+    tokenizer: TokenizerWrapper,
+    device: Device,
+    model_id: String,
+    eos_token_id: u32,
+}
+
+impl Gemma {
+    /// Clear the KV cache between generations
+    pub fn clear_kv_cache(&mut self) {
+        self.model.clear_kv_cache();
+    }
+    
+    /// Load a Gemma model from HuggingFace Hub
+    pub async fn from_pretrained(model_id: &str, device: Device) -> CandleResult<Self> {
+        let api = Api::new()
+            .map_err(|e| candle_core::Error::Msg(format!("Failed to create HF API: {}", e)))?;
+        
+        let repo = api.repo(Repo::model(model_id.to_string()));
+        
+        // Download model files
+        let config_filename = repo
+            .get("config.json")
+            .await
+            .map_err(|e| candle_core::Error::Msg(format!("Failed to download config: {}", e)))?;
+        
+        let tokenizer_filename = repo
+            .get("tokenizer.json")
+            .await
+            .map_err(|e| candle_core::Error::Msg(format!("Failed to download tokenizer: {}", e)))?;
+        
+        // Try different file patterns for model weights
+        let weights_filenames = if let Ok(single_file) = repo.get("model.safetensors").await {
+            vec![single_file]
+        } else {
+            // Try to find sharded model files
+            let mut sharded_files = Vec::new();
+            let mut index = 1;
+            loop {
+                // Try common shard counts for Gemma models
+                let mut found = false;
+                for total in [2, 3, 4, 5, 6, 7, 8, 10, 12] {
+                    let filename = format!("model-{:05}-of-{:05}.safetensors", index, total);
+                    if let Ok(file) = repo.get(&filename).await {
+                        sharded_files.push(file);
+                        found = true;
+                        break;
+                    }
+                }
+                if !found {
+                    break;
+                }
+                index += 1;
+            }
+            
+            if sharded_files.is_empty() {
+                return Err(candle_core::Error::Msg(
+                    "Could not find model weights. Tried: model.safetensors, model-*-of-*.safetensors".to_string()
+                ));
+            }
+            sharded_files
+        };
+        
+        // Load config
+        let config: Config = serde_json::from_reader(std::fs::File::open(config_filename)?)
+            .map_err(|e| candle_core::Error::Msg(format!("Failed to parse config: {}", e)))?;
+        
+        // Load tokenizer
+        let tokenizer = Tokenizer::from_file(tokenizer_filename)
+            .map_err(|e| candle_core::Error::Msg(format!("Failed to load tokenizer: {}", e)))?;
+        
+        // Gemma uses specific tokens
+        let eos_token_id = {
+            let vocab = tokenizer.get_vocab(true);
+            vocab.get("<eos>")
+                .or_else(|| vocab.get("<end_of_turn>"))
+                .copied()
+                .unwrap_or(1) // Default Gemma EOS
+        };
+        
+        // Load model weights
+        let vb = unsafe {
+            VarBuilder::from_mmaped_safetensors(&weights_filenames, DType::F32, &device)?
+        };
+        
+        let model = GemmaModel::new(false, &config, vb)?; // Don't use flash attention for now
+        
+        Ok(Self {
+            model,
+            tokenizer: TokenizerWrapper::new(tokenizer),
+            device,
+            model_id: model_id.to_string(),
+            eos_token_id,
+        })
+    }
+
+    /// Create from existing components (useful for testing)
+    pub fn new(
+        model: GemmaModel,
+        tokenizer: Tokenizer,
+        device: Device,
+        model_id: String,
+    ) -> Self {
+        let eos_token_id = {
+            let vocab = tokenizer.get_vocab(true);
+            vocab.get("<eos>")
+                .or_else(|| vocab.get("<end_of_turn>"))
+                .copied()
+                .unwrap_or(1)
+        };
+        
+        Self {
+            model,
+            tokenizer: TokenizerWrapper::new(tokenizer),
+            device,
+            model_id,
+            eos_token_id,
+        }
+    }
+
+    fn generate_tokens(
+        &mut self,
+        prompt_tokens: Vec<u32>,
+        config: &GenerationConfig,
+        mut callback: Option<impl FnMut(&str)>,
+    ) -> CandleResult<Vec<u32>> {
+        let mut text_gen = TextGeneration::from_config(config);
+        text_gen.set_eos_token_id(self.eos_token_id);
+        text_gen.set_tokens(prompt_tokens.clone());
+        
+        let mut all_tokens = prompt_tokens.clone();
+        let start_gen = all_tokens.len();
+        
+        for index in 0..config.max_length {
+            let context_size = if index > 0 { 1 } else { all_tokens.len() };
+            let start_pos = all_tokens.len().saturating_sub(context_size);
+            let ctxt = &all_tokens[start_pos..];
+            
+            let input = Tensor::new(ctxt, &self.device)?.unsqueeze(0)?;
+            let input = input.contiguous()?;
+            let logits = self.model.forward(&input, start_pos)?;
+            
+            let logits = logits.squeeze(0)?;
+            let logits = if logits.dims().len() == 2 {
+                let seq_len = logits.dim(0)?;
+                logits.narrow(0, seq_len - 1, 1)?.squeeze(0)?
+            } else {
+                logits
+            };
+            
+            let logits = logits.to_dtype(DType::F32)?;
+            
+            let next_token = text_gen.sample_next_token(
+                &logits,
+                Some((config.repetition_penalty, config.repetition_penalty_last_n)),
+            )?;
+            
+            all_tokens.push(next_token);
+            
+            // Stream callback
+            if let Some(ref mut cb) = callback {
+                let token_text = self.tokenizer.token_to_piece(next_token)?;
+                cb(&token_text);
+            }
+            
+            // Check stop conditions
+            if text_gen.should_stop(next_token, config.max_length) {
+                break;
+            }
+            
+            // Check stop sequences
+            let generated_text = self.tokenizer.decode(&all_tokens[start_gen..], true)?;
+            if text_gen.check_stop_sequences(&generated_text, &config.stop_sequences) {
+                break;
+            }
+        }
+        
+        Ok(if config.include_prompt {
+            all_tokens
+        } else {
+            all_tokens[start_gen..].to_vec()
+        })
+    }
+    
+    fn generate_tokens_decoded(
+        &mut self,
+        prompt_tokens: Vec<u32>,
+        config: &GenerationConfig,
+        mut callback: Option<impl FnMut(&str)>,
+    ) -> CandleResult<Vec<u32>> {
+        let mut text_gen = TextGeneration::from_config(config);
+        text_gen.set_eos_token_id(self.eos_token_id);
+        text_gen.set_tokens(prompt_tokens.clone());
+        
+        let mut all_tokens = prompt_tokens.clone();
+        let start_gen = all_tokens.len();
+        let mut previously_decoded = String::new();
+        
+        for index in 0..config.max_length {
+            let context_size = if index > 0 { 1 } else { all_tokens.len() };
+            let start_pos = all_tokens.len().saturating_sub(context_size);
+            let ctxt = &all_tokens[start_pos..];
+            
+            let input = Tensor::new(ctxt, &self.device)?.unsqueeze(0)?;
+            let input = input.contiguous()?;
+            let logits = self.model.forward(&input, start_pos)?;
+            
+            let logits = logits.squeeze(0)?;
+            let logits = if logits.dims().len() == 2 {
+                let seq_len = logits.dim(0)?;
+                logits.narrow(0, seq_len - 1, 1)?.squeeze(0)?
+            } else {
+                logits
+            };
+            
+            let logits = logits.to_dtype(DType::F32)?;
+            
+            let next_token = text_gen.sample_next_token(
+                &logits,
+                Some((config.repetition_penalty, config.repetition_penalty_last_n)),
+            )?;
+            
+            all_tokens.push(next_token);
+            
+            // Stream callback with incremental decoding
+            if let Some(ref mut cb) = callback {
+                let current_decoded = self.tokenizer.decode(&all_tokens[start_gen..], true)?;
+                
+                if current_decoded.len() > previously_decoded.len() {
+                    let new_text = &current_decoded[previously_decoded.len()..];
+                    cb(new_text);
+                    previously_decoded = current_decoded;
+                }
+            }
+            
+            // Check stop conditions
+            if text_gen.should_stop(next_token, config.max_length) {
+                break;
+            }
+            
+            // Check stop sequences
+            let generated_text = if callback.is_some() {
+                previously_decoded.clone()
+            } else {
+                self.tokenizer.decode(&all_tokens[start_gen..], true)?
+            };
+            
+            if text_gen.check_stop_sequences(&generated_text, &config.stop_sequences) {
+                break;
+            }
+        }
+        
+        Ok(if config.include_prompt {
+            all_tokens
+        } else {
+            all_tokens[start_gen..].to_vec()
+        })
+    }
+    
+    /// Apply Gemma chat template
+    pub fn apply_chat_template(&self, messages: &[serde_json::Value]) -> CandleResult<String> {
+        let mut prompt = String::new();
+        
+        // Gemma uses a specific format:
+        // <start_of_turn>user\n{user_message}<end_of_turn>
+        // <start_of_turn>model\n{model_message}<end_of_turn>
+        
+        for message in messages {
+            let role = message["role"].as_str().unwrap_or("");
+            let content = message["content"].as_str().unwrap_or("");
+            
+            match role {
+                "system" => {
+                    // Gemma doesn't have explicit system messages, prepend to first user message
+                    prompt.push_str(&format!("<start_of_turn>user\nSystem: {}\n", content));
+                }
+                "user" => {
+                    if !prompt.contains("<start_of_turn>user") || prompt.ends_with("<end_of_turn>\n") {
+                        prompt.push_str("<start_of_turn>user\n");
+                    }
+                    prompt.push_str(&format!("{}<end_of_turn>\n", content));
+                }
+                "assistant" | "model" => {
+                    prompt.push_str(&format!("<start_of_turn>model\n{}<end_of_turn>\n", content));
+                }
+                _ => {}
+            }
+        }
+        
+        // Add the model prompt
+        prompt.push_str("<start_of_turn>model\n");
+        
+        Ok(prompt)
+    }
+}
+
+impl TextGenerator for Gemma {
+    fn generate(
+        &mut self,
+        prompt: &str,
+        config: &GenerationConfig,
+    ) -> CandleResult<String> {
+        let prompt_tokens = self.tokenizer.encode(prompt, true)?;
+        let output_tokens = self.generate_tokens(prompt_tokens, config, None::<fn(&str)>)?;
+        self.tokenizer.decode(&output_tokens, true)
+    }
+
+    fn generate_stream(
+        &mut self,
+        prompt: &str,
+        config: &GenerationConfig,
+        mut callback: impl FnMut(&str),
+    ) -> CandleResult<String> {
+        let prompt_tokens = self.tokenizer.encode(prompt, true)?;
+        let output_tokens = self.generate_tokens_decoded(prompt_tokens, config, Some(&mut callback))?;
+        self.tokenizer.decode(&output_tokens, true)
+    }
+
+    fn model_name(&self) -> &str {
+        &self.model_id
+    }
+
+    fn device(&self) -> &Device {
+        &self.device
+    }
+    
+    fn clear_cache(&mut self) {
+        self.clear_kv_cache();
+    }
+}

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

@@ -3,6 +3,7 @@ use tokenizers::Tokenizer;
 
 pub mod mistral;
 pub mod llama;
+pub mod gemma;
 pub mod generation_config;
 pub mod text_generation;
 

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

@@ -1,7 +1,7 @@
 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, llama::Llama as RustLlama};
+use crate::llm::{GenerationConfig as RustGenerationConfig, TextGenerator, mistral::Mistral as RustMistral, llama::Llama as RustLlama, gemma::Gemma as RustGemma};
 use crate::ruby::{Result as RbResult, Device as RbDevice};
 
 // Use an enum to handle different model types instead of trait objects
@@ -9,6 +9,7 @@ use crate::ruby::{Result as RbResult, Device as RbDevice};
 enum ModelType {
     Mistral(RustMistral),
     Llama(RustLlama),
+    Gemma(RustGemma),
 }
 
 impl ModelType {
@@ -16,6 +17,7 @@ impl ModelType {
         match self {
             ModelType::Mistral(m) => m.generate(prompt, config),
             ModelType::Llama(m) => m.generate(prompt, config),
+            ModelType::Gemma(m) => m.generate(prompt, config),
         }
     }
 
@@ -28,6 +30,7 @@ impl ModelType {
         match self {
             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),
         }
     }
 
@@ -36,6 +39,7 @@ impl ModelType {
         match self {
             ModelType::Mistral(m) => m.model_name(),
             ModelType::Llama(m) => m.model_name(),
+            ModelType::Gemma(m) => m.model_name(),
         }
     }
     
@@ -43,6 +47,7 @@ impl ModelType {
         match self {
             ModelType::Mistral(m) => m.clear_cache(),
             ModelType::Llama(m) => m.clear_cache(),
+            ModelType::Gemma(m) => m.clear_cache(),
         }
     }
     
@@ -66,6 +71,7 @@ impl ModelType {
                 Ok(prompt)
             },
             ModelType::Llama(m) => m.apply_chat_template(messages),
+            ModelType::Gemma(m) => m.apply_chat_template(messages),
         }
     }
 }
@@ -208,16 +214,22 @@ impl LLM {
             })
             .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to load model: {}", e)))?;
             ModelType::Mistral(mistral)
-        } else if model_lower.contains("llama") || model_lower.contains("meta-llama") {
+        } else if model_lower.contains("llama") || model_lower.contains("meta-llama") || model_lower.contains("tinyllama") {
             let llama = rt.block_on(async {
                 RustLlama::from_pretrained(&model_id, candle_device).await
             })
             .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to load model: {}", e)))?;
             ModelType::Llama(llama)
+        } else if model_lower.contains("gemma") || model_lower.contains("google/gemma") {
+            let gemma = rt.block_on(async {
+                RustGemma::from_pretrained(&model_id, candle_device).await
+            })
+            .map_err(|e| Error::new(magnus::exception::runtime_error(), format!("Failed to load model: {}", e)))?;
+            ModelType::Gemma(gemma)
         } else {
             return Err(Error::new(
                 magnus::exception::runtime_error(),
-                format!("Unsupported model type: {}. Currently only Mistral and Llama models are supported.", model_id),
+                format!("Unsupported model type: {}. Currently Mistral, Llama, and Gemma models are supported.", model_id),
             ));
         };
         

data/lib/candle/version.rb

@@ -1,3 +1,3 @@
 module Candle
-  VERSION = "1.0.0.pre.5"
+  VERSION = "1.0.0.pre.6"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: red-candle
 version: !ruby/object:Gem::Version
-  version: 1.0.0.pre.5
+  version: 1.0.0.pre.6
 platform: ruby
 authors:
 - Christopher Petersen
@@ -9,7 +9,7 @@ authors:
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2025-07-09 00:00:00.000000000 Z
+date: 2025-07-10 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rb_sys
@@ -47,6 +47,7 @@ files:
 - ext/candle/extconf.rb
 - ext/candle/rustfmt.toml
 - ext/candle/src/lib.rs
+- ext/candle/src/llm/gemma.rs
 - ext/candle/src/llm/generation_config.rs
 - ext/candle/src/llm/llama.rs
 - ext/candle/src/llm/mistral.rs