Merge pull request #1 from assaydepot/main

Made Embeddings work

.gitignore

@@ -16,6 +16,9 @@
 # Ignore Byebug command history file.
 .byebug_history
 
+# Ignore the RustRover project file
+.idea
+
 ## Specific to RubyMotion:
 .dat*
 .repl_history
@@ -64,4 +67,4 @@ target
 *.o
 *.lock
 
-lib.py.rs
+lib.py.rs

README.md

@@ -18,6 +18,50 @@ x = x.reshape([3, 2])
 # Tensor[[3, 2], f32]
 ```
 
+```ruby
+require 'candle'
+model = Candle::Model.new
+embedding = model.embedding("Hi there!")
+```
+
+## A note on memory usage
+The `Candle::Model` defaults to the `jinaai/jina-embeddings-v2-base-en` model with the `sentence-transformers/all-MiniLM-L6-v2` tokenizer (both from [HuggingFace](https://huggingface.co)). With this configuration the model takes a little more than 3GB of memory running on my Mac. The memory stays with the instantiated `Candle::Model` class, if you instantiate more than one, you'll use more memory. Likewise, if you let it go out of scope and call the garbage collector, you'll free the memory. For example:
+
+```ruby
+> require 'candle'
+# Ruby memory = 25.9 MB
+> model = Candle::Model.new
+# Ruby memory = 3.50 GB
+> model2 = Candle::Model.new
+# Ruby memory = 7.04 GB
+> model2 = nil
+> GC.start
+# Ruby memory = 3.56 GB
+> model = nil
+> GC.start
+# Ruby memory = 55.2 MB
+```
+
+## A note on returned embeddings
+
+The code should match the same embeddings when generated from the python `transformers` library. For instance, locally I was able to generate the same embedding for the text "Hi there!" using the python code:
+
+```python
+from transformers import AutoModel
+model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True)
+sentence = ['Hi there!']
+embedding = model.encode(sentence)
+print(embedding)
+```
+
+And the following ruby:
+
+```ruby
+require 'candle'
+model = Candle::Model.new
+embedding = model.embedding("Hi there!")
+```
+
 ## Development
 
 FORK IT!
@@ -29,6 +73,7 @@ bundle
 bundle exec rake compile
 ```
 
