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
if 'google.colab' in sys.modules:
    !pip -q install umap-learn

# --- 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 if 'google.colab' in sys.modules: !pip -q install umap-learn

Chapter 4 — Embeddings Deep Dive¶

We use a real, open-source embedding model to turn text into meaning-vectors, and reproduce the lesson's key result: a question's nearest neighbor is its own answer.

First run downloads a small model (~90 MB), so it needs internet once, then works offline.

Setup¶

pip install sentence-transformers

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'))
from corpus import load_corpus, tokenize
import numpy as np

# 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')) from corpus import load_corpus, tokenize import numpy as np

In [ ]:

# Analysis helpers (pure numpy - independent of which model made the vectors)
def l2_normalize(mat):
    mat = np.asarray(mat, float)
    norms = np.linalg.norm(mat, axis=1, keepdims=True)
    return mat / np.clip(norms, 1e-12, None)

def cosine_sim_matrix(mat):
    m = l2_normalize(mat)
    return m @ m.T

def top_k(query_vec, matrix, k=3, exclude=None):
    q = l2_normalize(query_vec[None, :])[0]
    sims = l2_normalize(matrix) @ q
    order = np.argsort(-sims)
    if exclude is not None:
        order = [i for i in order if i != exclude]
    return [(int(i), float(sims[i])) for i in order[:k]]

# Analysis helpers (pure numpy - independent of which model made the vectors) def l2_normalize(mat): mat = np.asarray(mat, float) norms = np.linalg.norm(mat, axis=1, keepdims=True) return mat / np.clip(norms, 1e-12, None) def cosine_sim_matrix(mat): m = l2_normalize(mat) return m @ m.T def top_k(query_vec, matrix, k=3, exclude=None): q = l2_normalize(query_vec[None, :])[0] sims = l2_normalize(matrix) @ q order = np.argsort(-sims) if exclude is not None: order = [i for i in order if i != exclude] return [(int(i), float(sims[i])) for i in order[:k]]

1. Load the model and embed a few sentences¶

all-MiniLM-L6-v2 outputs 384-dimensional vectors and runs on CPU.

In [ ]:

from sentence_transformers import SentenceTransformer
model = SentenceTransformer('all-MiniLM-L6-v2')

vecs = model.encode(['an apple is a fruit', 'a banana is a fruit', 'a car has four wheels'])
print('shape:', vecs.shape, '  (n_sentences, dimensions)')
print('first 8 numbers of sentence 0:', np.round(vecs[0][:8], 3))

from sentence_transformers import SentenceTransformer model = SentenceTransformer('all-MiniLM-L6-v2') vecs = model.encode(['an apple is a fruit', 'a banana is a fruit', 'a car has four wheels']) print('shape:', vecs.shape, ' (n_sentences, dimensions)') print('first 8 numbers of sentence 0:', np.round(vecs[0][:8], 3))

2. Paraphrases are close; unrelated text is not¶

Different words, same meaning -> high cosine similarity.

In [ ]:

pair = model.encode(['hello, how are you?', "hi, how's it going?"])
unrel = model.encode(['hello, how are you?', 'the capital of Canada is Ottawa'])
def cos(a,b): return float(l2_normalize(a[None,:])[0] @ l2_normalize(b[None,:])[0])
print('paraphrase similarity :', round(cos(pair[0], pair[1]), 3))
print('unrelated similarity  :', round(cos(unrel[0], unrel[1]), 3))

pair = model.encode(['hello, how are you?', "hi, how's it going?"]) unrel = model.encode(['hello, how are you?', 'the capital of Canada is Ottawa']) def cos(a,b): return float(l2_normalize(a[None,:])[0] @ l2_normalize(b[None,:])[0]) print('paraphrase similarity :', round(cos(pair[0], pair[1]), 3)) print('unrelated similarity :', round(cos(unrel[0], unrel[1]), 3))

3. The 'aha' demo: questions find their answers¶

Embed Q/A sentences together; each question's nearest neighbor (excluding itself) is its answer -- despite few shared words.

In [ ]:

