Evaluating retrieval quality
Step 8 of 8 — the last step of the track. Step 7 ended at tag llm-ch07
with the whole pipeline behind a server: POST /ask in, a stream of
passages and tokens out. Nothing in those seven steps has measured
whether retrieval works. The service answers, and the answers look
right — but “looks right” is how quietly broken retrieval ships. A
chunker that fragments, an index that misses, a corpus that outgrows
brute force: all of it invisible behind a confident answer model,
because the answer model will produce something no matter what the
retriever hands it.
This step adds a way to ask the pipeline a question with a known
answer. llm-apps eval runs a small labelled set — six questions, each
labelled with the document that holds its answer — through the same
search a /ask request takes, and prints recall@1, recall@3 and mean
reciprocal rank next to a brute-force ranking of the same rows, so an
index miss has nowhere to hide. A second pass re-chunks the corpus
under three window sizes and re-ranks in memory, so the chunking trade
from step 4 — “a trade to measure, not a default to trust” — finally
gets measured. No new dependencies; the program has been carrying
everything it needs since step 3:
[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"rusqlite = { version = "0.41.1", features = ["bundled"] }serde_json = "1.0.154"tokenizers = "0.23.2"# Step 7: the server, and the client that talks to the answer model.# reqwest drops its TLS stack - everything this step talks to lives on# 127.0.0.1, and the TLS crates are the heaviest part of its tree.axum = "0.8"bytes = "1"futures-util = "0.3"reqwest = { version = "0.12", default-features = false, features = ["json", "stream"] }serde = { version = "1.0", features = ["derive"] }tokio = { version = "1", features = ["full"] }tokio-stream = "0.1"The labelled set
Section titled “The labelled set”The set is six questions written against the three documents the
corpus has had since step 6. Each entry pairs a question with the
index of the DOCS entry that should be retrieved for it — a
document-level label, not a chunk-level one, because at this corpus
size the honest claim a retrieval system can make is “I found the
right document”, and a label a person can check by re-reading three
paragraphs is worth more than a benchmark nobody re-reads.
/// The labelled set: a question, and the index of the DOCS entry that/// holds its answer. Kept small on purpose - a set whose labels can be/// checked by hand is worth more than a benchmark nobody re-reads.const QUESTIONS: [(&str, usize); 6] = [ ("How much memory does the model use at peak?", 0), ("What window size worked for short technical prose?", 1), ("How does the attention mask treat padding?", 2), ("Why does embedding a whole document hurt?", 0), ("How big is the word-piece vocabulary?", 2), ("What happens to recall when windows are too small?", 1),];Five of the six are questions the labelled document answers in its own words — memory budgeting in document 0, window sizes in document 1, padding and the vocabulary in document 2. The set is not designed to flatter: “How big is the word-piece vocabulary?” is the kind of short, size-flavoured question that retrieval on small corpora tends to mis-rank, and it stays in the set whether it passes or not. What the numbers say is decided by the run, not by the author.
Ranking through the server’s own path
Section titled “Ranking through the server’s own path”The one change to existing code: a search result now carries its row’s
chunk id, which the evaluation maps back to the document the chunk
came from. Retrieval itself is untouched — same query, same clusters,
same probe count a /ask request uses.
/// One search result: its cosine score against the question, the/// chunk's full text, the decoded vector - kept alongside the text/// because the dedupe pass needs it - and the row's chunk id, which/// the evaluation maps back to the document it came from.#[derive(Clone)]struct Hit { score: f32, text: String, vector: Vec<f32>, id: i64,}Mapping an id back to a document is bookkeeping the store’s own layout
defines: real chunks carry ids 0..real in chunking order, and the
near-copies of chunk i carry ids real + i * COPIES + j, the same
layout step 6 inserted them with. Copies inherit their original’s
document — a near-copy of a relevant chunk is a hit, not a
near-miss:
/// Synthetic near-copies per real chunk, as in step 6.const COPIES: usize = 3;
/// Which DOCS entry a chunk row came from. Real chunks carry ids/// `0..real` in chunking order; the near-copies of chunk `i` carry ids/// `real + i * COPIES + j`, same layout step 6 inserted them with.fn doc_of_id(id: i64, real: usize, chunks: &[(usize, String)]) -> usize { if (id as usize) < real { chunks[id as usize].0 } else { chunks[(id as usize - real) / COPIES].0 }}
/// Rank of the first hit that comes from the labelled document,/// 1-based; None when the ranking never reaches it.fn rank_of_label( ranked: &[(f32, i64)], label: usize, real: usize, chunks: &[(usize, String)],) -> Option<usize> { ranked .iter() .position(|&(_, id)| doc_of_id(id, real, chunks) == label) .map(|r| r + 1)}
/// `-` reads better than `>15` in a table column that is three/// characters wide.fn fmt_rank(rank: Option<usize>) -> String { match rank { Some(r) => r.to_string(), None => "-".to_string(), }}Two metrics carry the summary. Recall@k asks whether the labelled document appears anywhere in the top k — a pass/fail the reader can recompute by eye. Mean reciprocal rank is sharper about where the hit lands: the reciprocal of the first relevant rank, averaged over the set, so a hit at rank 1 counts ten times a hit at rank 10.
fn main() { // Two rankings of one question's hits: (score, from the labelled document). let hit_at_1 = vec![(0.83, true), (0.71, false), (0.64, false)]; let hit_at_3 = vec![(0.79, false), (0.75, false), (0.70, true)];
let mrr = |ranked: &[(f64, bool)]| -> f64 { ranked .iter() .position(|&(_, hit)| hit) .map(|r| 1.0 / (r + 1) as f64) .unwrap_or(0.0) }; println!("hit at rank 1: MRR {:.2}", mrr(&hit_at_1)); println!("hit at rank 3: MRR {:.2}", mrr(&hit_at_3));}Run this in the Rust Playground
hit at rank 1: MRR 1.00hit at rank 3: MRR 0.33The evaluation itself — eval and eval_window in the full program
below — embeds all six questions in one batch, ranks each two ways,
and prints one line per question. The indexed column walks
search_index, the exact path a /ask request takes: centroids
ranked, two clusters probed, candidates capped at 15. The brute
column scores every row in the store and sorts. When the two disagree,
the index is what missed — never the corpus:
model loaded in 83.9ms on Cpustored 13 chunk vectors from 3 documents, 3 near-copies each (52 rows)index built: 8 clusters, sizes [2, 2, 4, 4, 4, 8, 12, 16]question doc indexed bruteHow much memory does the model use at peak? 0 1 1What window size worked for short technical prose? 1 1 1How does the attention mask treat padding? 2 1 1Why does embedding a whole document hurt? 0 1 1How big is the word-piece vocabulary? 2 5 5What happens to recall when windows are too small? 1 1 1recall@1 5/6, recall@3 5/6, MRR 0.87 over the indexed searchWhat the numbers say
Section titled “What the numbers say”Five of six questions rank the labelled document first. The miss is the one the set kept in on purpose: “How big is the word-piece vocabulary?” lands at rank 5. The two columns agree — 5 indexed, 5 brute — so this is not the index dropping a cluster; the corpus genuinely scores four rows above the labelled one. On these six questions, probing 2 of 8 clusters lost nothing; that agreement is a property of a 52-row store, not a guarantee, and the harness exists so the same question can be asked again when the corpus is real.
MRR 0.87 is the set’s one miss written out: five hits at rank 1 and one at rank 5 give (5 × 1 + 1/5) / 6. Recall@3’s 5/6 says the same thing with less resolution.
The interesting number is where the miss moves. The second pass — the same six questions ranked in memory over the 13 real chunks, no store, no index, no near-copies — puts the vocabulary question at rank 2, not 5. Between the two settings the only thing that changed is the 39 jittered rows the store carries, and copies of the labelled document count as hits, so the rows that overtake the vocabulary chunk are all other documents’ near-copies: originals that scored within a hair of it, each now flanked by three copies one percent of cosine away. The top of this ranking is packed tight enough that rank 2 and rank 5 are the same result rounded differently. That is worth knowing before tuning anything: on a corpus this small, one ranking position is noise, and the labelled set’s real finding is the pile-up, not the miss.
The chunking trade, measured
Section titled “The chunking trade, measured”Step 4 chose 30-word windows with a 6-word overlap and said outright
that the choice was a trade to measure. eval_window is the
measurement: for each window size it re-chunks the three documents,
embeds the chunks and the questions, and ranks in memory — no store,
no index, so the number isolates the window and stride from everything
step 5 added on top. It returns each question’s rank, because a
summary line hides which question moved, and which one moved is the
interesting part.
the chunking trade, brute force over the real chunks: window 30 stride 24: recall@3 6/6, MRR 0.92, ranks 1,1,1,1,2,1 window 15 stride 10: recall@3 6/6, MRR 0.92, ranks 1,1,1,1,2,1 window 60 stride 54: recall@3 5/6, MRR 0.88, ranks 1,1,1,1,4,1Halving the window to 15 words with stride 10 changes nothing — the same five questions at rank 1, the vocabulary question still at 2. Doubling the window to 60 does move a number: the vocabulary question falls to rank 4, out of the top 3. Six questions over three documents cannot say which trade is right, and this set does not pretend to; it can say that the choice moves the numbers, in which direction it moved here, and — because the ranks are printed per question — exactly which question paid for it. When the corpus stops being three paragraphs of this site’s own prose, the same run answers the same question about the real one.
Running it
Section titled “Running it”cargo run -q -- evalThe evaluation exits after the summary — no ollama, no port, nothing
listened on. Quality is a property of retrieval, not of the answer
model downstream, and a measurement that needs a second process to
talk to is a measurement that runs less often. With no argument the
program is unchanged from step 7: the branch in main below runs
eval and returns, and every other path serves POST /ask exactly as
the previous step built it.
//! Steps 3-7 are the pipeline: embed, chunk, store, index, prompt,//! serve. Step 8 asks whether retrieval actually works. `llm-apps//! eval` runs a small labelled set - questions paired with the//! document that holds their answer - through the same search the//! server uses, and prints recall@1, recall@3 and MRR next to a//! brute-force ranking of the same rows. A second pass re-chunks and//! re-ranks in memory under three window sizes, so the chunking trade//! from step 4 gets measured instead of asserted. With no argument the//! program is still the step 7 server: POST /ask, streaming over SSE.
use anyhow::Result;use axum::{ extract::State, response::sse::{Event, KeepAlive, KeepAliveStream, Sse}, routing::post, Json, Router,};use candle_core::{Device, Tensor};use candle_nn::VarBuilder;use candle_transformers::models::bert::{BertModel, Config, DTYPE};use futures_util::StreamExt;use hf_hub::api::sync::Api;use rusqlite::{params, params_from_iter, Connection};use serde::Deserialize;use serde_json::json;use serde_json::Value;use std::convert::Infallible;use std::sync::{Arc, Mutex};use tokenizers::Tokenizer;use tokio_stream::wrappers::ReceiverStream;
const MODEL_ID: &str = "sentence-transformers/all-MiniLM-L6-v2";
/// The answer model: step 2's ollama runtime, streaming. This program/// never loads its weights - generation is an HTTP call away.const ANSWER_MODEL: &str = "qwen2.5:0.5b";const OLLAMA_URL: &str = "http://127.0.0.1:11434/api/generate";
/// Hard cap on the answer's length, in tokens, passed to ollama. A/// question-answering endpoint that never stops writing is a bug.const MAX_ANSWER_TOKENS: u32 = 120;const DEFAULT_PORT: u16 = 8091;
/// Three documents, three topics: the runtime's memory footprint, the/// chunking trade, and the tokenizer. Same corpus as step 6.const DOCS: [&str; 3] = [ "Running a sentence encoder on your own machine is mostly an exercise in memory budgeting. The runtime in this project is candle, which loads the safetensors weights through memory mapping, so nothing is copied until a layer actually runs. Startup cost is dominated entirely by paging the model in on the first forward pass. Its peak memory footprint is about 340 megabytes. That is with three sentences in the batch; a long document pushes the number up. Tokenization itself is cheap, microseconds per sentence, and it happens before anything reaches the model. Embedding a whole document at once is the mistake this step fixes: the vector averages every topic together and retrieval degrades quietly.", "Chunking is where retrieval quality is decided. Windows of thirty words with a six-word overlap worked well for short technical prose: small enough that one vector describes one idea, overlapped just enough that a sentence split across a boundary still lands whole in at least one chunk. Bigger windows average more topics into each vector and recall suffers quietly. Smaller windows fragment the answer across chunks and the prompt fills with fragments. No window size is right for every corpus; the overlap and the stride are a trade to measure, not a default to trust.", "Tokenization happens before anything reaches the model, and it is the cheap part of the pipeline: microseconds per sentence on CPU. The tokenizer splits text into word pieces drawn from a fixed vocabulary of about thirty thousand entries, and the count it reports is the number of positions the model actually sees. A batch of sentences is padded to its longest member, and the attention mask marks the padding so the pooled embedding ignores it. That count is also the only honest way to budget a context window, which is what this step uses it for.",];
/// How wide the net is cast, how close two chunks must be before the/// second counts as a duplicate of the first, and how many tokens the/// context window may hold. Same dials as step 6.const CANDIDATES: usize = 15;const DUP_THRESHOLD: f32 = 0.98;const CONTEXT_TOKENS: usize = 100;
/// Index size (clusters) and how many of them a query visits.const K: usize = 8;const PROBE: usize = 2;/// Synthetic near-copies per real chunk, as in step 6.const COPIES: usize = 3;
/// Cut `doc` into word windows of `window` words, advancing `stride`/// words each time. Same chunker as step 4.fn chunk(doc: &str, window: usize, stride: usize) -> Vec<String> { let words: Vec<&str> = doc.split_whitespace().collect(); let mut out = Vec::new(); let mut start = 0; while start < words.len() { let end = (start + window).min(words.len()); out.push(words[start..end].join(" ")); if end == words.len() { break; } start += stride; } out}
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)}
/// Tokenize, forward, mean-pool: one Vec<f32> of 384 numbers per input.fn embed( model: &BertModel, device: &Device, tokenizer: &Tokenizer, texts: &[String],) -> Result<Vec<Vec<f32>>> { let encodings = tokenizer .encode_batch(texts.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]; texts.len()];
let ids = Tensor::new(token_ids, device)?; let mask = Tensor::new(attention.clone(), device)?; let types = Tensor::new(types, device)?;
let hidden = model.forward(&ids, &types, Some(&mask))?.to_vec3::<f32>()?; Ok(hidden .iter() .zip(attention.iter()) .map(|(tokens, mask)| { 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 }) .collect())}
/// A vector becomes a BLOB: 384 little-endian f32s, 1536 bytes.fn to_blob(v: &[f32]) -> Vec<u8> { v.iter().flat_map(|f| f.to_le_bytes()).collect()}
/// And back. `chunks_exact` panics on a corrupted row rather than/// silently mis-decoding it.fn from_blob(bytes: &[u8]) -> Vec<f32> { bytes .chunks_exact(4) .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) .collect()}
/// Deterministic pseudo-noise, so every run of this program prints the/// same numbers. Amplitude 0.01 keeps each copy within a hair of its/// original - a stand-in for the near-neighbours a real corpus produces.fn jitter(seed: usize, dim: usize) -> f32 { ((seed * 31 + dim * 7) as f32 * 0.173).sin() * 0.01}
/// Lloyd's k-means with a deterministic start: seed 0, then the vector/// farthest from it, then repeatedly the vector farthest from everything/// chosen so far. Runs to convergence or `iters` passes, whichever first.fn kmeans(vectors: &[Vec<f32>], k: usize, iters: usize) -> (Vec<Vec<f32>>, Vec<usize>) { let dist = |a: &Vec<f32>, b: &Vec<f32>| 1.0 - cosine(a, b); let mut seeds: Vec<usize> = vec![0]; while seeds.len() < k { let (mut best, mut bd) = (0usize, -1.0f32); for (i, v) in vectors.iter().enumerate() { if seeds.contains(&i) { continue; } let d: f32 = seeds.iter().map(|&s| dist(v, &vectors[s])).sum(); if d > bd { bd = d; best = i; } } seeds.push(best); } let mut centroids: Vec<Vec<f32>> = seeds.iter().map(|&s| vectors[s].clone()).collect(); let mut assign = vec![0usize; vectors.len()]; for _ in 0..iters { let mut changed = false; for (i, v) in vectors.iter().enumerate() { let nearest = centroids .iter() .enumerate() .max_by(|a, b| cosine(v, a.1).total_cmp(&cosine(v, b.1))) .map(|(ci, _)| ci) .unwrap(); if assign[i] != nearest { changed = true; assign[i] = nearest; } } for (ci, c) in centroids.iter_mut().enumerate() { let members: Vec<&Vec<f32>> = vectors .iter() .enumerate() .filter(|(i, _)| assign[*i] == ci) .map(|(_, v)| v) .collect(); if members.is_empty() { continue; } for d in 0..c.len() { c[d] = members.iter().map(|v| v[d]).sum::<f32>() / members.len() as f32; } } if !changed { break; } } (centroids, assign)}
/// The labelled set: a question, and the index of the DOCS entry that/// holds its answer. Kept small on purpose - a set whose labels can be/// checked by hand is worth more than a benchmark nobody re-reads.const QUESTIONS: [(&str, usize); 6] = [ ("How much memory does the model use at peak?", 0), ("What window size worked for short technical prose?", 1), ("How does the attention mask treat padding?", 2), ("Why does embedding a whole document hurt?", 0), ("How big is the word-piece vocabulary?", 2), ("What happens to recall when windows are too small?", 1),];
/// One search result: its cosine score against the question, the/// chunk's full text, the decoded vector - kept alongside the text/// because the dedupe pass needs it - and the row's chunk id, which/// the evaluation maps back to the document it came from.#[derive(Clone)]struct Hit { score: f32, text: String, vector: Vec<f32>, id: i64,}
/// The token count under the model's own tokenizer.fn token_count(tokenizer: &Tokenizer, text: &str) -> usize { tokenizer .encode(text, false) .map(|e| e.get_ids().len()) .unwrap_or(0)}
/// Indexed search, as step 5 built it: rank the centroids, then read/// only the rows assigned to the `probe` nearest clusters. Returns the/// clusters visited, the rows read, and up to `top` hits.fn search_index( conn: &Connection, q: &[f32], probe: usize, top: usize,) -> Result<(Vec<(i64, f32)>, usize, Vec<Hit>)> { let mut stmt = conn.prepare("SELECT id, vector FROM centroids")?; let mut rows = stmt.query([])?; let mut clusters: Vec<(i64, f32)> = Vec::new(); while let Some(row) = rows.next()? { let id: i64 = row.get(0)?; let c = from_blob(&row.get::<_, Vec<u8>>(1)?); clusters.push((id, cosine(&c, q))); } clusters.sort_by(|a, b| b.1.total_cmp(&a.1)); let chosen: Vec<i64> = clusters[..probe].iter().map(|c| c.0).collect();
let list = chosen .iter() .enumerate() .map(|(i, _)| format!("?{}", i + 1)) .collect::<Vec<_>>() .join(", "); let mut stmt = conn.prepare(&format!( "SELECT v.chunk_id, v.embedding, c.text FROM vectors v JOIN chunks c ON c.id = v.chunk_id JOIN assignments a ON a.chunk_id = v.chunk_id WHERE a.centroid_id IN ({list})" ))?; let mut rows = stmt.query(params_from_iter(chosen.iter()))?; let mut scored: Vec<Hit> = Vec::new(); while let Some(row) = rows.next()? { let id: i64 = row.get(0)?; let vector = from_blob(&row.get::<_, Vec<u8>>(1)?); let text: String = row.get(2)?; let score = cosine(&vector, q); scored.push(Hit { score, text, vector, id, }); } let scanned = scored.len(); scored.sort_by(|a, b| b.score.total_cmp(&a.score)); scored.truncate(top); let visited: Vec<(i64, f32)> = clusters[..probe].to_vec(); Ok((visited, scanned, scored))}
/// True when `v` is within `threshold` cosine of anything already/// kept. The threshold is deliberately generous: it exists to catch/// near-copies, not related chunks.fn is_duplicate(v: &[f32], kept: &[Vec<f32>], threshold: f32) -> bool { kept.iter().any(|k| cosine(v, k) > threshold)}
/// Walk the ranked candidates, keep the first of every distinct chunk,/// and stop at the first passage that no longer fits the budget.fn select( candidates: &[Hit], tokenizer: &Tokenizer, threshold: f32, budget: usize,) -> (Vec<Hit>, usize, usize) { let mut kept: Vec<Hit> = Vec::new(); let mut vectors: Vec<Vec<f32>> = Vec::new(); let mut used = 0usize; let (mut dupes, mut over) = (0usize, 0usize); for cand in candidates { if is_duplicate(&cand.vector, &vectors, threshold) { dupes += 1; continue; } let tokens = token_count(tokenizer, &cand.text); if used + tokens > budget { over = candidates.len() - kept.len() - dupes; break; } used += tokens; vectors.push(cand.vector.clone()); kept.push(cand.clone()); } (kept, dupes, over)}
/// The prompt the answer model will respond to: an instruction, the/// numbered passages, and the question last.fn build_prompt(passages: &[String], question: &str) -> String { let mut prompt = String::from( "Answer the question using only the passages below. \ If the passages do not contain the answer, say so.\n\n", ); for (i, passage) in passages.iter().enumerate() { prompt.push_str(&format!("[{}] {}\n\n", i + 1, passage)); } prompt.push_str("Question: "); prompt.push_str(question); prompt}
/// Which DOCS entry a chunk row came from. Real chunks carry ids/// `0..real` in chunking order; the near-copies of chunk `i` carry ids/// `real + i * COPIES + j`, same layout step 6 inserted them with.fn doc_of_id(id: i64, real: usize, chunks: &[(usize, String)]) -> usize { if (id as usize) < real { chunks[id as usize].0 } else { chunks[(id as usize - real) / COPIES].0 }}
/// Rank of the first hit that comes from the labelled document,/// 1-based; None when the ranking never reaches it.fn rank_of_label( ranked: &[(f32, i64)], label: usize, real: usize, chunks: &[(usize, String)],) -> Option<usize> { ranked .iter() .position(|&(_, id)| doc_of_id(id, real, chunks) == label) .map(|r| r + 1)}
/// `-` reads better than `>15` in a table column that is three/// characters wide.fn fmt_rank(rank: Option<usize>) -> String { match rank { Some(r) => r.to_string(), None => "-".to_string(), }}
/// The evaluation itself. Each labelled question is embedded once, then/// ranked two ways: through `search_index`, the exact path a `/ask`/// request takes, and by brute force over every row in the store. When/// the two disagree, the index - not the corpus - is what missed.fn eval( conn: &Connection, model: &BertModel, device: &Device, tokenizer: &Tokenizer, chunks: &[(usize, String)], ids: &[i64], all: &[Vec<f32>],) -> Result<()> { let real = chunks.len(); let questions: Vec<String> = QUESTIONS.iter().map(|(q, _)| q.to_string()).collect(); let qv = embed(model, device, tokenizer, &questions)?;
println!( "{:<52} {:>3} {:>7} {:>6}", "question", "doc", "indexed", "brute" ); let (mut top1, mut top3, mut mrr) = (0usize, 0usize, 0.0f64); for (qi, &(question, label)) in QUESTIONS.iter().enumerate() { let (_, _, hits) = search_index(conn, &qv[qi], PROBE, CANDIDATES)?; let indexed = hits.iter().map(|h| (h.score, h.id)).collect::<Vec<_>>(); let brute: Vec<(f32, i64)> = all .iter() .zip(ids) .map(|(v, &id)| (cosine(v, &qv[qi]), id)) .collect(); let mut brute = brute; brute.sort_by(|a, b| b.0.total_cmp(&a.0));
let irank = rank_of_label(&indexed, label, real, chunks); let brink = rank_of_label(&brute, label, real, chunks); if irank == Some(1) { top1 += 1; } if matches!(irank, Some(r) if r <= 3) { top3 += 1; } if let Some(r) = irank { mrr += 1.0 / r as f64; } println!( "{:<52} {:>3} {:>7} {:>6}", format!("{question:.50}"), label, fmt_rank(irank), fmt_rank(brink) ); } let n = QUESTIONS.len(); println!( "recall@1 {top1}/{n}, recall@3 {top3}/{n}, MRR {:.2} over the indexed search", mrr / n as f64 );
// The chunking trade, measured. Re-chunk and re-embed for each // window size and rank in memory - no store, no index - so the // number isolates the window and stride from everything step 5 // added on top. println!("\nthe chunking trade, brute force over the real chunks:"); for (window, stride) in [(30usize, 24usize), (15, 10), (60, 54)] { let (h3, n, mrr, ranks) = eval_window(model, device, tokenizer, window, stride)?; let ranks = ranks.iter().map(|r| fmt_rank(*r)).collect::<Vec<_>>().join(","); println!(" window {window:>2} stride {stride:>2}: recall@3 {h3}/{n}, MRR {mrr:.2}, ranks {ranks}"); } Ok(())}
/// One window size, end to end: chunk, embed, rank the labelled/// questions against the real chunks in memory. Returns hits in the/// top 3, the question count, mean reciprocal rank, and each/// question's rank - the summary hides which question moved, and/// which one moved is the interesting part.fn eval_window( model: &BertModel, device: &Device, tokenizer: &Tokenizer, window: usize, stride: usize,) -> Result<(usize, usize, f64, Vec<Option<usize>>)> { let mut chunks: Vec<(usize, String)> = Vec::new(); for (d, doc) in DOCS.iter().enumerate() { for c in chunk(doc, window, stride) { chunks.push((d, c)); } } let texts: Vec<String> = chunks.iter().map(|(_, c)| c.clone()).collect(); let vectors = embed(model, device, tokenizer, &texts)?; let questions: Vec<String> = QUESTIONS.iter().map(|(q, _)| q.to_string()).collect(); let qv = embed(model, device, tokenizer, &questions)?;
let (mut hits3, mut mrr) = (0usize, 0.0f64); let mut ranks: Vec<Option<usize>> = Vec::new(); for (qi, &(_, label)) in QUESTIONS.iter().enumerate() { let mut ranked: Vec<(f32, usize)> = vectors .iter() .enumerate() .map(|(i, v)| (cosine(v, &qv[qi]), i)) .collect(); ranked.sort_by(|a, b| b.0.total_cmp(&a.0)); let rank = ranked .iter() .position(|&(_, i)| chunks[i].0 == label) .map(|r| r + 1); if let Some(r) = rank { if r <= 3 { hits3 += 1; } mrr += 1.0 / r as f64; } ranks.push(rank); } Ok((hits3, QUESTIONS.len(), mrr / QUESTIONS.len() as f64, ranks))}
/// State shared by every request: the store and its index (steps 3-5),/// plus what it takes to embed a question on the way in. The SQLite/// connection is behind a mutex because `Connection` is not `Sync`.struct App { conn: Mutex<Connection>, model: BertModel, device: Device, tokenizer: Tokenizer,}
/// The request body.#[derive(Deserialize)]struct Ask { question: String,}
async fn ask( State(app): State<Arc<App>>, Json(req): Json<Ask>,) -> Sse<KeepAliveStream<ReceiverStream<Result<Event, Infallible>>>> { let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(32); tokio::spawn(relay(app, req.question, tx)); Sse::new(ReceiverStream::new(rx)).keep_alive(KeepAlive::default())}
/// The request pipeline, off the request task entirely. Retrieval runs/// on the blocking pool - it is CPU work against a mutex-guarded/// connection - and ollama's token stream is relayed from the async/// side. The client sees the passages first, then the tokens.async fn relay( app: Arc<App>, question: String, tx: tokio::sync::mpsc::Sender<Result<Event, Infallible>>,) { let embedder = app.clone(); let work = tokio::task::spawn_blocking(move || -> Result<(Vec<Hit>, usize, String)> { let conn = embedder.conn.lock().unwrap(); let q = embed( &embedder.model, &embedder.device, &embedder.tokenizer, &[question.clone()], )? .remove(0); let (_, _, hits) = search_index(&conn, &q, PROBE, CANDIDATES)?; let (kept, _, _) = select(&hits, &embedder.tokenizer, DUP_THRESHOLD, CONTEXT_TOKENS); let prompt = build_prompt( &kept.iter().map(|h| h.text.clone()).collect::<Vec<_>>(), &question, ); let tokens = token_count(&embedder.tokenizer, &prompt); Ok((kept, tokens, prompt)) });
let (kept, prompt_tokens, prompt) = match work.await { Ok(Ok(x)) => x, Ok(Err(e)) => { send_error(&tx, format!("{e}")).await; return; } Err(e) => { send_error(&tx, format!("{e}")).await; return; } };
// The retrieved chunks are named first, one event each, so a client // can show its sources while the answer is still being written. if !send_passages(&kept, &app.tokenizer, &tx).await { return; // client went away }
let body = json!({ "model": ANSWER_MODEL, "prompt": prompt, "stream": true, // Greedy decoding: a factual question over retrieved passages // has one best answer, and sampling at the default temperature // makes two identical requests disagree. "options": { "num_predict": MAX_ANSWER_TOKENS, "temperature": 0.0 }, }); let response = match reqwest::Client::new().post(OLLAMA_URL).json(&body).send().await { Ok(r) if r.status().is_success() => r, Ok(r) => { send_error(&tx, format!("ollama returned {}", r.status())).await; return; } Err(e) => { send_error(&tx, format!("ollama unreachable: {e}")).await; return; } };
// Ollama streams newline-delimited JSON; relay_lines turns each // complete line into an SSE event. let stream = response.bytes_stream(); relay_lines(stream, &tx, prompt_tokens).await;}
/// The retrieved chunks are named first, one event each, so a client/// can show its sources while the answer is still being written./// Returns false when the client has gone away.async fn send_passages( kept: &[Hit], tokenizer: &Tokenizer, tx: &tokio::sync::mpsc::Sender<Result<Event, Infallible>>,) -> bool { for (i, h) in kept.iter().enumerate() { let passage = json!({ "rank": i + 1, "score": h.score, "tokens": token_count(tokenizer, &h.text), "text": h.text, }); if tx .send(Ok(Event::default().event("passage").data(passage.to_string()))) .await .is_err() { return false; } } true}
/// Consume ollama's newline-delimited JSON stream and relay each/// complete line as an SSE event. Each line carries one token in its/// "response" field; the line with "done": true carries the final/// counts and ends the answer.async fn relay_lines( mut stream: impl futures_util::Stream<Item = reqwest::Result<bytes::Bytes>> + Unpin, tx: &tokio::sync::mpsc::Sender<Result<Event, Infallible>>, prompt_tokens: usize,) { let mut buf = String::new(); while let Some(chunk) = stream.next().await { match chunk { Ok(bytes) => buf.push_str(&String::from_utf8_lossy(&bytes)), Err(e) => { send_error(tx, format!("{e}")).await; return; } } while let Some(pos) = buf.find('\n') { let line: String = buf.drain(..=pos).collect(); let line = line.trim(); if line.is_empty() { continue; } let v: Value = match serde_json::from_str(line) { Ok(v) => v, Err(_) => continue, }; if let Some(token) = v.get("response").and_then(|r| r.as_str()) { if !token.is_empty() { let ev = Event::default().event("token").data(token); if tx.send(Ok(ev)).await.is_err() { return; // client went away } } } if v.get("done").and_then(|d| d.as_bool()) == Some(true) { let summary = json!({ "prompt_tokens": prompt_tokens, "answer_tokens": v.get("eval_count"), }); let _ = tx .send(Ok(Event::default().event("done").data(summary.to_string()))) .await; return; } } }}
async fn send_error(tx: &tokio::sync::mpsc::Sender<Result<Event, Infallible>>, msg: String) { let msg = msg.replace('\n', " "); let _ = tx .send(Ok(Event::default().event("error").data(msg))) .await;}
#[tokio::main]async 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 mut tokenizer = Tokenizer::from_file(repo.get("tokenizer.json")?).map_err(|e| anyhow::anyhow!("{e}"))?; // The published tokenizer.json ships with fixed padding to 128; // counting needs it off, as step 6 found. tokenizer.with_padding(None); 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());
// Same chunker as steps 4-6: 30-word windows, 6-word overlap. let mut chunks: Vec<(usize, String)> = Vec::new(); for (d, doc) in DOCS.iter().enumerate() { for c in chunk(doc, 30, 24) { chunks.push((d, c)); } }
// The store lives in the temp directory; delete any copy from an // earlier run so startup always begins from an empty database. let db_path = std::env::temp_dir().join("llm-step7-vectors.db"); let _ = std::fs::remove_file(&db_path); let conn = Connection::open(&db_path)?; conn.execute_batch( "CREATE TABLE chunks (id INTEGER PRIMARY KEY, text TEXT NOT NULL); CREATE TABLE vectors ( chunk_id INTEGER PRIMARY KEY REFERENCES chunks(id), embedding BLOB NOT NULL); CREATE TABLE centroids (id INTEGER PRIMARY KEY, vector BLOB NOT NULL); CREATE TABLE assignments ( chunk_id INTEGER PRIMARY KEY REFERENCES chunks(id), centroid_id INTEGER NOT NULL REFERENCES centroids(id));", )?;
let texts: Vec<String> = chunks.iter().map(|(_, c)| c.clone()).collect(); let vectors = embed(&model, &device, &tokenizer, &texts)?;
let real = chunks.len(); for (i, v) in vectors.iter().enumerate() { conn.execute( "INSERT INTO chunks (id, text) VALUES (?1, ?2)", params![i as i64, &chunks[i].1], )?; conn.execute( "INSERT INTO vectors (chunk_id, embedding) VALUES (?1, ?2)", params![i as i64, to_blob(v)], )?; } // Near-copies of every chunk, as in step 6. for (i, v) in vectors.iter().enumerate() { for j in 0..COPIES { let id = (real + i * COPIES + j) as i64; let noisy: Vec<f32> = v .iter() .enumerate() .map(|(d, x)| x + jitter(id as usize, d)) .collect(); conn.execute( "INSERT INTO chunks (id, text) VALUES (?1, ?2)", params![id, &chunks[i].1], )?; conn.execute( "INSERT INTO vectors (chunk_id, embedding) VALUES (?1, ?2)", params![id, to_blob(&noisy)], )?; } } let total: i64 = conn.query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0))?; println!( "stored {real} chunk vectors from {} documents, {COPIES} near-copies each ({total} rows)", DOCS.len() );
// The index from step 5, built the same way. let mut stmt = conn.prepare("SELECT chunk_id, embedding FROM vectors")?; let mut rows = stmt.query([])?; let mut ids: Vec<i64> = Vec::new(); let mut all: Vec<Vec<f32>> = Vec::new(); while let Some(row) = rows.next()? { ids.push(row.get(0)?); all.push(from_blob(&row.get::<_, Vec<u8>>(1)?)); } // Both must drop before `conn` is moved into the shared state. drop(rows); drop(stmt); let (centroids, assign) = kmeans(&all, K, 20); for (ci, c) in centroids.iter().enumerate() { conn.execute( "INSERT INTO centroids (id, vector) VALUES (?1, ?2)", params![ci as i64, to_blob(c)], )?; } for (&chunk_id, &ci) in ids.iter().zip(assign.iter()) { conn.execute( "INSERT INTO assignments (chunk_id, centroid_id) VALUES (?1, ?2)", params![chunk_id, ci as i64], )?; } let mut sizes = vec![0usize; K]; for &ci in &assign { sizes[ci] += 1; } sizes.sort(); println!("index built: {K} clusters, sizes {sizes:?}");
// `eval` measures the store exactly as the server is about to use // it - same index, same probe count - then exits. No ollama, no // port: quality is a property of retrieval, not of the answer // model downstream. if std::env::args().nth(1).as_deref() == Some("eval") { eval(&conn, &model, &device, &tokenizer, &chunks, &ids, &all)?; return Ok(()); }
let port: u16 = match std::env::args().nth(1) { Some(p) => p.parse()?, None => DEFAULT_PORT, }; let addr = format!("127.0.0.1:{port}"); let listener = tokio::net::TcpListener::bind(&addr).await?; let app = Arc::new(App { conn: Mutex::new(conn), model, device, tokenizer, }); let router = Router::new().route("/ask", post(ask)).with_state(app); println!("ready: POST http://{addr}/ask with {{\"question\": \"...\"}}"); axum::serve(listener, router).await?; Ok(())}The tag
Section titled “The tag”The step’s end state is tagged in the examples repository:
git checkout llm-ch08At that tag the project does everything steps 2-7 did — loads the
embedding model, chunks, stores vectors in SQLite, builds the index,
serves POST /ask with streaming answers — and measures itself: six
labelled questions, recall@1, recall@3, MRR, a brute-force cross-check
of the index, and a window-size ablation, all from one command and no
dependencies beyond the ones step 3 already pulled in.
The track is complete: eight steps, tags llm-ch02 through llm-ch08,
one local RAG service built end to end. The part worth keeping when
the corpus grows past three paragraphs is the shape, not the numbers:
a set small enough to check by hand, questions kept in it that are
expected to be hard, and metrics printed per question so a summary
never decides anything on the reader’s behalf. The
track overview lists every step in order.
Related
Section titled “Related”- An HTTP API for questions — the previous step, whose
/askpath this evaluation measures unchanged - Chunking documents for retrieval — the window-and-stride trade step 4 asserted and this step measured
- A vector store in SQLite — the index the brute-force column cross-checks