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 2 — Keyword Search from Scratch¶

We build a real keyword search engine step by step: inverted index → TF-IDF → BM25, all in plain Python over our own corpus, then cross-check BM25 against a library. We finish with the failure demo that motivates embeddings.

In [ ]:

# Bootstrap: find the repo root (folder containing data/sample_corpus.json) and import 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
docs = load_corpus()
print(f'Loaded {len(docs)} documents from', os.path.join(ROOT, 'data', 'sample_corpus.json'))

# Bootstrap: find the repo root (folder containing data/sample_corpus.json) and import 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 docs = load_corpus() print(f'Loaded {len(docs)} documents from', os.path.join(ROOT, 'data', 'sample_corpus.json'))

In [ ]:

from collections import Counter, defaultdict
import math

# Tokenize the whole corpus once (this is our 'indexing' step).
tok = [tokenize(d['text']) for d in docs]
print('D4 tokens:', tok[4][:12], '...')

from collections import Counter, defaultdict import math # Tokenize the whole corpus once (this is our 'indexing' step). tok = [tokenize(d['text']) for d in docs] print('D4 tokens:', tok[4][:12], '...')

1. The inverted index¶

Map every term to the documents that contain it (and how many times). This is what makes lookups instant.

In [ ]:

inverted = defaultdict(dict)
for i, toks in enumerate(tok):
    for term, count in Counter(toks).items():
        inverted[term][i] = count

print("postings for 'blue' :", inverted['blue'])   # {doc_id: count}
print("postings for 'whale':", inverted['whale'])

# Answering 'which docs contain blue?' is now a single dict lookup:
print('\nDocs containing "blue":', [docs[i]['title'] for i in inverted['blue']])

inverted = defaultdict(dict) for i, toks in enumerate(tok): for term, count in Counter(toks).items(): inverted[term][i] = count print("postings for 'blue' :", inverted['blue']) # {doc_id: count} print("postings for 'whale':", inverted['whale']) # Answering 'which docs contain blue?' is now a single dict lookup: print('\nDocs containing "blue":', [docs[i]['title'] for i in inverted['blue']])

2. Term Frequency and Inverse Document Frequency¶

Not all shared words matter equally. TF rewards words frequent in a document; IDF rewards words rare across the corpus. Let's reproduce the worked example from the chapter README on the toy corpus so you can see the numbers.

In [ ]:

toy = ['the grass is green', 'the sky is blue', 'the capital of canada is ottawa', 'a whale is a mammal', 'tomorrow is saturday']
toy_tok = [tokenize(t) for t in toy]
Ntoy = len(toy_tok)

toy_df = defaultdict(int)
for t in toy_tok:
    for term in set(t):
        toy_df[term] += 1

def toy_idf(term):
    return math.log(Ntoy / toy_df[term]) if toy_df[term] else float('nan')

print(f'{"term":8}{"df":>4}{"idf=ln(N/df)":>16}')
for term in ['is', 'the', 'grass']:
    print(f'{term:8}{toy_df[term]:>4}{toy_idf(term):>16.3f}')

# Score D0 for query 'what color is the grass' via sum of TF*IDF over present query terms
q = tokenize('what color is the grass')
f0 = Counter(toy_tok[0])
score_d0 = sum(f0[term]*toy_idf(term) for term in q if term in f0)
print(f'\nTF-IDF score(D0) = {score_d0:.3f}  (driven almost entirely by "grass")')

toy = ['the grass is green', 'the sky is blue', 'the capital of canada is ottawa', 'a whale is a mammal', 'tomorrow is saturday'] toy_tok = [tokenize(t) for t in toy] Ntoy = len(toy_tok) toy_df = defaultdict(int) for t in toy_tok: for term in set(t): toy_df[term] += 1 def toy_idf(term): return math.log(Ntoy / toy_df[term]) if toy_df[term] else float('nan') print(f'{"term":8}{"df":>4}{"idf=ln(N/df)":>16}') for term in ['is', 'the', 'grass']: print(f'{term:8}{toy_df[term]:>4}{toy_idf(term):>16.3f}') # Score D0 for query 'what color is the grass' via sum of TF*IDF over present query terms q = tokenize('what color is the grass') f0 = Counter(toy_tok[0]) score_d0 = sum(f0[term]*toy_idf(term) for term in q if term in f0) print(f'\nTF-IDF score(D0) = {score_d0:.3f} (driven almost entirely by "grass")')

