audrey.feldroy.com

The experimental notebooks of Audrey M. Roy Greenfeld. This website and all its notebooks are open-source at github.com/audreyfeldroy/audrey.feldroy.com


# Semantic Search With Sentence Transformers and a Bi-Encoder Model

by Audrey M. Roy Greenfeld | Mon, Apr 14, 2025

Here I use sentence transformers and a bi-encoder model to encode my notebooks as embeddings and implement semantic search.


Setup

from fastcore.utils import *
from pathlib import Path
from sentence_transformers import SentenceTransformer, util

Initialize Bi-Encoder

Here we download a bi-encoder model to use for the precomputed embeddings.

bienc_model = SentenceTransformer('all-MiniLM-L6-v2')

Get All Notebook Paths

We put each notebook to be searched into a list.

def get_nb_paths(): 
    root = Path() if IN_NOTEBOOK else Path("nbs/")
    return L(root.glob("*.ipynb")).sorted(reverse=True)
nb_paths = get_nb_paths()
nb_paths
(#74) [Path('2025-04-15-Semantic-Search-With-Sentence-Transformers-and-a-Cross-Encoder-Model.ipynb'),Path('2025-04-14-Semantic-Search-With-Sentence-Transformers-and-a-Bi-Encoder-Model.ipynb'),Path('2025-04-08-Get-a-Jupyter-Notebook-Filename-From-Itself.ipynb'),Path('2025-03-20-Minimal-Screen-Recording-on-macOS-No-Third-Party-Apps-Required.ipynb'),Path('2025-03-14-Pi.ipynb'),Path('2025-02-20-One-Liner-to-Clean-Python-Bytecode.ipynb'),Path('2025-02-15-Building-a-Better-Title-Caser-Part-2-Using-an-Ollama-Modelfile.ipynb'),Path('2025-02-14-Building-a-Better-Title-Caser-Part-1-Beyond-Python-str-title.ipynb'),Path('2025-02-13-Excavating-a-Lost-CLI-Tool.ipynb'),Path('2025-02-12-My-Self-Analysis-of-How-to-Get-Back-to-Posting-Every-Day.ipynb'),Path('2025-02-09-An-Informationally-Dense-Index-Page.ipynb'),Path('2025-02-08-This-Notebook-Is-Also-a-Keylogger.ipynb'),Path('2025-02-07-This-Site-Is-Now-Powered-by-This-Notebook-Part-6.ipynb'),Path('2025-02-06-Creating-an-Accessible-Inline-Nav-FastTag.ipynb'),Path('2025-02-05-Create-a-CLI-Tool-With-Fastcore-Script.ipynb'),Path('2025-02-04-How-to-Turn-a-Jupyter-Notebook-Into-a-Python-Script.ipynb'),Path('2025-02-03-FastHTML-and-MonsterUI-Time-Converter.ipynb'),Path('2025-02-02-Text-Embeddings-and-Cosine-Similarity.ipynb'),Path('2025-02-01-Auto-Renaming-My-Untitled-ipynb-Files-With-Gemini.ipynb'),Path('2025-01-31-Performance-Optimization-Moving-HTML-Class-Injection-from-lxml-to-Mistletoe.ipynb')...]

Create an Embedding for Each Notebook

Now we can turn that list of notebook paths into embeddings by:

  1. Opening each notebook file
  2. Putting notebook content into a list of notebooks
  3. Passing the notebook list to the bi-encoder model to generate a list of embeddings
def read_nb_simple(nb_path):
    with open(nb_path, 'r', encoding='utf-8') as f:
        return f.read()
nbs = L(nb_paths).map(read_nb_simple)
nb_embs = bienc_model.encode(nbs)
len(nb_embs)
74
print(nb_embs.shape)

Encode the Query String

If we search for a particular query string, that string needs to be encoded as an embedding using the same bi-encoder as before. Then we can compare it to the notebook embeddings.

q = "Web search"
q_emb = bienc_model.encode(q)
q_emb[:10]
array([-0.0328431 , -0.00064043, -0.06456785, 0.01314389, -0.02520958, 0.02097196, 0.03034499, 0.05960393, -0.03566388, -0.03963251], dtype=float32)

Create a Cosine Similarities Tensor

Sentence Transformers provides a function to get the similarity between the query and each of the notebook embeddings. It defaults to cosine similarity. We use it to get a tensor of how similar the query embedding is to each notebook.

sims = bienc_model.similarity(q_emb, nb_embs)
sims
tensor([[ 0.2022, 0.1729, -0.0333, -0.0670, 0.0745, 0.0353, 0.0670, 0.0779, 0.0709, 0.0786, 0.1814, 0.0342, 0.0047, 0.0716, 0.0518, -0.0770, 0.1202, 0.1646, 0.0790, 0.0343, 0.0567, 0.0842, 0.0070, 0.1067, 0.0751, -0.0592, -0.0341, -0.0082, 0.0048, 0.0697, 0.0034, 0.0660, 0.1866, 0.0680, 0.0811, 0.0612, 0.1918, 0.2615, 0.2304, 0.1414, 0.0626, 0.1566, 0.0056, 0.1292, 0.0197, 0.1162, -0.0663, 0.0835, 0.0663, 0.0659, 0.0946, 0.1104, 0.0101, 0.1370, 0.0635, 0.0044, 0.0777, -0.0330, -0.0023, 0.0593, 0.0358, 0.0823, 0.0667, 0.0458, 0.1565, 0.1318, 0.1485, 0.1480, 0.0771, 0.0885, 0.0954, 0.0929, 0.0607, 0.1207]])
sims.shape
torch.Size([1, 74])

Get Top 10 Similar Results

Sentence Transformers also provides a semantic search utility that returns search results:

hits = util.semantic_search(q_emb, nb_embs, top_k=10)
hits
[[{'corpus_id': 37, 'score': 0.2615070044994354}, {'corpus_id': 38, 'score': 0.2304057776927948}, {'corpus_id': 0, 'score': 0.20217692852020264}, {'corpus_id': 36, 'score': 0.1918058693408966}, {'corpus_id': 32, 'score': 0.18661072850227356}, {'corpus_id': 10, 'score': 0.18138977885246277}, {'corpus_id': 1, 'score': 0.17291493713855743}, {'corpus_id': 17, 'score': 0.16460277140140533}, {'corpus_id': 41, 'score': 0.15658749639987946}, {'corpus_id': 64, 'score': 0.15653304755687714}]]

Let's display the search results:

L(hits[0])
(#10) [{'corpus_id': 37, 'score': 0.2615070044994354},{'corpus_id': 38, 'score': 0.2304057776927948},{'corpus_id': 0, 'score': 0.20217692852020264},{'corpus_id': 36, 'score': 0.1918058693408966},{'corpus_id': 32, 'score': 0.18661072850227356},{'corpus_id': 10, 'score': 0.18138977885246277},{'corpus_id': 1, 'score': 0.17291493713855743},{'corpus_id': 17, 'score': 0.16460277140140533},{'corpus_id': 41, 'score': 0.15658749639987946},{'corpus_id': 64, 'score': 0.15653304755687714}]
def print_search_result(hit): print(f"{hit['score']:.4f} {nb_paths[hit['corpus_id']]}")
L(hits[0]).map(print_search_result)
(#10) [None,None,None,None,None,None,None,None,None,None]

Define a Function

Putting together everything we've figured out:

def bienc_search_nbs(q):
    nb_paths = get_nb_paths()
    nbs = L(nb_paths).map(read_nb_simple)
    nb_embs = bienc_model.encode(nbs)
    q_emb = bienc_model.encode(q)
    hits = util.semantic_search(q_emb, nb_embs, top_k=10)
    L(hits[0]).map(print_search_result)

We can try out our biencoder-based semantic search function:

bienc_search_nbs("search")

Reflection

A bi-encoder is nice because it allows you to pregenerate embeddings and later use those for comparison. But I'm reading that it's not as accurate as a cross-encoder. In the next post we'll see if that's true.

© 2024-2025 Audrey M. Roy Greenfeld