qa = [
    'what color is the sky?', 'the sky is blue.',
    'what is an apple?', 'an apple is a fruit.',
    'where does the bear live?', 'the bear lives in the woods.',
    'where is the world cup?', 'the world cup is in qatar.',
]
emb = model.encode(qa)
questions = [i for i in range(len(qa)) if qa[i].endswith('?')]
for qi in questions:
    best_i, score = top_k(emb[qi], emb, k=1, exclude=qi)[0]
    print(f'Q: {qa[qi]:32s} -> nearest: {qa[best_i]!r}  (cos={score:.2f})')

qa = [ 'what color is the sky?', 'the sky is blue.', 'what is an apple?', 'an apple is a fruit.', 'where does the bear live?', 'the bear lives in the woods.', 'where is the world cup?', 'the world cup is in qatar.', ] emb = model.encode(qa) questions = [i for i in range(len(qa)) if qa[i].endswith('?')] for qi in questions: best_i, score = top_k(emb[qi], emb, k=1, exclude=qi)[0] print(f'Q: {qa[qi]:32s} -> nearest: {qa[best_i]!r} (cos={score:.2f})')

4. The migraine query keyword search missed (Chapter 2)¶

Embed our course corpus and query it by meaning. The migraine passage should now rise to the top.

In [ ]:

docs = load_corpus()
corpus_emb = model.encode([d['text'] for d in docs])
q = model.encode(['strong pain in the side of the head'])[0]
print('Query: strong pain in the side of the head\n')
for i, s in top_k(q, corpus_emb, k=3):
    print(f'  {s:.3f}  D{i}: {docs[i]["title"]}')
print('\n(Compare with Chapter 2, where BM25 buried the migraine passage.)')

docs = load_corpus() corpus_emb = model.encode([d['text'] for d in docs]) q = model.encode(['strong pain in the side of the head'])[0] print('Query: strong pain in the side of the head\n') for i, s in top_k(q, corpus_emb, k=3): print(f' {s:.3f} D{i}: {docs[i]["title"]}') print('\n(Compare with Chapter 2, where BM25 buried the migraine passage.)')

5. Cluster by meaning and project to 2-D¶

KMeans groups the corpus into topics; PCA squashes 384-D vectors to 2-D so we can see the structure.

In [ ]:

from sklearn.cluster import KMeans
from sklearn.decomposition import PCA
import matplotlib.pyplot as plt

labels = KMeans(n_clusters=5, n_init=10, random_state=0).fit_predict(corpus_emb)
xy = PCA(n_components=2, random_state=0).fit_transform(corpus_emb)
plt.figure(figsize=(8,6))
plt.scatter(xy[:,0], xy[:,1], c=labels, cmap='tab10')
for i, d in enumerate(docs): plt.annotate(d['title'], (xy[i,0], xy[i,1]), fontsize=8)
plt.title('Course corpus embedded, clustered, and projected to 2-D'); plt.show()

from sklearn.cluster import KMeans from sklearn.decomposition import PCA import matplotlib.pyplot as plt labels = KMeans(n_clusters=5, n_init=10, random_state=0).fit_predict(corpus_emb) xy = PCA(n_components=2, random_state=0).fit_transform(corpus_emb) plt.figure(figsize=(8,6)) plt.scatter(xy[:,0], xy[:,1], c=labels, cmap='tab10') for i, d in enumerate(docs): plt.annotate(d['title'], (xy[i,0], xy[i,1]), fontsize=8) plt.title('Course corpus embedded, clustered, and projected to 2-D'); plt.show()

6. The semantic atlas: map a whole vocabulary¶

PCA above is a linear shadow of the space. UMAP is non-linear: it preserves each point's local neighbourhood, so clusters separate much more crisply. Let's build a small vocabulary of our own across eight everyday themes and see whether the model groups them without being told the categories exist.

In [ ]:

# UMAP is optional: install with `pip install umap-learn`.
# If it isn't available we fall back to PCA so this notebook always runs.
try:
    import umap
    HAVE_UMAP = True
except ImportError:
    HAVE_UMAP = False
    print("umap-learn not installed - falling back to PCA.")
    print("For the full effect:  pip install umap-learn")

SEED = 42

