Embeddings and vector similarity

Step 3 of 8. Step 2 ended with a candle project that compiles and runs on the CPU but deliberately touches no model — the runtime, before the weights. This step loads the first real model into it: all-MiniLM-L6-v2, the sentence encoder fixed in step 1, which turns any sentence into 384 numbers. By the end you have those numbers coming out of a real forward pass and the cosine similarity between every pair of sentences — the retrieval primitive the whole service is built on. The step ends at tag llm-ch03.

What an embedding is

A sentence encoder is a neural network with a fixed output: feed it text, get back a vector of a fixed length — 384 for this model — chosen so that sentences with similar meanings point in similar directions. “How do I run a model on my own machine?” and “what is the fastest way to host a model locally?” end up pointing almost the same way; “add a column to a SQLite table” points somewhere else entirely.

Direction, not distance, is the comparison that matters, and the standard measure of “same direction” is the cosine similarity: the dot product of the two vectors divided by both their lengths. Identical directions score 1.0, unrelated ones score around 0.0, opposites score −1.0. That is the entire arithmetic a RAG service uses to decide which document chunk is relevant to a question — and after this step it is arithmetic you have run on real vectors, not read about.

The dependencies

Step 2’s Cargo.toml had exactly one dependency. This step needs the rest of the stack: the model definitions, the tokenizer, and the Hub client that downloads the weights on first run.

[package]
name = "llm-apps"
version = "0.1.0"
edition = "2024"

[dependencies]
anyhow = "1.0.104"
# The runtime the whole track runs on. `default-features = false` keeps the
# CPU path only - no Metal, no CUDA - which is everything this track needs
# and the shortest dependency tree candle has.
candle-core = { version = "0.11.0", default-features = false }
candle-nn = "0.11.0"
candle-transformers = "0.11.0"
hf-hub = "0.3.2"
serde_json = "1.0.154"
tokenizers = "0.23.2"

The three candle crates split the work the way candle itself does. candle-core is the tensors and the CPU device; candle-nn provides the layer primitives (VarBuilder, which loads weights into them); and candle-transformers holds pre-written model architectures — its bert module is a from-scratch Rust implementation of BERT that this model is a variant of. tokenizers is Hugging Face’s tokenizer, hf-hub fetches model files from the Hub, and anyhow catches the mixed error types those crates throw without forcing a wrapper per call.

The whole program

Ninety-odd lines load the model, embed three sentences, and print the cosine similarity between each pair. The file in full, since this step replaces step 2’s ten-liner entirely:

//! Step 3: real embeddings, and the distance between them.
//!
//! Loads `all-MiniLM-L6-v2` (a 384-dimension sentence encoder), embeds three
//! sentences, and prints the cosine similarity between every pair - the
//! retrieval primitive the whole track is built on.

use anyhow::Result;
use candle_core::{Device, Tensor};
use candle_nn::VarBuilder;
use candle_transformers::models::bert::{BertModel, Config, DTYPE};
use hf_hub::api::sync::Api;
use tokenizers::Tokenizer;

const MODEL_ID: &str = "sentence-transformers/all-MiniLM-L6-v2";

const SENTENCES: [&str; 3] = [
    "How do I run a language model on my own machine?",
    "What is the fastest way to host a model locally?",
    "Add a new column to an existing SQLite table.",
];

