red-candle 1.5.0 → 1.6.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: a15236dd78a04cfcc6d7978ca088934987b84a599371c1508383d147172abcc3
-  data.tar.gz: f2661ad1220f566dd8a7d9eab1d74bd30007ed022c9381daf2697d12d88fac96
+  metadata.gz: 00c0870d599db76ba556ab880f69219bc32d72698c3570cd308df3923b1332f9
+  data.tar.gz: bffbd02fa1fe11813a304ed16771e6a72ee01f85ccb445eb67c5da6bf696c670
 SHA512:
-  metadata.gz: d74f26e2c0527f25424245442b5b74cf1575505a42a2c1c823be31606433a6c9e4736afc6d8062bd54642cbc8373deaae197bfc7163cbaa392d457a844e67628
-  data.tar.gz: 1ffc6d400ba4ffe9c23603fb42adae088142446f1d0bc03d0365a2bfcde3c1d1beab51eec01481d1a6e89f686c235e2ec58f0703daed73fabcf5cfd7f7e22d95
+  metadata.gz: c9ab39f01e4777c3d9906592cf62b8b1d2d03fd04be4ba82d8ddee837678ac477be88b1a0a48fa898e443f7469681e26f25297d4d310c39fb5222137d3b17d87
+  data.tar.gz: 22fe87f0fb64754a0a19bcc632dff72b0cc43838d73797aa991c961b65c9ecb7d9bd7aacbde51baf4063011081f4111a64f53726b51c87aca3e48214304e0119

data/README.md

@@ -273,6 +273,85 @@ See [STRUCTURED_GENERATION.md](docs/STRUCTURED_GENERATION.md) for detailed docum
 
 **Note on Reliability**: Structured generation constrains the model's output tokens, but success rates vary by model size and schema complexity. Smaller models (< 7B parameters) may occasionally produce incomplete or invalid JSON, especially with complex schemas. Consider implementing retry logic or fallback strategies in production applications. Larger models generally perform much better with structured generation.
 