def project(vectors, n_components=2, n_neighbors=12, min_dist=0.25):
    """Reduce embeddings to 2-D/3-D. UMAP when available, else PCA."""
    vectors = np.asarray(vectors)
    if not HAVE_UMAP:
        return PCA(n_components=n_components, random_state=SEED).fit_transform(vectors), "PCA"
    n_neighbors = max(2, min(n_neighbors, len(vectors) - 1))
    reducer = umap.UMAP(
        n_components=n_components,
        n_neighbors=n_neighbors,   # how much local vs global structure to keep
        min_dist=min_dist,         # how tightly points may pack together
        metric="cosine",           # match the metric the embedding model was trained with
        random_state=SEED,         # reproducible layout
    )
    return reducer.fit_transform(vectors), "UMAP"

# UMAP is optional: install with `pip install umap-learn`. # If it isn't available we fall back to PCA so this notebook always runs. try: import umap HAVE_UMAP = True except ImportError: HAVE_UMAP = False print("umap-learn not installed - falling back to PCA.") print("For the full effect: pip install umap-learn") SEED = 42 def project(vectors, n_components=2, n_neighbors=12, min_dist=0.25): """Reduce embeddings to 2-D/3-D. UMAP when available, else PCA.""" vectors = np.asarray(vectors) if not HAVE_UMAP: return PCA(n_components=n_components, random_state=SEED).fit_transform(vectors), "PCA" n_neighbors = max(2, min(n_neighbors, len(vectors) - 1)) reducer = umap.UMAP( n_components=n_components, n_neighbors=n_neighbors, # how much local vs global structure to keep min_dist=min_dist, # how tightly points may pack together metric="cosine", # match the metric the embedding model was trained with random_state=SEED, # reproducible layout ) return reducer.fit_transform(vectors), "UMAP"

In [ ]:

# Our own vocabulary - eight themes, nothing labelled for the model.
ATLAS = {
    "Fruit":       ["apple","banana","orange","mango","strawberry","pineapple","grape",
                    "peach","watermelon","cherry","lemon","blueberry"],
    "Vehicles":    ["car","bus","bicycle","motorcycle","truck","train","tram","scooter",
                    "aeroplane","helicopter","ferry","submarine"],
    "Buildings":   ["house","castle","cathedral","skyscraper","cottage","barn","lighthouse",
                    "museum","library","stadium","hospital","warehouse"],
    "Sports":      ["football","tennis","cricket","basketball","swimming","boxing","rugby",
                    "golf","cycling","marathon","surfing","skiing"],
    "Music":       ["guitar","piano","violin","drums","trumpet","flute","cello","saxophone",
                    "harp","clarinet","orchestra","melody"],
    "Weather":     ["rain","snow","thunderstorm","fog","hurricane","drought","blizzard",
                    "sunshine","hail","humidity","frost","lightning"],
    "Programming": ["python","javascript","compiler","database","algorithm","debugging",
                    "recursion","variable","function","repository","framework","syntax"],
    "Cooking":     ["roasting","simmer","marinate","whisk","saucepan","recipe","seasoning",
                    "baking","chopping","frying","dough","garnish"],
}

words  = [w for group in ATLAS.values() for w in group]
themes = [t for t, group in ATLAS.items() for _ in group]

atlas_emb = model.encode(words)
print(f"{len(words)} terms -> {atlas_emb.shape[1]}-dimensional vectors")

# Our own vocabulary - eight themes, nothing labelled for the model. ATLAS = { "Fruit": ["apple","banana","orange","mango","strawberry","pineapple","grape", "peach","watermelon","cherry","lemon","blueberry"], "Vehicles": ["car","bus","bicycle","motorcycle","truck","train","tram","scooter", "aeroplane","helicopter","ferry","submarine"], "Buildings": ["house","castle","cathedral","skyscraper","cottage","barn","lighthouse", "museum","library","stadium","hospital","warehouse"], "Sports": ["football","tennis","cricket","basketball","swimming","boxing","rugby", "golf","cycling","marathon","surfing","skiing"], "Music": ["guitar","piano","violin","drums","trumpet","flute","cello","saxophone", "harp","clarinet","orchestra","melody"], "Weather": ["rain","snow","thunderstorm","fog","hurricane","drought","blizzard", "sunshine","hail","humidity","frost","lightning"], "Programming": ["python","javascript","compiler","database","algorithm","debugging", "recursion","variable","function","repository","framework","syntax"], "Cooking": ["roasting","simmer","marinate","whisk","saucepan","recipe","seasoning", "baking","chopping","frying","dough","garnish"], } words = [w for group in ATLAS.values() for w in group] themes = [t for t, group in ATLAS.items() for _ in group] atlas_emb = model.encode(words) print(f"{len(words)} terms -> {atlas_emb.shape[1]}-dimensional vectors")