fn main() -> Result<()> {
    let device = Device::Cpu;
    let started = std::time::Instant::now();

    let api = Api::new()?;
    let repo = api.model(MODEL_ID.to_string());
    let config: Config = serde_json::from_str(&std::fs::read_to_string(repo.get("config.json")?)?)?;
    let tokenizer = Tokenizer::from_file(repo.get("tokenizer.json")?)
        .map_err(|e| anyhow::anyhow!("{e}"))?;
    let weights = repo.get("model.safetensors")?;
    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[weights], DTYPE, &device)? };
    let model = BertModel::load(vb, &config)?;
    println!("model loaded in {:.1?} on {device:?}", started.elapsed());

    // Tokenize all three sentences into one padded batch.
    let encodings = tokenizer
        .encode_batch(SENTENCES.to_vec(), true)
        .map_err(|e| anyhow::anyhow!("{e}"))?;
    let len = encodings.iter().map(|e| e.get_ids().len()).max().unwrap();
    let token_ids: Vec<Vec<u32>> = encodings
        .iter()
        .map(|e| {
            let mut ids = e.get_ids().to_vec();
            ids.resize(len, 0);
            ids
        })
        .collect();
    let attention: Vec<Vec<u32>> = encodings
        .iter()
        .map(|e| {
            let mut mask = e.get_attention_mask().to_vec();
            mask.resize(len, 0);
            mask
        })
        .collect();
    let types = vec![vec![0u32; len]; SENTENCES.len()];

    let ids = Tensor::new(token_ids, &device)?;
    let mask = Tensor::new(attention.clone(), &device)?;
    let types = Tensor::new(types, &device)?;

    // Forward pass: one embedding vector per token, [batch, seq, 384].
    let hidden = model.forward(&ids, &types, Some(&mask))?.to_vec3::<f32>()?;

    // Mean-pool the token vectors over real tokens (padding contributes
    // nothing): one plain Vec<f32> per sentence.
    fn sentence_vector(tokens: &[Vec<f32>], mask: &[u32]) -> Vec<f32> {
        let mut v = vec![0f32; tokens[0].len()];
        let mut n = 0f32;
        for (t, &keep) in tokens.iter().zip(mask) {
            if keep == 1 {
                n += 1.0;
                for (v, &x) in v.iter_mut().zip(t) {
                    *v += x;
                }
            }
        }
        v.iter_mut().for_each(|x| *x /= n);
        v
    }
    let vectors: Vec<Vec<f32>> = hidden
        .iter()
        .zip(attention.iter())
        .map(|(tokens, mask)| sentence_vector(tokens, mask))
        .collect();

    let (rows, dim) = (vectors.len(), vectors[0].len());
    println!("each sentence is now {dim} numbers; sentence 0 starts with:");
    for v in vectors[0].iter().take(4) {
        println!("  {v:.6}");
    }

    // Cosine similarity: dot product divided by both lengths. The vectors
    // stay ordinary floats from here on - no more tensors.
    fn cosine(a: &[f32], b: &[f32]) -> f32 {
        let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
        let len_a = a.iter().map(|x| x * x).sum::<f32>().sqrt();
        let len_b = b.iter().map(|x| x * x).sum::<f32>().sqrt();
        dot / (len_a * len_b)
    }
    for i in 0..rows {
        for j in (i + 1)..rows {
            println!("s{i} vs s{j}: {:.4}", cosine(&vectors[i], &vectors[j]));
        }
    }

    Ok(())
}

Run it with cargo run. The first run downloads the model’s three files — config.json (the layer sizes), tokenizer.json (the vocabulary), and model.safetensors, the ~90 MB of weights — into ~/.cache/huggingface/, where every later run finds them already on disk. On this machine the second run took:

model loaded in 80.1ms on Cpu
each sentence is now 384 numbers; sentence 0 starts with:
  0.229124
  -0.718850
  -0.117356
  -0.303682
s0 vs s1: 0.4523
s0 vs s2: -0.0014
s1 vs s2: 0.0002

(The load time is whatever the disk hands back; the numbers below it are deterministic — the same weights and the same sentences produce the same vectors every run.)

What the code does, move by move

Fetch the files. Api::new() points at the Hub, api.model(MODEL_ID) at the repository, and repo.get(...) returns a local path — downloading the file if it is not already in the cache. Note that device from step 2 appears again in the same role: a value passed into the loader, Device::Cpu by the track’s contract.

Tokenize a batch. encode_batch turns all three sentences into token ids in one call. Sentences have different lengths, so the batch is padded to the longest one, and each row of the attention mask records which positions are real tokens and which are padding — the model must not read the padding, and neither must the pooling step.

Forward. BertModel::forward returns one 384-number vector per token — shape [3, 128, 384] here, 128 being the longest sentence’s token count. That is what a transformer’s output layer actually is; a sentence vector is a convention applied on top.

Mean-pool, then compare. The convention this model’s authors chose is mean pooling: average the per-token vectors over the real (non-padding) tokens, giving one Vec<f32> of 384 numbers per sentence. Everything after that is plain Rust — a dot product and two lengths, exactly the cosine formula, on ordinary floats. No tensor library is needed for similarity, and that is worth seeing once before step 5 pushes thousands of vectors through it.

Reading the numbers

s0 and s1 are two wordings of the same intent — running a model locally — and they score 0.4523. s0 and s2 share no intent, and score −0.0014: essentially zero, two vectors at right angles. That contrast is retrieval. At indexing time the service embeds every document chunk once; at question time it embeds the question and takes the chunks with the highest cosine similarity. Steps 4 and 5 build exactly that loop — chunking first, storage after — and the ranker at the bottom of it is the ten-line cosine function above.

Two details are worth internalising now. First, the diagonal of the similarity matrix is 1.0 by construction — a vector’s cosine with itself is its length squared over both lengths squared — so the program only prints the off-diagonal pairs. Second, the raw embedding values are small and carry no unit: they are coordinates in a learned space, and only their relations (comparing vectors to each other) mean anything. The first four numbers of sentence 0 printed above are not features a person can read; they are addresses.

The tag

The step’s end state is tagged in the examples repository:

git checkout llm-ch03

At that tag the project loads a real model on the CPU and computes cosine similarities between real sentence vectors — the retrieval primitive, working end to end. Step 4 adds the other half of the index side: splitting documents into chunks worth embedding, and what overlap costs and buys. The track overview keeps the full list current as articles land.

Running a model locally — the previous step, where the project and its device were set up
What we’re building: a local RAG service — the map, if this is your first step in the track