- [Sentence embeddings](https://github.com/pgvector/pgvector-rust/blob/master/examples/candle/src/main.rs) with Candle

[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"

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>()?)

}