+## Tool Calling
+
+Red-candle supports tool/function calling, enabling models to invoke external functions during generation. This works best with models fine-tuned for tool calling, such as Qwen3.
+
+### Defining Tools
+
+```ruby
+get_weather = Candle::Tool.new(
+  name: "get_weather",
+  description: "Get the current weather for a city",
+  parameters: {
+    type: "object",
+    properties: { city: { type: "string", description: "City name" } },
+    required: ["city"]
+  }
+) { |args| { city: args["city"], temperature: 72, condition: "sunny" } }
+```
+
+### Extracting Tool Calls
+
+`chat_with_tools` injects tool definitions into the system prompt, generates a response, and parses any `<tool_call>` tags from the output. It does **not** feed results back to the model — it just tells you what the model wants to call. You decide what to do with it:
+
+```ruby
+llm = Candle::LLM.from_pretrained("Qwen/Qwen3-0.6B")
+
+messages = [{ role: "user", content: "What's the weather in San Francisco?" }]
+result = llm.chat_with_tools(messages, tools: [get_weather],
+  config: Candle::GenerationConfig.deterministic(max_length: 500))
+
+if result.has_tool_calls?
+  result.tool_calls.each do |tc|
+    puts "#{tc.name}(#{tc.arguments})"
+    output = get_weather.call(tc.arguments)
+    puts "=> #{output}"
+  end
+else
+  puts result.text_response
+end
+```
+
+Pass `execute: true` to automatically run the tools (but still no round-trip back to the model):
+
+```ruby
+result = llm.chat_with_tools(messages, tools: [get_weather], execute: true,
+  config: Candle::GenerationConfig.deterministic(max_length: 500))
+
+result.tool_results.each do |tr|
+  puts "#{tr[:tool_call].name} => #{tr[:result]}"
+end
+```
+
+### Agent (Multi-Turn Tool Loop)
+
+`Candle::Agent` completes the round-trip: generate → parse tool calls → execute → feed results back to the model → repeat until the model produces a final text answer or hits `max_iterations`. This is a convenience wrapper for quick prototyping — for production use, frameworks like [RubyLLM](https://github.com/crmne/ruby_llm) manage this loop for you via the [ruby_llm-red_candle](https://github.com/scientist-labs/ruby_llm-red_candle) plugin:
+
+```ruby
+agent = Candle::Agent.new(llm, tools: [get_weather, lookup_price], max_iterations: 5)
+result = agent.run("What's the weather in Paris, and how much does a widget cost?",
+  config: Candle::GenerationConfig.deterministic(max_length: 1000))
+
+puts result.response         # Final text answer from the model
+puts result.iterations       # Number of generate cycles
+puts result.tool_calls_made  # Number of tools invoked
+```
+
+### Model Recommendations
+
+Tool calling quality depends heavily on model size:
+
+| Model | Tool Calling Quality |
+|-------|---------------------|
+| **Qwen3-8B GGUF** (~5 GB) | Calls correct tools, self-corrects errors, but may hallucinate values from tool results |
+| **Qwen3-4B GGUF** (~2.5 GB) | Calls correct tools, occasional reasoning errors |
+| **Qwen3-0.6B** (~1.2 GB) | Single-turn works, needs `max_length: 500+` for thinking |
+| SmolLM2-360M | Does not work |
+| TinyLlama-1.1B | Does not work (not fine-tuned for tool calling) |
+
+**Tip:** Qwen3 models use a `<think>` reasoning block before producing tool calls. Set `max_length` high enough (500+ for 0.6B, 1000+ for larger models) to allow room for both thinking and the tool call.
+
 ## ⚠️ Model Format Requirements
 
 ### EmbeddingModels and Rerankers: Safetensors Only
@@ -355,12 +434,26 @@ embedding = model.embedding("Hi there!")
 
 Red-Candle includes support for cross-encoder reranking models, which can be used to reorder documents by relevance to a query. This is particularly useful for improving search results or implementing retrieval-augmented generation (RAG) systems.
 
+### Supported Models
+
+Red-Candle supports both **BERT** and **XLM-RoBERTa** reranker architectures. The model type is auto-detected from `config.json`.
+
+| Model | Architecture | Params | Notes |
+|-------|-------------|--------|-------|
+| `BAAI/bge-reranker-base` | XLM-RoBERTa | 278M | Recommended — strong quality, multilingual |
+| `BAAI/bge-reranker-large` | XLM-RoBERTa | 560M | Best quality, higher resource usage |
+| `BAAI/bge-reranker-v2-m3` | XLM-RoBERTa | 278M | Multilingual, very strong |
+| `cross-encoder/ms-marco-MiniLM-L-12-v2` | BERT | 33M | Lightweight, English only |
+
 ### Basic Usage
 
 ```ruby
 require 'candle'
 
-# Initialize the reranker with a cross-encoder model
+# Initialize the reranker (BGE reranker recommended for quality)
+reranker = Candle::Reranker.from_pretrained("BAAI/bge-reranker-base")
+
+# Or use the lighter BERT-based model
 reranker = Candle::Reranker.from_pretrained("cross-encoder/ms-marco-MiniLM-L-12-v2")
 
 # Or with custom max_length for truncation (default is 512)
@@ -449,7 +542,7 @@ results = reranker.rerank_with_pooling(query, documents, "cls")
 results = reranker.rerank_sigmoid_with_pooling(query, documents, "cls")
 ```
 
-Note: The default "pooler" method is recommended as it matches how cross-encoder models are trained. Other pooling methods may produce different ranking results.
+Note: Pooling methods only apply to BERT-based models. XLM-RoBERTa models (e.g., BGE rerankers) have a built-in classification head and ignore the `pooling_method` parameter. For BERT models, the default "pooler" method is recommended as it matches how cross-encoder models are trained.
 
 ### CUDA Support
 
@@ -467,11 +560,11 @@ Cross-encoder reranking models differ from bi-encoder embedding models:
 - **Bi-encoders** (like the embedding models above) encode queries and documents separately into dense vectors
 - **Cross-encoders** process the query and document together, allowing for more nuanced relevance scoring
 
-The reranker uses a BERT-based architecture that:
-1. Concatenates the query and document with special tokens: `[CLS] query [SEP] document [SEP]`
-2. Processes them jointly through BERT layers
-3. Applies a pooler layer (dense + tanh) to the [CLS] token
-4. Uses a classifier layer to produce a single relevance score
+The reranker concatenates the query and document with special tokens and processes them jointly through transformer layers to produce a single relevance score.
+
+**BERT models** (e.g., MiniLM): Use a pooler layer (dense + tanh) on the [CLS] token, then a classifier layer. Pooling method is configurable (`pooler`, `cls`, `mean`).
+
+**XLM-RoBERTa models** (e.g., BGE rerankers): Use a built-in classification head that returns logits directly. The `pooling_method` parameter is ignored — the model handles its own pooling internally.
 
 This joint processing allows cross-encoders to capture subtle semantic relationships between queries and documents, making them more accurate for reranking tasks, though at the cost of higher computational requirements.
 

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

@@ -319,9 +319,9 @@ impl Llama {
                     if i == 1 || (i == 0 && system_message.is_empty()) {
                         // First user message
                         if !system_message.is_empty() {
-                            prompt.push_str(&format!("<s>[INST] <<SYS>>\n{}\n<</SYS>>\n\n{} [/INST]", system_message, content));
+                            prompt.push_str(&format!("[INST] <<SYS>>\n{}\n<</SYS>>\n\n{} [/INST]", system_message, content));
                         } else {
-                            prompt.push_str(&format!("<s>[INST] {} [/INST]", content));
+                            prompt.push_str(&format!("[INST] {} [/INST]", content));
                         }
                     } else {
                         prompt.push_str(&format!(" [INST] {} [/INST]", content));
@@ -340,7 +340,7 @@ impl Llama {
     fn apply_llama3_template(&self, messages: &[serde_json::Value]) -> CandleResult<String> {
         let mut prompt = String::new();
         
-        prompt.push_str("<|begin_of_text|>");
+        // BOS token is added by the tokenizer's encode(prompt, add_special_tokens=true)
         
         for message in messages {
             let role = message["role"].as_str().unwrap_or("");

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

@@ -386,9 +386,9 @@ impl QuantizedGGUF {
                 "user" => {
                     if i == 1 || (i == 0 && system_message.is_empty()) {
                         if !system_message.is_empty() {
-                            prompt.push_str(&format!("<s>[INST] <<SYS>>\n{}\n<</SYS>>\n\n{} [/INST]", system_message, content));
+                            prompt.push_str(&format!("[INST] <<SYS>>\n{}\n<</SYS>>\n\n{} [/INST]", system_message, content));
                         } else {
-                            prompt.push_str(&format!("<s>[INST] {} [/INST]", content));
+                            prompt.push_str(&format!("[INST] {} [/INST]", content));
                         }
                     } else {
                         prompt.push_str(&format!(" [INST] {} [/INST]", content));
@@ -406,7 +406,7 @@ impl QuantizedGGUF {
     
     fn apply_llama3_template(&self, messages: &[serde_json::Value]) -> CandleResult<String> {
         let mut prompt = String::new();
-        prompt.push_str("<|begin_of_text|>");
+        // BOS token is added by the tokenizer's encode(prompt, add_special_tokens=true)
         
         for message in messages {
             let role = message["role"].as_str().unwrap_or("");
@@ -481,15 +481,18 @@ impl QuantizedGGUF {
                 "assistant" => {
                     prompt.push_str(&format!("<|im_start|>assistant\n{}<|im_end|>\n", content));
                 }
+                "tool" => {
+                    prompt.push_str(&format!("<|im_start|>tool\n{}<|im_end|>\n", content));
+                }
                 _ => {}
             }
         }
-        
+
         // Add generation prompt
         prompt.push_str("<|im_start|>assistant\n");
         Ok(prompt)
     }
-    
+
     fn apply_phi_template(&self, messages: &[serde_json::Value]) -> CandleResult<String> {
         let mut prompt = String::new();
         

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

@@ -128,6 +128,9 @@ impl Qwen {
                 "assistant" => {
                     prompt.push_str(&format!("<|im_start|>assistant\n{}<|im_end|>\n", content));
                 }
+                "tool" => {
+                    prompt.push_str(&format!("<|im_start|>tool\n{}<|im_end|>\n", content));
+                }
                 _ => {}
             }
         }

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

@@ -1,18 +1,59 @@
 use magnus::{function, method, prelude::*, Error, RModule, RArray, RHash, Ruby};
-use candle_transformers::models::bert::{BertModel, Config};
+use candle_transformers::models::bert::{BertModel, Config as BertConfig};
+use candle_transformers::models::xlm_roberta::{
+    XLMRobertaForSequenceClassification, Config as XLMRobertaConfig,
+};
+use candle_transformers::models::debertav2::{
+    DebertaV2Model, DebertaV2ContextPooler, Config as DebertaV2Config,
+};
+use candle_transformers::models::modernbert::{
+    ModernBert, Config as ModernBertConfig,
+};
+use candle_transformers::models::qwen3::{
+    ModelForCausalLM as Qwen3Model, Config as Qwen3Config,
+};
 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::{EncodeInput, Tokenizer};
+use std::cell::RefCell;
 use crate::ruby::{Device, Result};
 use crate::tokenizer::{TokenizerWrapper, loader::TokenizerLoader};
 
+enum RerankerModel {
+    Bert {
+        model: BertModel,
+        pooler: Linear,
+        classifier: Linear,
+    },
+    XLMRoberta {
+        model: XLMRobertaForSequenceClassification,
+        pad_token_id: u32,
+    },
+    DeBERTa {
+        model: DebertaV2Model,
+        pooler: DebertaV2ContextPooler,
+        classifier: Linear,
+        pad_token_id: u32,
+    },
+    ModernBert {
+        model: ModernBert,
+        head_dense: Linear,
+        head_norm: candle_nn::LayerNorm,
+        classifier: Linear,
+        pad_token_id: u32,
+    },
+    Qwen3 {
+        model: RefCell<Qwen3Model>,
+        yes_token_id: u32,
+        no_token_id: u32,
+    },
+}
+
 #[magnus::wrap(class = "Candle::Reranker", free_immediately, size)]
 pub struct Reranker {
-    model: BertModel,
+    model: RerankerModel,
     tokenizer: TokenizerWrapper,
-    pooler: Linear,
-    classifier: Linear,
     device: CoreDevice,
     model_id: String,
 }
@@ -28,7 +69,7 @@ impl Reranker {
         let ruby = Ruby::get().unwrap();
         let runtime_error = ruby.exception_runtime_error();
 
-        let result = (|| -> std::result::Result<(BertModel, TokenizerWrapper, Linear, Linear), Box<dyn std::error::Error + Send + Sync>> {
+        let result = (|| -> std::result::Result<(RerankerModel, TokenizerWrapper), Box<dyn std::error::Error + Send + Sync>> {
             let api = Api::new()?;
             let repo = api.repo(Repo::new(model_id.clone(), RepoType::Model));
 
@@ -37,9 +78,10 @@ impl Reranker {
             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)?;
+            // Read raw config to detect model type
+            let config_str = std::fs::read_to_string(&config_filename)?;
+            let raw_config: serde_json::Value = serde_json::from_str(&config_str)?;
+            let model_type = raw_config["model_type"].as_str().unwrap_or("bert");
 
             // Setup tokenizer with padding AND truncation
             let tokenizer = Tokenizer::from_file(tokenizer_filename)?;
@@ -51,21 +93,71 @@ impl Reranker {
                 VarBuilder::from_mmaped_safetensors(&[weights_filename], DType::F32, &device)?
             };
 
-            // 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"))?;
+            let model = match model_type {
+                "xlm-roberta" => {
+                    let config: XLMRobertaConfig = serde_json::from_str(&config_str)?;
+                    let pad_token_id = config.pad_token_id;
+                    let model = XLMRobertaForSequenceClassification::new(1, &config, vb)?;
+                    RerankerModel::XLMRoberta { model, pad_token_id }
+                }
+                "deberta-v2" => {
+                    let config: DebertaV2Config = serde_json::from_str(&config_str)?;
+                    let pad_token_id = config.pad_token_id.unwrap_or(0) as u32;
+                    let model = DebertaV2Model::load(vb.pp("deberta"), &config)?;
+                    let pooler = DebertaV2ContextPooler::load(vb.clone(), &config)?;
+                    let pooler_hidden_size = config.pooler_hidden_size.unwrap_or(config.hidden_size);
+                    let num_labels = config.id2label.as_ref().map_or(1, |m| m.len());
+                    let classifier = candle_nn::linear(pooler_hidden_size, num_labels, vb.pp("classifier"))?;
+                    RerankerModel::DeBERTa { model, pooler, classifier, pad_token_id }
+                }
+                "qwen3" => {
+                    let config: Qwen3Config = serde_json::from_str(&config_str)?;
+                    let model = Qwen3Model::new(&config, vb)?;
+
+                    // Look up "yes" and "no" token IDs from the tokenizer
+                    let yes_token_id: u32 = tokenizer
+                        .encode("yes", false)
+                        .ok()
+                        .and_then(|enc| enc.get_ids().first().copied())
+                        .unwrap_or(9693);
+                    let no_token_id: u32 = tokenizer
+                        .encode("no", false)
+                        .ok()
+                        .and_then(|enc| enc.get_ids().first().copied())
+                        .unwrap_or(2152);
+
+                    RerankerModel::Qwen3 {
+                        model: RefCell::new(model),
+                        yes_token_id,
+                        no_token_id,
+                    }
+                }
+                "modernbert" => {
+                    let config: ModernBertConfig = serde_json::from_str(&config_str)?;
+                    let pad_token_id = config.pad_token_id;
+                    let model = ModernBert::load(vb.clone(), &config)?;
+                    // ModernBertHead::load is private, so load the head layers manually
+                    let head_vb = vb.pp("head");
+                    let head_dense = candle_nn::linear_no_bias(config.hidden_size, config.hidden_size, head_vb.pp("dense"))?;
+                    let head_norm = candle_nn::layer_norm_no_bias(config.hidden_size, config.layer_norm_eps, head_vb.pp("norm"))?;
+                    let classifier = candle_nn::linear(config.hidden_size, 1, vb.pp("classifier"))?;
+                    RerankerModel::ModernBert { model, head_dense, head_norm, classifier, pad_token_id }
+                }
+                _ => {
+                    let config: BertConfig = serde_json::from_str(&config_str)?;
+                    let model = BertModel::load(vb.pp("bert"), &config)?;
+                    let pooler = candle_nn::linear(config.hidden_size, config.hidden_size, vb.pp("bert.pooler.dense"))?;
+                    let classifier = candle_nn::linear(config.hidden_size, 1, vb.pp("classifier"))?;
+                    RerankerModel::Bert { model, pooler, classifier }
+                }
+            };
 
-            Ok((model, TokenizerWrapper::new(tokenizer), pooler, classifier))
+            Ok((model, TokenizerWrapper::new(tokenizer)))
         })();
 
         match result {
-            Ok((model, tokenizer, pooler, classifier)) => {
-                Ok(Self { model, tokenizer, pooler, classifier, device, model_id })
+            Ok((model, tokenizer)) => {
+                Ok(Self { model, tokenizer, device, model_id })
             }
             Err(e) => Err(Error::new(runtime_error, format!("Failed to load model: {}", e))),
         }
@@ -164,49 +256,161 @@ impl Reranker {
             .map_err(|e| Error::new(runtime_error, format!("Failed to create tensor: {}", e)))?;
         let token_type_ids = Tensor::new(token_type_ids, &self.device)
             .map_err(|e| Error::new(runtime_error, format!("Failed to create token type ids tensor: {}", e)))?;
-        let attention_mask = token_ids.ne(0u32)
-            .map_err(|e| Error::new(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(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)
-                let cls_embeddings = self.extract_cls_embeddings(&embeddings)?;
-                let pooled = self.pooler.forward(&cls_embeddings)
+
+        // Compute scores based on model type
+        let scores = match &self.model {
+            RerankerModel::Bert { model, pooler, classifier } => {
+                let attention_mask = token_ids.ne(0u32)
+                    .map_err(|e| Error::new(runtime_error, format!("Failed to create attention mask: {}", e)))?;
+
+                // Forward pass through BERT
+                let embeddings = model.forward(&token_ids, &token_type_ids, Some(&attention_mask))
+                    .map_err(|e| Error::new(runtime_error, format!("Model forward pass failed: {}", e)))?;
+
+                // Apply pooling based on the specified method
+                let pooled_embeddings = match pooling_method.as_str() {
+                    "pooler" => {
+                        let cls_embeddings = self.extract_cls_embeddings(&embeddings)?;
+                        let pooled = pooler.forward(&cls_embeddings)
+                            .map_err(|e| Error::new(runtime_error, format!("Pooler forward failed: {}", e)))?;
+                        pooled.tanh()
+                            .map_err(|e| Error::new(runtime_error, format!("Tanh activation failed: {}", e)))?
+                    },
+                    "cls" => {
+                        self.extract_cls_embeddings(&embeddings)?
+                    },
+                    "mean" => {
+                        let (_batch, seq_len, _hidden) = embeddings.dims3()
+                            .map_err(|e| Error::new(runtime_error, format!("Failed to get tensor dimensions: {}", e)))?;
+                        let sum = embeddings.sum(1)
+                            .map_err(|e| Error::new(runtime_error, format!("Failed to sum embeddings: {}", e)))?;
+                        (sum / (seq_len as f64))
+                            .map_err(|e| Error::new(runtime_error, format!("Failed to compute mean: {}", e)))?
+                    },
+                    _ => return Err(Error::new(runtime_error,
+                        format!("Unknown pooling method: {}. Use 'pooler', 'cls', or 'mean'", pooling_method)))
+                };
+
+                let pooled_embeddings = pooled_embeddings.contiguous()
+                    .map_err(|e| Error::new(runtime_error, format!("Failed to make pooled_embeddings contiguous: {}", e)))?;
+                let logits = classifier.forward(&pooled_embeddings)
+                    .map_err(|e| Error::new(runtime_error, format!("Classifier forward failed: {}", e)))?;
+                logits.squeeze(1)
+                    .map_err(|e| Error::new(runtime_error, format!("Failed to squeeze tensor: {}", e)))?
+            }
+            RerankerModel::XLMRoberta { model, pad_token_id } => {
+                let attention_mask = token_ids.ne(*pad_token_id)
+                    .map_err(|e| Error::new(runtime_error, format!("Failed to create attention mask: {}", e)))?;
+
+                // XLMRobertaForSequenceClassification returns logits directly
+                let logits = model.forward(&token_ids, &attention_mask, &token_type_ids)
+                    .map_err(|e| Error::new(runtime_error, format!("Model forward pass failed: {}", e)))?;
+                logits.squeeze(1)
+                    .map_err(|e| Error::new(runtime_error, format!("Failed to squeeze tensor: {}", e)))?
+            }
+            RerankerModel::DeBERTa { model, pooler, classifier, pad_token_id } => {
+                let attention_mask = token_ids.ne(*pad_token_id)
+                    .map_err(|e| Error::new(runtime_error, format!("Failed to create attention mask: {}", e)))?;
+
+                // Forward through DeBERTa encoder
+                let encoder_output = model.forward(&token_ids, Some(token_type_ids.clone()), Some(attention_mask))
+                    .map_err(|e| Error::new(runtime_error, format!("Model forward pass failed: {}", e)))?;
+
+                // Pool and classify
+                let pooled = pooler.forward(&encoder_output)
                     .map_err(|e| Error::new(runtime_error, format!("Pooler forward failed: {}", e)))?;
-                pooled.tanh()
-                    .map_err(|e| Error::new(runtime_error, format!("Tanh activation failed: {}", e)))?
-            },
-            "cls" => {
-                // Just use the [CLS] token embeddings directly (no pooler layer)
-                self.extract_cls_embeddings(&embeddings)?
-            },
-            "mean" => {
-                // Mean pooling across all tokens
-                let (_batch, seq_len, _hidden) = embeddings.dims3()
-                    .map_err(|e| Error::new(runtime_error, format!("Failed to get tensor dimensions: {}", e)))?;
-                let sum = embeddings.sum(1)
-                    .map_err(|e| Error::new(runtime_error, format!("Failed to sum embeddings: {}", e)))?;
-                (sum / (seq_len as f64))
-                    .map_err(|e| Error::new(runtime_error, format!("Failed to compute mean: {}", e)))?
-            },
-            _ => return Err(Error::new(runtime_error,
-                format!("Unknown pooling method: {}. Use 'pooler', 'cls', or 'mean'", pooling_method)))
+                let logits = classifier.forward(&pooled)
+                    .map_err(|e| Error::new(runtime_error, format!("Classifier forward failed: {}", e)))?;
+                logits.squeeze(1)
+                    .map_err(|e| Error::new(runtime_error, format!("Failed to squeeze tensor: {}", e)))?
+            }
+            RerankerModel::ModernBert { model, head_dense, head_norm, classifier, pad_token_id } => {
+                let attention_mask = token_ids.ne(*pad_token_id)
+                    .map_err(|e| Error::new(runtime_error, format!("Failed to create attention mask: {}", e)))?;
+                let attention_mask_f32 = attention_mask.to_dtype(DType::F32)
+                    .map_err(|e| Error::new(runtime_error, format!("Failed to convert attention mask: {}", e)))?;
+
+                // Forward through ModernBERT encoder
+                let encoder_output = model.forward(&token_ids, &attention_mask_f32)
+                    .map_err(|e| Error::new(runtime_error, format!("Model forward pass failed: {}", e)))?;
+
+                // CLS pooling, then head (dense + GELU + norm) + classifier
+                let cls = encoder_output.i((.., 0, ..))
+                    .map_err(|e| Error::new(runtime_error, format!("Failed to extract CLS: {}", e)))?
+                    .contiguous()
+                    .map_err(|e| Error::new(runtime_error, format!("Failed to make contiguous: {}", e)))?;
+                let hidden = head_dense.forward(&cls)
+                    .map_err(|e| Error::new(runtime_error, format!("Head dense failed: {}", e)))?;
+                let hidden = hidden.gelu_erf()
+                    .map_err(|e| Error::new(runtime_error, format!("GELU activation failed: {}", e)))?;
+                let hidden = head_norm.forward(&hidden)
+                    .map_err(|e| Error::new(runtime_error, format!("Head norm failed: {}", e)))?;
+                let logits = classifier.forward(&hidden)
+                    .map_err(|e| Error::new(runtime_error, format!("Classifier forward failed: {}", e)))?;
+                logits.squeeze(1)
+                    .map_err(|e| Error::new(runtime_error, format!("Failed to squeeze tensor: {}", e)))?
+            }
+            RerankerModel::Qwen3 { model, yes_token_id, no_token_id } => {
+                // Qwen3 reranker: decoder-based yes/no scoring
+                // Process each document individually (causal LM, not batch encoder)
+                let mut scores_vec: Vec<f32> = Vec::with_capacity(documents.len());
+                let mut model = model.borrow_mut();
+
+                for doc in &documents {
+                    // Build the Qwen3 reranker prompt
+                    let prompt = format!(
+                        "<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be \"yes\" or \"no\".<|im_end|>\n<|im_start|>user\n<Instruct>: Given a web search query, retrieve relevant passages that answer the query\n<Query>: {}\n<Document>: {}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n",
+                        query, doc
+                    );
+
+                    // Tokenize the prompt
+                    let encoding = self.tokenizer.inner().encode(prompt.as_str(), false)
+                        .map_err(|e| Error::new(runtime_error, format!("Tokenization failed: {}", e)))?;
+                    let input_ids: Vec<u32> = encoding.get_ids().to_vec();
+
+                    // Clear KV cache for each document
+                    model.clear_kv_cache();
+
+                    // Forward pass — get logits for the last token position
+                    let input_tensor = Tensor::new(&input_ids[..], &self.device)
+                        .map_err(|e| Error::new(runtime_error, format!("Failed to create tensor: {}", e)))?
+                        .unsqueeze(0)
+                        .map_err(|e| Error::new(runtime_error, format!("Failed to unsqueeze: {}", e)))?;
+
+                    let logits = model.forward(&input_tensor, 0)
+                        .map_err(|e| Error::new(runtime_error, format!("Model forward pass failed: {}", e)))?;
+
+                    // logits shape: [1, 1, vocab_size] → flatten to [vocab_size]
+                    let logits = logits.flatten_all()
+                        .map_err(|e| Error::new(runtime_error, format!("Failed to flatten: {}", e)))?
+                        .to_dtype(DType::F32)
+                        .map_err(|e| Error::new(runtime_error, format!("Failed to convert dtype: {}", e)))?;
+
+                    // Extract yes/no logits and compute score
+                    let yes_logit: f32 = logits.i(*yes_token_id as usize)
+                        .map_err(|e| Error::new(runtime_error, format!("Failed to get yes logit: {}", e)))?
+                        .to_scalar()
+                        .map_err(|e| Error::new(runtime_error, format!("Failed to convert yes logit: {}", e)))?;
+                    let no_logit: f32 = logits.i(*no_token_id as usize)
+                        .map_err(|e| Error::new(runtime_error, format!("Failed to get no logit: {}", e)))?
+                        .to_scalar()
+                        .map_err(|e| Error::new(runtime_error, format!("Failed to convert no logit: {}", e)))?;
+
+                    // softmax over [yes, no] → P(yes)
+                    let max_logit = yes_logit.max(no_logit);
+                    let yes_exp = (yes_logit - max_logit).exp();
+                    let no_exp = (no_logit - max_logit).exp();
+                    let score = yes_exp / (yes_exp + no_exp);
+
+                    scores_vec.push(score);
+                }
+
+                // Build scores tensor for uniform handling below
+                Tensor::new(scores_vec.as_slice(), &self.device)
+                    .map_err(|e| Error::new(runtime_error, format!("Failed to create scores tensor: {}", e)))?
+            }
         };
 
-        // 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(runtime_error, format!("Failed to make pooled_embeddings contiguous: {}", e)))?;
-        let logits = self.classifier.forward(&pooled_embeddings)
-            .map_err(|e| Error::new(runtime_error, format!("Classifier forward failed: {}", e)))?;
-        let scores = logits.squeeze(1)
-            .map_err(|e| Error::new(runtime_error, format!("Failed to squeeze tensor: {}", e)))?;
-
         // Optionally apply sigmoid activation
         let scores = if apply_sigmoid {
             sigmoid(&scores)

data/lib/candle/agent.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Candle
+  class Agent
+    MAX_ITERATIONS = 10
+
+    attr_reader :llm, :tools, :system_prompt, :max_iterations
+
+    def initialize(llm, tools:, system_prompt: nil, max_iterations: MAX_ITERATIONS)
+      @llm = llm
+      @tools = tools
+      @system_prompt = system_prompt
+      @max_iterations = max_iterations
+    end
+
+    def run(user_message, **options)
+      messages = []
+      messages << { role: "system", content: @system_prompt } if @system_prompt
+      messages << { role: "user", content: user_message }
+
+      iterations = 0
+      loop do
+        iterations += 1
+        if iterations > @max_iterations
+          raise AgentMaxIterationsError,
+            "Agent exceeded maximum iterations (#{@max_iterations})"
+        end
+
+        result = @llm.chat_with_tools(messages, tools: @tools, execute: true, **options)
+
+        if result.has_tool_calls?
+          # If the model produced a substantial text answer alongside tool calls,
+          # treat it as a final response (model is done, trailing tool calls are noise).
+          # Strip <think> blocks so they don't count toward the length check.
+          text_without_thinking = result.text_response&.gsub(/<think>.*?<\/think>/m, "")&.strip
+          if text_without_thinking && text_without_thinking.length > 50
+            return AgentResult.new(
+              response: result.text_response,
+              messages: messages,
+              iterations: iterations,
+              tool_calls_made: messages.count { |m| m[:role] == "tool" }
+            )
+          end
+
+          messages << { role: "assistant", content: result.raw_response }
+
+          result.tool_results.each do |tr|
+            tool_name = tr[:tool_call]&.name || "unknown"
+            tool_output = tr[:error] ? "Error: #{tr[:error]}" : JSON.generate(tr[:result])
+            messages << { role: "tool", content: "[#{tool_name}] #{tool_output}" }
+          end
+        else
+          return AgentResult.new(
+            response: result.text_response || result.raw_response,
+            messages: messages,
+            iterations: iterations,
+            tool_calls_made: messages.count { |m| m[:role] == "tool" }
+          )
+        end
+      end
+    end
+  end
+
+  AgentResult = Struct.new(:response, :messages, :iterations, :tool_calls_made, keyword_init: true)
+  AgentMaxIterationsError = Class.new(StandardError)
+end

data/lib/candle/llm.rb

@@ -252,7 +252,7 @@ module Candle
       base_model
     end
     
-    # Simple chat interface for instruction models
+    # Chat interface — always returns a String
     def chat(messages, **options)
       prompt = apply_chat_template(messages)
       generate(prompt, **options)
@@ -263,7 +263,48 @@ module Candle
       prompt = apply_chat_template(messages)
       generate_stream(prompt, **options, &block)
     end
-    
+
+    # Chat with tool calling — always returns a ToolCallResult
+    # Set execute: true to automatically run the tools (default: false)
+    def chat_with_tools(messages, tools:, execute: false, **options)
+      tool_prompt = build_tool_system_prompt(tools)
+      augmented = inject_tool_instructions(messages, tool_prompt)
+
+      raw_response = chat(augmented, **options)
+
+      result = ToolCallParser.parse(raw_response, available_tools: tools)
+
+      if result.has_tool_calls? && execute
+        tool_results = result.tool_calls.map do |tool_call|
+          tool = tools.find { |t| t.name == tool_call.name }
+          unless tool
+            next { tool_call: tool_call, result: nil, error: "Unknown tool: #{tool_call.name}" }
+          end
+
+          begin
+            output = tool.call(tool_call.arguments)
+            { tool_call: tool_call, result: output, error: nil }
+          rescue Exception => e
+            { tool_call: tool_call, result: nil, error: e.message }
+          end
+        end
+
+        ToolCallResult.new(
+          tool_calls: result.tool_calls,
+          tool_results: tool_results,
+          text_response: result.text_response,
+          raw_response: raw_response
+        )
+      else
+        ToolCallResult.new(
+          tool_calls: result.tool_calls,
+          tool_results: [],
+          text_response: result.has_tool_calls? ? result.text_response : raw_response,
+          raw_response: raw_response
+        )
+      end
+    end
+
     # Inspect method for debugging and exploration
     def inspect
       opts = options rescue {}
@@ -354,6 +395,27 @@ module Candle
 
     private
 
+    def build_tool_system_prompt(tools)
+      tool_defs = tools.map { |t| JSON.generate(t.to_tool_definition) }.join("\n\n")
+      "You are a helpful assistant with access to the following tools:\n\n" \
+        "#{tool_defs}\n\n" \
+        "When you need to use a tool, respond with a tool call in the following format:\n" \
+        "<tool_call>\n" \
+        "{\"name\": \"tool_name\", \"arguments\": {\"arg1\": \"value1\"}}\n" \
+        "</tool_call>\n\n" \
+        "If you don't need to use a tool, respond normally with text."
+    end
+
+    def inject_tool_instructions(messages, tool_prompt)
+      msgs = messages.map { |m| m.dup }
+      if msgs.first && msgs.first[:role] == "system"
+        msgs.first[:content] = "#{tool_prompt}\n\n#{msgs.first[:content]}"
+      else
+        msgs.unshift({ role: "system", content: tool_prompt })
+      end
+      msgs
+    end
+
     # Extract JSON content from generated text, handling stop tokens and extra content
     def extract_json_content(text)
       # Remove any content after common stop tokens

data/lib/candle/tool.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module Candle
+  class Tool
+    attr_reader :name, :description, :parameters
+
+    def initialize(name:, description:, parameters: {}, &block)
+      @name = name
+      @description = description
+      @parameters = parameters
+      @callable = block
+    end
+
+    def call(arguments)
+      @callable.call(arguments)
+    end
+
+    def to_tool_definition
+      {
+        "type" => "function",
+        "function" => {
+          "name" => @name,
+          "description" => @description,
+          "parameters" => @parameters
+        }
+      }
+    end
+  end
+
+  ToolCall = Struct.new(:name, :arguments, keyword_init: true)
+
+  ToolCallResult = Struct.new(
+    :tool_calls,
+    :tool_results,
+    :text_response,
+    :raw_response,
+    keyword_init: true
+  ) do
+    def has_tool_calls?
+      tool_calls && !tool_calls.empty?
+    end
+
+    def success?
+      tool_results.all? { |r| r[:error].nil? }
+    end
+  end
+end

data/lib/candle/tool_call_parser.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Candle
+  class ToolCallParser
+    DEFAULT_PATTERN = /<tool_call>\s*(.*?)\s*<\/tool_call>/m
+
+    attr_reader :pattern
+
+    def initialize(pattern: DEFAULT_PATTERN)
+      @pattern = pattern
+    end
+
+    ParseResult = Struct.new(:text_response, :tool_calls, keyword_init: true) do
+      def has_tool_calls?
+        tool_calls && !tool_calls.empty?
+      end
+    end
+
+    def parse(text, available_tools: [])
+      tool_calls = []
+
+      text.scan(@pattern) do |match|
+        json_str = match[0].strip
+        begin
+          parsed = JSON.parse(json_str)
+          name = parsed["name"]
+          arguments = parsed["arguments"] || parsed["parameters"] || {}
+
+          next unless name
+          if available_tools.empty? || available_tools.any? { |t| t.name == name }
+            tool_calls << ToolCall.new(name: name, arguments: arguments)
+          end
+        rescue JSON::ParserError
+          # Skip malformed tool calls
+        end
+      end
+
+      # Deduplicate identical tool calls (models sometimes repeat the same call)
+      tool_calls.uniq! { |tc| [tc.name, tc.arguments] }
+
+      remaining_text = text.gsub(@pattern, "").strip
+      remaining_text = nil if remaining_text.empty?
+
+      ParseResult.new(
+        text_response: remaining_text,
+        tool_calls: tool_calls
+      )
+    end
+
+    # Convenience class method using the default pattern
+    def self.parse(text, available_tools: [])
+      new.parse(text, available_tools: available_tools)
+    end
+  end
+end

data/lib/candle/version.rb

@@ -1,5 +1,5 @@
 # :nocov:
 module Candle
-  VERSION = "1.5.0"
+  VERSION = "1.6.1"
 end
 # :nocov:

data/lib/candle.rb

@@ -5,7 +5,10 @@ require_relative "candle/device_utils"
 require_relative "candle/embedding_model_type"
 require_relative "candle/embedding_model"
 require_relative "candle/reranker"
+require_relative "candle/tool"
+require_relative "candle/tool_call_parser"
 require_relative "candle/llm"
+require_relative "candle/agent"
 require_relative "candle/tokenizer"
 require_relative "candle/ner"
 require_relative "candle/build_info"

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: red-candle
 version: !ruby/object:Gem::Version
-  version: 1.5.0
+  version: 1.6.1
 platform: ruby
 authors:
 - Christopher Petersen
@@ -254,6 +254,7 @@ files:
 - ext/candle/tests/device_tests.rs
 - ext/candle/tests/tensor_tests.rs
 - lib/candle.rb
+- lib/candle/agent.rb
 - lib/candle/build_info.rb
 - lib/candle/device_utils.rb
 - lib/candle/embedding_model.rb
@@ -264,6 +265,8 @@ files:
 - lib/candle/reranker.rb
 - lib/candle/tensor.rb
 - lib/candle/tokenizer.rb
+- lib/candle/tool.rb
+- lib/candle/tool_call_parser.rb
 - lib/candle/version.rb
 - lib/red-candle.rb
 homepage: https://github.com/scientist-labs/red-candle