Added Candle example [skip ci]

Author: ankane

README.md

@@ -18,6 +18,7 @@ Or check out some examples:
 
 - [Embeddings](https://github.com/pgvector/pgvector-rust/blob/master/examples/openai/src/main.rs) with OpenAI
 - [Binary embeddings](https://github.com/pgvector/pgvector-rust/blob/master/examples/cohere/src/main.rs) with Cohere
+- [Sentence embeddings](https://github.com/pgvector/pgvector-rust/blob/master/examples/candle/src/main.rs) with Candle
 - [Recommendations](https://github.com/pgvector/pgvector-rust/blob/master/examples/disco/src/main.rs) with Disco
 - [Bulk loading](https://github.com/pgvector/pgvector-rust/blob/master/examples/loading/src/main.rs) with `COPY`
 

examples/candle/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "example"
+version = "0.1.0"
+edition = "2021"
+publish = false
+
+[dependencies]
+candle-core = "0.6"
+candle-nn = "0.6"
+candle-transformers = "0.6"
+hf-hub = "0.3"
+pgvector = { path = "../..", features = ["postgres"] }
+postgres = "0.19"
+serde_json = "1"
+tokenizers = "0.19"

examples/candle/src/main.rs

@@ -0,0 +1,95 @@
+// https://github.com/huggingface/candle/tree/main/candle-examples/examples/bert
+// https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2
+
+use candle_core::{Device, Tensor};
+use candle_nn::VarBuilder;
+use candle_transformers::models::bert::{BertModel, Config, DTYPE};
+use hf_hub::api::sync::Api;
+use pgvector::Vector;
+use postgres::{Client, NoTls};
+use std::error::Error;
+use std::fs::read_to_string;
+use tokenizers::{PaddingParams, PaddingStrategy, Tokenizer};
+
+fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
+    let mut client = Client::configure()
+        .host("localhost")
+        .dbname("pgvector_example")
+        .user(std::env::var("USER")?.as_str())
+        .connect(NoTls)?;
+
+    client.execute("CREATE EXTENSION IF NOT EXISTS vector", &[])?;
+    client.execute("DROP TABLE IF EXISTS documents", &[])?;
+    client.execute(
+        "CREATE TABLE documents (id serial PRIMARY KEY, content text, embedding vector(384))",
+        &[],
+    )?;
+
+    let model = EmbeddingModel::new("sentence-transformers/all-MiniLM-L6-v2")?;
+
+    let input = [
+        "The dog is barking",
+        "The cat is purring",
+        "The bear is growling",
+    ];
+    let embeddings = input
+        .iter()
+        .map(|text| model.embed(text))
+        .collect::<Result<Vec<_>, _>>()?;
+
+    for (content, embedding) in input.iter().zip(embeddings) {
+        client.execute(
+            "INSERT INTO documents (content, embedding) VALUES ($1, $2)",
+            &[&content, &Vector::from(embedding)],
+        )?;
+    }
+
+    let document_id = 2;
+    for row in client.query("SELECT content FROM documents WHERE id != $1 ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = $1) LIMIT 5", &[&document_id])? {
+        let content: &str = row.get(0);
+        println!("{}", content);
+    }
+
+    Ok(())
+}
+
+struct EmbeddingModel {
+    tokenizer: Tokenizer,
+    model: BertModel,
+}
+
+impl EmbeddingModel {
+    pub fn new(model_id: &str) -> Result<Self, Box<dyn Error + Send + Sync>> {
+        let api = Api::new()?;
+        let repo = api.model(model_id.to_string());
+        let tokenizer_path = repo.get("tokenizer.json")?;
+        let config_path = repo.get("config.json")?;
+        let weights_path = repo.get("model.safetensors")?;
+
+        let mut tokenizer = Tokenizer::from_file(tokenizer_path)?;
+        let padding = PaddingParams {
+            strategy: PaddingStrategy::BatchLongest,
+            ..Default::default()
+        };
+        tokenizer.with_padding(Some(padding));
+
+        let device = Device::Cpu;
+        let config: Config = serde_json::from_str(&read_to_string(config_path)?)?;
+        let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[weights_path], DTYPE, &device)? };
+        let model = BertModel::load(vb, &config)?;
+
+        Ok(Self { tokenizer, model })
+    }
+
+    // embed one at a time since BertModel does not support attention mask
+    // https://github.com/huggingface/candle/issues/1798
+    fn embed(&self, text: &str) -> Result<Vec<f32>, Box<dyn Error + Send + Sync>> {
+        let tokens = self.tokenizer.encode(text, true)?;
+        let token_ids = Tensor::new(vec![tokens.get_ids().to_vec()], &self.model.device)?;
+        let token_type_ids = token_ids.zeros_like()?;
+        let embeddings = self.model.forward(&token_ids, &token_type_ids)?;
+        let embeddings = (embeddings.sum(1)? / (embeddings.dim(1)? as f64))?;
+        let embeddings = embeddings.broadcast_div(&embeddings.sqr()?.sum_keepdim(1)?.sqrt()?)?;
+        Ok(embeddings.squeeze(0)?.to_vec1::<f32>()?)
+    }
+}