In [ ]:

xy, method = project(atlas_emb, n_components=2)

colors = plt.cm.tab10(np.linspace(0, 1, len(ATLAS)))
plt.figure(figsize=(12, 8))
for k, theme in enumerate(ATLAS):
    idx = [i for i, t in enumerate(themes) if t == theme]
    plt.scatter(xy[idx, 0], xy[idx, 1], s=70, color=colors[k],
                edgecolors="white", linewidths=0.8, label=theme)
    # name each cluster just above its points
    plt.text(xy[idx, 0].mean(), xy[idx, 1].max() + 0.25, theme.upper(),
             ha="center", fontsize=11, fontweight="bold", color=colors[k])

plt.title(f"A semantic atlas - {len(words)} terms, {method} projection of "
          f"{atlas_emb.shape[1]}-D embeddings")
plt.legend(loc="center left", bbox_to_anchor=(1.01, 0.5), frameon=False)
plt.xticks([]); plt.yticks([]); plt.tight_layout(); plt.show()

xy, method = project(atlas_emb, n_components=2) colors = plt.cm.tab10(np.linspace(0, 1, len(ATLAS))) plt.figure(figsize=(12, 8)) for k, theme in enumerate(ATLAS): idx = [i for i, t in enumerate(themes) if t == theme] plt.scatter(xy[idx, 0], xy[idx, 1], s=70, color=colors[k], edgecolors="white", linewidths=0.8, label=theme) # name each cluster just above its points plt.text(xy[idx, 0].mean(), xy[idx, 1].max() + 0.25, theme.upper(), ha="center", fontsize=11, fontweight="bold", color=colors[k]) plt.title(f"A semantic atlas - {len(words)} terms, {method} projection of " f"{atlas_emb.shape[1]}-D embeddings") plt.legend(loc="center left", bbox_to_anchor=(1.01, 0.5), frameon=False) plt.xticks([]); plt.yticks([]); plt.tight_layout(); plt.show()

Nobody handed the model these eight categories. It grouped them purely from how the words are used in language — meaning became location.

Try the in-between terms: "apple pie" should land between Fruit and Cooking, "racing car" between Vehicles and Sports. The space is continuous, not a set of boxes.

The same atlas in 3-D¶

Two dimensions is a harsh squeeze of 384. A third axis often untangles clusters that 2-D forces on top of each other — and it's a useful reminder that the real space has hundreds of axes.

In [ ]:

xyz, method3 = project(atlas_emb, n_components=3)

fig = plt.figure(figsize=(11, 8))
ax = fig.add_subplot(111, projection="3d")
for k, theme in enumerate(ATLAS):
    idx = [i for i, t in enumerate(themes) if t == theme]
    ax.scatter(xyz[idx, 0], xyz[idx, 1], xyz[idx, 2],
               s=55, color=colors[k], edgecolors="white", linewidths=0.5, label=theme)

ax.set_title(f"The atlas in three dimensions ({method3})")
ax.legend(loc="upper left", bbox_to_anchor=(-0.05, 0.9), frameon=False, fontsize=9)
ax.set_xticklabels([]); ax.set_yticklabels([]); ax.set_zticklabels([])
ax.view_init(elev=22, azim=-60)   # <- change these to spin the view
plt.tight_layout(); plt.show()

