red-candle 1.8.1 → 1.8.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: e2d45880b82de259bacdef9c2ce582a3d98da006966e3c211b5f2c63a76a474a
-  data.tar.gz: 2f239dbfc74f6c924fba41b7ed53437416459ceed6bfa2b6d5501d82a778c071
+  metadata.gz: e7b7ecd1fe8f4551498e553e005c09ffb75f858b107525cbc8bd77948a040ff6
+  data.tar.gz: cf84d7690cfdc13eaf210745b38e73bbf72806a4f94be2c1a0315d1d3827ff8d
 SHA512:
-  metadata.gz: 18f03813074b73b797a82f9d2917441b701e920a3e2a1afef8b7fb4e863ebbf2acbda06decb90a9014766fbaef4554f0e88e66744fe28e86d4e489e5a7db8dcf
-  data.tar.gz: 97482fdd896803f4bb3a4a8142ad04995f9a25b53c8876979c77e392881ae43b72f497797501f5fc24323108b1edb5a6206933e0a8eb30ba02879f393df4a507
+  metadata.gz: 8f33e7cad2a4678a77f54c6997b31ebd3880a5916a473f8f15bdb58b8bfe78960077efa713ebab265dfa25acc0bd28bd952bb8a7cd48549593f6aee107cbc887
+  data.tar.gz: f4ba719c4f721fa23155bf11df9bd779d423eba03676ba5c3aa1c92c29e0551455267b07bc6f15b2996cd2a8592cdf7c471aed6d11d1706c0f5371de054f739d

data/README.md

@@ -107,6 +107,10 @@ device = Candle::Device.cpu     # or Candle::Device.metal, Candle::Device.cuda
 model = Candle::EmbeddingModel.from_pretrained("jinaai/jina-embeddings-v2-base-en", device: device)
 embedding = model.embedding("Hi there!")
 
+# Cap input length (in tokens) to cut memory; longer inputs are truncated.
+# JinaBERT otherwise reserves ~3.2 GB for its full 8192-token context.
+model = Candle::EmbeddingModel.from_pretrained(max_length: 512)
+
 # Reranker also supports device selection
 reranker = Candle::Reranker.from_pretrained("cross-encoder/ms-marco-MiniLM-L-12-v2", device: device)
 results = reranker.rerank("query", ["doc1", "doc2", "doc3"])

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

@@ -72,10 +72,11 @@ pub struct EmbeddingModelInner {
     model: Option<EmbeddingModelVariant>,
     tokenizer: Option<TokenizerWrapper>,
     embedding_size: Option<usize>,
+    max_length: Option<usize>,
 }
 
 impl EmbeddingModel {
-    pub fn new(model_id: Option<String>, tokenizer: Option<String>, device: Option<Device>, model_type: Option<String>, embedding_size: Option<usize>) -> Result<Self> {
+    pub fn new(model_id: Option<String>, tokenizer: Option<String>, device: Option<Device>, model_type: Option<String>, embedding_size: Option<usize>, max_length: Option<usize>) -> Result<Self> {
         let device = device.unwrap_or(Device::best()).as_device()?;
         let model_type = model_type
             .and_then(|mt| EmbeddingModelType::from_string(&mt))
@@ -86,14 +87,15 @@ impl EmbeddingModel {
             tokenizer_id: tokenizer.clone(),
             model_type: Some(model_type),
             model: match model_id.as_ref() {
-                Some(id) => Some(Self::build_embedding_model(id, device, model_type, embedding_size)?),
+                Some(id) => Some(Self::build_embedding_model(id, device, model_type, embedding_size, max_length)?),
                 None => None
             },
             tokenizer: match tokenizer {
-                Some(tid) => Some(Self::build_tokenizer(tid)?),
+                Some(tid) => Some(Self::build_tokenizer(tid, max_length)?),
                 None => None
             },
             embedding_size,
+            max_length,
         }))
     }
 
@@ -177,7 +179,7 @@ impl EmbeddingModel {
         }
     }
 