is has IDF 0 (it's in every doc), grass has the highest IDF (it's rare). TF-IDF ranks for the right reason.

3. BM25 from scratch¶

BM25 improves TF-IDF with saturation (k1) and length normalization (b). We build it on the real corpus.

In [ ]:

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):
    # BM25's smoothed IDF
    return math.log(1 + (N - df[term] + 0.5) / (df[term] + 0.5))

def bm25_search(query, k1=1.5, b=0.75, topn=5):
    q = tokenize(query)
    scores = []
    for i, t in enumerate(tok):
        f = Counter(t)
        s = 0.0
        for term in q:
            if term not in f:
                continue
            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])[:topn]

print('Query: what is the most viewed televised event')
for i, s in bm25_search('what is the most viewed televised event'):
    print(f'  {s:5.2f}  D{i}: {docs[i]["title"]}')

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): # BM25's smoothed IDF return math.log(1 + (N - df[term] + 0.5) / (df[term] + 0.5)) def bm25_search(query, k1=1.5, b=0.75, topn=5): q = tokenize(query) scores = [] for i, t in enumerate(tok): f = Counter(t) s = 0.0 for term in q: if term not in f: continue 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])[:topn] print('Query: what is the most viewed televised event') for i, s in bm25_search('what is the most viewed televised event'): print(f' {s:5.2f} D{i}: {docs[i]["title"]}')

BM25 confidently ranks the two sports/TV passages on top — exactly what we want from first-stage keyword retrieval.

4. Cross-check against a library¶

Our from-scratch BM25 should agree (up to parameter/IDF-variant details) with the popular rank-bm25 package. If it isn't installed, this cell just skips.

In [ ]:

try:
    from rank_bm25 import BM25Okapi
    bm = BM25Okapi(tok)
    q = tokenize('what is the most viewed televised event')
    import numpy as np
    order = np.argsort(bm.get_scores(q))[::-1][:5]
    print('rank-bm25 top results:')
    for i in order:
        print(f'  D{i}: {docs[i]["title"]}')
except ImportError:
    print('rank-bm25 not installed - run: pip install rank-bm25')

try: from rank_bm25 import BM25Okapi bm = BM25Okapi(tok) q = tokenize('what is the most viewed televised event') import numpy as np order = np.argsort(bm.get_scores(q))[::-1][:5] print('rank-bm25 top results:') for i in order: print(f' D{i}: {docs[i]["title"]}') except ImportError: print('rank-bm25 not installed - run: pip install rank-bm25')

5. The failure demo¶

Now the query BM25 cannot handle — the answer is phrased with different words. Watch where the migraine passage lands.

In [ ]:

query = 'strong pain in the side of the head'
print('Query:', query, '\n')
ranking = bm25_search(query, topn=len(docs))
for i, s in ranking[:5]:
    print(f'  {s:5.2f}  D{i}: {docs[i]["title"]}')

migraine_id = next(d['id'] for d in docs if d['title'] == 'Migraine symptoms')
pos = [i for i,_ in ranking].index(migraine_id)
print(f'\nThe migraine passage (D{migraine_id}) that truly matches the meaning is at rank {pos+1}.')

query = 'strong pain in the side of the head' print('Query:', query, '\n') ranking = bm25_search(query, topn=len(docs)) for i, s in ranking[:5]: print(f' {s:5.2f} D{i}: {docs[i]["title"]}') migraine_id = next(d['id'] for d in docs if d['title'] == 'Migraine symptoms') pos = [i for i,_ in ranking].index(migraine_id) print(f'\nThe migraine passage (D{migraine_id}) that truly matches the meaning is at rank {pos+1}.')

BM25 latches onto surface words like pain and misses the passage that actually describes the condition. No amount of parameter tuning fixes this — it's a fundamental limit of matching words.

Next: Chapter 3 begins turning text into meaning-vectors so 'side of the head' and 'temple' can finally match.

Exercises¶

1. Sweep k1 in {0.5, 1.5, 3.0} and b in {0, 0.5, 1.0}. How do rankings shift on a long vs. short document?
2. Add stop-word removal to indexing. Which queries improve, which don't?
3. Extend bm25_search to also match on the title field with extra weight (a mini BM25F).
4. Construct a query where BM25 beats what you'd expect a meaning-based system to do (hint: exact names, codes, rare terms).