xyz, method3 = project(atlas_emb, n_components=3) fig = plt.figure(figsize=(11, 8)) ax = fig.add_subplot(111, projection="3d") for k, theme in enumerate(ATLAS): idx = [i for i, t in enumerate(themes) if t == theme] ax.scatter(xyz[idx, 0], xyz[idx, 1], xyz[idx, 2], s=55, color=colors[k], edgecolors="white", linewidths=0.5, label=theme) ax.set_title(f"The atlas in three dimensions ({method3})") ax.legend(loc="upper left", bbox_to_anchor=(-0.05, 0.9), frameon=False, fontsize=9) ax.set_xticklabels([]); ax.set_yticklabels([]); ax.set_zticklabels([]) ax.view_init(elev=22, azim=-60) # <- change these to spin the view plt.tight_layout(); plt.show()

UMAP vs PCA on identical vectors¶

Same 384-D embeddings, two projections. Worth understanding why they differ before you trust either one.

In [ ]:

fig, axes = plt.subplots(1, 2, figsize=(14, 6))

pca_xy = PCA(n_components=2, random_state=SEED).fit_transform(atlas_emb)
panels = [(xy, method), (pca_xy, "PCA")] if HAVE_UMAP else [(pca_xy, "PCA"), (pca_xy, "PCA")]

for ax, (coords, name) in zip(axes, panels):
    for k, theme in enumerate(ATLAS):
        idx = [i for i, t in enumerate(themes) if t == theme]
        ax.scatter(coords[idx, 0], coords[idx, 1], s=55, color=colors[k],
                   edgecolors="white", linewidths=0.6, label=theme)
    ax.set_title(name, fontweight="bold")
    ax.set_xticks([]); ax.set_yticks([])

axes[1].legend(loc="center left", bbox_to_anchor=(1.01, 0.5), frameon=False, fontsize=9)
plt.tight_layout(); plt.show()

fig, axes = plt.subplots(1, 2, figsize=(14, 6)) pca_xy = PCA(n_components=2, random_state=SEED).fit_transform(atlas_emb) panels = [(xy, method), (pca_xy, "PCA")] if HAVE_UMAP else [(pca_xy, "PCA"), (pca_xy, "PCA")] for ax, (coords, name) in zip(axes, panels): for k, theme in enumerate(ATLAS): idx = [i for i, t in enumerate(themes) if t == theme] ax.scatter(coords[idx, 0], coords[idx, 1], s=55, color=colors[k], edgecolors="white", linewidths=0.6, label=theme) ax.set_title(name, fontweight="bold") ax.set_xticks([]); ax.set_yticks([]) axes[1].legend(loc="center left", bbox_to_anchor=(1.01, 0.5), frameon=False, fontsize=9) plt.tight_layout(); plt.show()

How to read the difference

	UMAP	PCA
Kind	non-linear, neighbourhood-preserving	linear projection
Clusters	crisply separated	often overlapping
Distance within a cluster	meaningful	meaningful
Distance between clusters	not meaningful	meaningful
Reproducible	only with a fixed random_state	always
Cost	slower, needs tuning	instant

⚠️ The classic UMAP mistake: reading the gaps between clusters as "how different" those topics are. UMAP optimises local neighbourhoods and freely distorts global distances — two clusters sitting far apart may be no more unrelated than two sitting close. If you need distances you can reason about, use PCA — or better, compute cosine similarity on the original vectors, as we did earlier in this notebook.

Projections are for looking, not for deciding. Search always runs on the full 384-D vectors; nothing in your retrieval pipeline should ever run on a 2-D projection.

Takeaway & exercises¶

Nearest embedding to a query == search by meaning. That's dense retrieval -- built for real in Chapter 5.

Exercises

1. Add your own Q/A pairs to the demo. Does every question still find its answer?
2. Swap the model for all-mpnet-base-v2 (768-dim). Re-embed everything. Do neighbors improve?
3. Try a query with an exact rare term (a name/code). Where might keyword search still win? (Foreshadows Chapter 8.)
4. Add a theme of your own to ATLAS (films? animals? chemistry?) and re-plot. Does it claim its own region?
5. Add ambiguous terms -- "apple pie", "racing car", "football stadium" -- and see where they land between clusters.
6. Change UMAP's n_neighbors (try 2, 5, 50). Small values emphasise local detail, large values global shape. Where does the map stop being trustworthy?