Runs in your browser — no install. (Works once the repo is public.)

In [ ]:

# --- Colab setup (skipped when running locally) ---
import os, sys
if 'google.colab' in sys.modules and not os.path.exists('data/sample_corpus.json'):
    !git clone -q https://github.com/mdhabibi/llm-search-handbook.git
    %cd llm-search-handbook
    !pip -q install -r requirements.txt

# --- Colab setup (skipped when running locally) --- import os, sys if 'google.colab' in sys.modules and not os.path.exists('data/sample_corpus.json'): !git clone -q https://github.com/mdhabibi/llm-search-handbook.git %cd llm-search-handbook !pip -q install -r requirements.txt

Chapter 5 — Dense Retrieval (Semantic Search Engine)¶

We turn embeddings into a real search engine: embed the corpus once, embed the query, return the nearest documents. Then we race it against BM25.

First run downloads the small embedding model (internet needed once).

Setup¶

pip install sentence-transformers rank-bm25

In [ ]:

# Bootstrap: locate repo root and import shared helpers
import sys, os
d = os.getcwd()
while d != os.path.dirname(d) and not os.path.exists(os.path.join(d, 'data', 'sample_corpus.json')):
    d = os.path.dirname(d)
ROOT = d; sys.path.insert(0, os.path.join(ROOT, 'src'))
import numpy as np
from corpus import load_corpus, tokenize

# Bootstrap: locate repo root and import shared helpers import sys, os d = os.getcwd() while d != os.path.dirname(d) and not os.path.exists(os.path.join(d, 'data', 'sample_corpus.json')): d = os.path.dirname(d) ROOT = d; sys.path.insert(0, os.path.join(ROOT, 'src')) import numpy as np from corpus import load_corpus, tokenize

1. A reusable semantic search engine¶

We use the SemanticSearch class from src/semantic_search.py. It takes any encode function, embeds the corpus into a normalized matrix, and searches by cosine (= dot product on normalized vectors).

In [ ]:

from semantic_search import SemanticSearch
from sentence_transformers import SentenceTransformer

model = SentenceTransformer('all-MiniLM-L6-v2')
docs = load_corpus()
engine = SemanticSearch(model.encode).index(docs)
print('Indexed', len(docs), 'documents into a', engine.matrix.shape, 'matrix.')

from semantic_search import SemanticSearch from sentence_transformers import SentenceTransformer model = SentenceTransformer('all-MiniLM-L6-v2') docs = load_corpus() engine = SemanticSearch(model.encode).index(docs) print('Indexed', len(docs), 'documents into a', engine.matrix.shape, 'matrix.')

2. Query in plain English¶

No shared-word requirement — ask by meaning.

In [ ]:

def show(query, k=3):
    print('Q:', query)
    for i, score, doc in engine.search(query, k=k):
        print(f'   {score:.3f}  D{i}: {doc["title"]}')
    print()

show('why is grass green?')
show('a large ocean animal')
show('capital city of a country')

def show(query, k=3): print('Q:', query) for i, score, doc in engine.search(query, k=k): print(f' {score:.3f} D{i}: {doc["title"]}') print() show('why is grass green?') show('a large ocean animal') show('capital city of a country')

3. The migraine query — finally solved¶

Our running example since Chapter 1. BM25 buried it; semantic search should surface it.

In [ ]:

show('strong pain in the side of the head')
print('The migraine passage should now be the top hit -- meaning matched, no shared words needed.')

show('strong pain in the side of the head') print('The migraine passage should now be the top hit -- meaning matched, no shared words needed.')

4. Dense vs. BM25, side by side¶

Neither is strictly better. We build a quick BM25 (from Chapter 2) and compare on the same queries.

In [ ]:

from collections import Counter, defaultdict
import math
tok = [tokenize(d['text']) for d in docs]
N = len(tok); avgdl = sum(len(t) for t in tok)/N
df = defaultdict(int)
for t in tok:
    for term in set(t): df[term]+=1
def idf(term): return math.log(1 + (N-df[term]+0.5)/(df[term]+0.5))
def bm25(query, k=3, k1=1.5, b=0.75):
    q = tokenize(query); scores=[]
    for i,t in enumerate(tok):
        f=Counter(t); s=0.0
        for term in q:
            if term in f:
                s += idf(term)*(f[term]*(k1+1))/(f[term]+k1*(1-b+b*len(t)/avgdl))
        scores.append((i,s))
    return sorted(scores,key=lambda x:-x[1])[:k]

def compare(query):
    print('QUERY:', query)
    print('  BM25   :', [docs[i]['title'] for i,_ in bm25(query)])
    print('  Dense  :', [doc['title'] for _,_,doc in engine.search(query)])
    print()

compare('strong pain in the side of the head')   # dense wins (synonyms)
compare('a large ocean animal')                  # dense wins (concept)
compare('the world cup final')                   # both do well (shared words)

from collections import Counter, defaultdict import math tok = [tokenize(d['text']) for d in docs] N = len(tok); avgdl = sum(len(t) for t in tok)/N df = defaultdict(int) for t in tok: for term in set(t): df[term]+=1 def idf(term): return math.log(1 + (N-df[term]+0.5)/(df[term]+0.5)) def bm25(query, k=3, k1=1.5, b=0.75): q = tokenize(query); scores=[] for i,t in enumerate(tok): f=Counter(t); s=0.0 for term in q: if term in f: s += idf(term)*(f[term]*(k1+1))/(f[term]+k1*(1-b+b*len(t)/avgdl)) scores.append((i,s)) return sorted(scores,key=lambda x:-x[1])[:k] def compare(query): print('QUERY:', query) print(' BM25 :', [docs[i]['title'] for i,_ in bm25(query)]) print(' Dense :', [doc['title'] for _,_,doc in engine.search(query)]) print() compare('strong pain in the side of the head') # dense wins (synonyms) compare('a large ocean animal') # dense wins (concept) compare('the world cup final') # both do well (shared words)

5. A case where keyword search still wins¶

Semantic models can drift on exact rare tokens (names, codes, IDs). Our corpus mentions the exact library name Lucene inside the search-engine passage.

In [ ]:

compare('Lucene')
print('Exact rare tokens favor keyword matching -- motivation for HYBRID search (Chapter 8).')

compare('Lucene') print('Exact rare tokens favor keyword matching -- motivation for HYBRID search (Chapter 8).')

Takeaway & exercises¶

Dense retrieval searches by meaning; keyword search nails exactness. Chapter 6 scales dense retrieval to millions of vectors; Chapter 8 fuses both.

Exercises

1. Add 5 of your own passages, re-index, and query them.
2. Find two more queries where dense beats BM25, and one where BM25 beats dense.
3. Print the full cosine scores — are they comparable across different queries? (They're not; use them only to rank.)