+
 Implemented with [Magnus](https://github.com/matsadler/magnus), with reference to [Polars Ruby](https://github.com/ankane/polars-ruby)
 
 Policies

ext/candle/src/lib.rs

@@ -1,6 +1,6 @@
 use magnus::{function, method, prelude::*, Ruby};
 
-use crate::model::{candle_utils, ModelConfig, RbDType, RbDevice, RbQTensor, RbResult, RbTensor};
+use crate::model::{candle_utils, RbModel, RbDType, RbDevice, RbQTensor, RbResult, RbTensor};
 
 pub mod model;
 
@@ -22,6 +22,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
  rb_tensor.define_method("dtype", method!(RbTensor::dtype, 0))?;
  rb_tensor.define_method("device", method!(RbTensor::device, 0))?;
  rb_tensor.define_method("rank", method!(RbTensor::rank, 0))?;
+ rb_tensor.define_method("elem_count", method!(RbTensor::elem_count, 0))?;
  rb_tensor.define_method("sin", method!(RbTensor::sin, 0))?;
  rb_tensor.define_method("cos", method!(RbTensor::cos, 0))?;
  rb_tensor.define_method("log", method!(RbTensor::log, 0))?;
@@ -93,10 +94,10 @@ fn init(ruby: &Ruby) -> RbResult<()> {
  rb_qtensor.define_method("dequantize", method!(RbQTensor::dequantize, 0))?;
 
  let rb_model = rb_candle.define_class("Model", Ruby::class_object(ruby))?;
- rb_model.define_singleton_method("new", function!(ModelConfig::new, 0))?;
- rb_model.define_method("embedding", method!(ModelConfig::embedding, 1))?;
- rb_model.define_method("to_s", method!(ModelConfig::__str__, 0))?;
- rb_model.define_method("inspect", method!(ModelConfig::__repr__, 0))?;
+ rb_model.define_singleton_method("new", function!(RbModel::new, 0))?;
+ rb_model.define_method("embedding", method!(RbModel::embedding, 1))?;
+ rb_model.define_method("to_s", method!(RbModel::__str__, 0))?;
+ rb_model.define_method("inspect", method!(RbModel::__repr__, 0))?;
 
  Ok(())
 }

ext/candle/src/model/mod.rs

@@ -1,5 +1,5 @@
-mod config;
-pub use config::*;
+mod rb_model;
+pub use rb_model::*;
 
 mod errors;
 pub use errors::*;

ext/candle/src/model/config.rs → ext/candle/src/model/rb_model.rs

@@ -16,24 +16,24 @@ use crate::model::RbResult;
 use tokenizers::Tokenizer;
 
 #[magnus::wrap(class = "Candle::Model", free_immediately, size)]
-pub struct ModelConfig(pub ModelConfigInner);
+pub struct RbModel(pub RbModelInner);
 
-pub struct ModelConfigInner {
+pub struct RbModelInner {
  device: Device,
  tokenizer_path: Option<String>,
  model_path: Option<String>,
  model: Option<BertModel>,
  tokenizer: Option<Tokenizer>,
 }
 
-impl ModelConfig {
+impl RbModel {
  pub fn new() -> RbResult<Self> {
  Self::new2(Some("jinaai/jina-embeddings-v2-base-en".to_string()), Some("sentence-transformers/all-MiniLM-L6-v2".to_string()), Some(Device::Cpu))
  }
 
  pub fn new2(model_path: Option<String>, tokenizer_path: Option<String>, device: Option<Device>) -> RbResult<Self> {
  let device = device.unwrap_or(Device::Cpu);
- Ok(ModelConfig(ModelConfigInner {
+ Ok(RbModel(RbModelInner {
  device: device.clone(),
  model_path: model_path.clone(),
  tokenizer_path: tokenizer_path.clone(),
@@ -92,10 +92,14 @@ impl ModelConfig {
  ))
  .get("tokenizer.json")
  .map_err(wrap_hf_err)?;
- let tokenizer = tokenizers::Tokenizer::from_file(tokenizer_path)
- // .with_padding(None)
- // .with_truncation(None)
+ let mut tokenizer = tokenizers::Tokenizer::from_file(tokenizer_path)
  .map_err(wrap_std_err)?;
+ let pp = tokenizers::PaddingParams {
+ strategy: tokenizers::PaddingStrategy::BatchLongest,
+ ..Default::default()
+ };
+ tokenizer.with_padding(Some(pp));
+
  Ok(tokenizer)
  }
 
@@ -105,9 +109,6 @@ impl ModelConfig {
  model: &BertModel,
  tokenizer: &Tokenizer,
  ) -> Result<Tensor, Error> {
- let start: std::time::Instant = std::time::Instant::now();
- // let tokenizer_impl = tokenizer
- // .map_err(wrap_std_err)?;
  let tokens = tokenizer
  .encode(prompt, true)
  .map_err(wrap_std_err)?
@@ -117,16 +118,29 @@ impl ModelConfig {
  .map_err(wrap_candle_err)?
  .unsqueeze(0)
  .map_err(wrap_candle_err)?;
- println!("Loaded and encoded {:?}", start.elapsed());
- let start: std::time::Instant = std::time::Instant::now();
+
+ // let start: std::time::Instant = std::time::Instant::now();
  let result = model.forward(&token_ids).map_err(wrap_candle_err)?;
- // println!("{result}");
- println!("Took {:?}", start.elapsed());
- Ok(result)
+
+ // Apply some avg-pooling by taking the mean embedding value for all tokens (including padding)
+ let (_n_sentence, n_tokens, _hidden_size) = result.dims3()
+ .map_err(wrap_candle_err)?;
+ let sum = result.sum(1)
+ .map_err(wrap_candle_err)?;
+ let embeddings = (sum / (n_tokens as f64))
+ .map_err(wrap_candle_err)?;
+ // let embeddings = Self::normalize_l2(&embeddings).map_err(wrap_candle_err)?;
+
+ Ok(embeddings)
+ }
+
+ #[allow(dead_code)]
+ fn normalize_l2(v: &Tensor) -> Result<Tensor, candle_core::Error> {
+ v.broadcast_div(&v.sqr()?.sum_keepdim(1)?.sqrt()?)
  }
 
  pub fn __repr__(&self) -> String {
- format!("Candle::Model(path={})", self.0.model_path.as_deref().unwrap_or("None"))
+ format!("#<Candle::Model model_path: {} tokenizer_path: {})", self.0.model_path.as_deref().unwrap_or("nil"), self.0.tokenizer_path.as_deref().unwrap_or("nil"))
  }
 
  pub fn __str__(&self) -> String {
@@ -143,14 +157,14 @@ impl ModelConfig {
 
 // #[test]
 // fn test_build_model_and_tokenizer() {
-// let config = super::ModelConfig::build();
+// let config = super::RbModel::build();
 // let (_model, tokenizer) = config.build_model_and_tokenizer().unwrap();
 // assert_eq!(tokenizer.get_vocab_size(true), 30522);
 // }
 
 // #[test]
 // fn test_embedding() {
-// let config = super::ModelConfig::build();
+// let config = super::RbModel::build();
 // // let (_model, tokenizer) = config.build_model_and_tokenizer().unwrap();
 // // assert_eq!(config.embedding("Scientist.com is a marketplace for pharmaceutical services.")?, None);
 // }

ext/candle/src/model/rb_tensor.rs

@@ -40,7 +40,6 @@ impl RbTensor {
  ))
  }
 
- // FXIME: Do not use `to_f64` here.
  pub fn values(&self) -> RbResult<Vec<f64>> {
  let values = self
  .0
@@ -83,6 +82,12 @@ impl RbTensor {
  self.0.rank()
  }
 
+ /// The number of elements stored in this tensor.
+ /// &RETURNS&: int
+ pub fn elem_count(&self) -> usize {
+ self.0.elem_count()
+ }
+
  pub fn __repr__(&self) -> String {
  format!("{}", self.0)
  }

ext/candle/src/model/utils.rs

@@ -71,6 +71,7 @@ pub fn candle_utils(rb_candle: magnus::RModule) -> Result<(), Error> {
 
 /// Applies the Softmax function to a given tensor.#
 /// &RETURNS&: Tensor
+#[allow(dead_code)]
 fn softmax(tensor: RbTensor, dim: i64) -> RbResult<RbTensor> {
  let dim = actual_dim(&tensor, dim).map_err(wrap_candle_err)?;
  let sm = candle_nn::ops::softmax(&tensor.0, dim).map_err(wrap_candle_err)?;
@@ -79,6 +80,7 @@ fn softmax(tensor: RbTensor, dim: i64) -> RbResult<RbTensor> {
 
 /// Applies the Sigmoid Linear Unit (SiLU) function to a given tensor.
 /// &RETURNS&: Tensor
+#[allow(dead_code)]
 fn silu(tensor: RbTensor) -> RbResult<RbTensor> {
  let s = candle_nn::ops::silu(&tensor.0).map_err(wrap_candle_err)?;
  Ok(RbTensor(s))

lib/candle.rb

@@ -1 +1,2 @@
 require_relative 'candle/candle'
+require_relative 'candle/tensor'

lib/candle/tensor.rb

@@ -0,0 +1,17 @@
+module Candle
+ class Tensor
+ include Enumerable
+
+ def each
+ if self.rank == 1
+ self.values.each do |value|
+ yield value
+ end
+ else
+ shape.first.times do |i|
+ yield self[i]
+ end
+ end
+ end
+ end
+end