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

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)
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]

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
sims.shape

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

Let's display the search results:

L(hits[0])
def print_search_result(hit): print(f"{hit['score']:.4f} {nb_paths[hit['corpus_id']]}")
L(hits[0]).map(print_search_result)

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.