-    fn build_embedding_model(model_id: &str, device: CoreDevice, model_type: EmbeddingModelType, embedding_size: Option<usize>) -> Result<EmbeddingModelVariant> {
+    fn build_embedding_model(model_id: &str, device: CoreDevice, model_type: EmbeddingModelType, embedding_size: Option<usize>, max_length: Option<usize>) -> Result<EmbeddingModelVariant> {
         use hf_hub::{api::sync::Api, Repo, RepoType};
         let api = Api::new().map_err(wrap_hf_err)?;
         let repo = Repo::new(model_id.to_string(), RepoType::Model);
@@ -194,6 +196,11 @@ impl EmbeddingModel {
                 let final_emb_dim = Self::resolve_embedding_size(Path::new(&model_path), embedding_size)?;
                 let mut config = JinaConfig::v2_base();
                 config.hidden_size = final_emb_dim;
+                // JinaBERT precomputes an ALiBi bias of heads x len x len floats
+                // (3.2 GB at the default 8192), so size it to the inputs we accept.
+                if let Some(max_length) = max_length {
+                    config.max_position_embeddings = max_length;
+                }
                 let vb = unsafe {
                     VarBuilder::from_mmaped_safetensors(&[model_path], CoreDType::F32, &device)
                         .map_err(wrap_candle_err)?
@@ -268,7 +275,7 @@ impl EmbeddingModel {
         }
     }
 
-    fn build_tokenizer(tokenizer_id: String) -> Result<TokenizerWrapper> {
+    fn build_tokenizer(tokenizer_id: String, max_length: Option<usize>) -> Result<TokenizerWrapper> {
         use hf_hub::{api::sync::Api, Repo, RepoType};
         let tokenizer_path = Api::new()
                 .map_err(wrap_hf_err)?
@@ -282,6 +289,10 @@ impl EmbeddingModel {
             .map_err(wrap_std_err)?;
         
         let tokenizer = TokenizerLoader::with_padding(tokenizer, None);
+        let tokenizer = match max_length {
+            Some(max_length) => TokenizerLoader::with_truncation(tokenizer, max_length),
+            None => tokenizer,
+        };
         Ok(TokenizerWrapper::new(tokenizer))
     }
 
@@ -451,6 +462,11 @@ impl EmbeddingModel {
         if let Some(size) = self.0.embedding_size {
             hash.aset("embedding_size", size)?;
         }
+
+        // Add max_length
+        if let Some(max_length) = self.0.max_length {
+            hash.aset("max_length", max_length)?;
+        }
         
         Ok(hash)
     }
@@ -459,7 +475,7 @@ impl EmbeddingModel {
 pub fn init(rb_candle: RModule) -> Result<()> {
     let ruby = Ruby::get().unwrap();
     let rb_embedding_model = rb_candle.define_class("EmbeddingModel", ruby.class_object())?;
-    rb_embedding_model.define_singleton_method("_create", function!(EmbeddingModel::new, 5))?;
+    rb_embedding_model.define_singleton_method("_create", function!(EmbeddingModel::new, 6))?;
     // Expose embedding with an optional pooling_method argument (default: "pooled")
     rb_embedding_model.define_method("_embedding", method!(EmbeddingModel::embedding, 2))?;
     rb_embedding_model.define_method("embeddings", method!(EmbeddingModel::embeddings, 1))?;

data/lib/candle/embedding_model.rb

@@ -15,8 +15,9 @@ module Candle
     # @param tokenizer [String, nil] The tokenizer to use (defaults to using the model's tokenizer)
     # @param model_type [String, nil] The type of embedding model (auto-detected if nil)
     # @param embedding_size [Integer, nil] Override for the embedding size (optional)
+    # @param max_length [Integer, nil] Longest input in tokens; longer inputs are truncated (optional)
     # @return [EmbeddingModel] A new EmbeddingModel instance
-    def self.from_pretrained(model_id = DEFAULT_MODEL_PATH, device: Candle::Device.best, tokenizer: nil, model_type: nil, embedding_size: nil)
+    def self.from_pretrained(model_id = DEFAULT_MODEL_PATH, device: Candle::Device.best, tokenizer: nil, model_type: nil, embedding_size: nil, max_length: nil)
       # Auto-detect model type based on model_id if not provided
       if model_type.nil?
         model_type = case model_id.downcase
@@ -34,7 +35,7 @@ module Candle
       # Use model_id as tokenizer if not specified (usually what you want)
       tokenizer_id = tokenizer || model_id
       
-      _create(model_id, tokenizer_id, device, model_type, embedding_size)
+      _create(model_id, tokenizer_id, device, model_type, embedding_size, max_length)
     end
     
     # Constructor for creating a new EmbeddingModel with optional parameters
@@ -44,13 +45,15 @@ module Candle
     # @param device [Candle::Device, Candle::Device.cpu] The device to use for computation
     # @param model_type [String, nil] The type of embedding model to use
     # @param embedding_size [Integer, nil] Override for the embedding size (optional)
+    # @param max_length [Integer, nil] Longest input in tokens; longer inputs are truncated (optional)
     def self.new(model_path: DEFAULT_MODEL_PATH,
       tokenizer_path: DEFAULT_TOKENIZER_PATH,
       device: Candle::Device.best,
       model_type: DEFAULT_EMBEDDING_MODEL_TYPE,
-      embedding_size: nil)
+      embedding_size: nil,
+      max_length: nil)
       $stderr.puts "[DEPRECATION] `EmbeddingModel.new` is deprecated. Please use `EmbeddingModel.from_pretrained` instead."
-      _create(model_path, tokenizer_path, device, model_type, embedding_size)
+      _create(model_path, tokenizer_path, device, model_type, embedding_size, max_length)
     end
     # Returns the embedding for a string using the specified pooling method.
     # @param str [String] The input text

data/lib/candle/version.rb

@@ -1,5 +1,5 @@
 # :nocov:
 module Candle
-  VERSION = "1.8.1"
+  VERSION = "1.8.2"
 end
 # :nocov:

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: red-candle
 version: !ruby/object:Gem::Version
-  version: 1.8.1
+  version: 1.8.2
 platform: ruby
 authors:
 - Christopher Petersen