Dense Retrieval: Semantic Search & Bi-Encoder Implementation › Code Implementation

!uv pip install sentence-transformers scikit-learn

import numpy as np
import matplotlib.pyplot as plt
from collections import Counter
from sklearn.decomposition import PCA
from sentence_transformers import SentenceTransformer

# Load a pre-trained bi-encoder model
# This model was trained on over 1 billion sentence pairs
# using contrastive learning on diverse text collections
model = SentenceTransformer('all-MiniLM-L6-v2')

# Sample document corpus about machine learning
documents = [
    "Neural networks learn hierarchical representations of data through multiple layers.",
    "Gradient descent optimizes model parameters by iteratively updating weights.",
    "Transformers use self-attention to capture dependencies regardless of distance.",
    "Random forests combine multiple decision trees for reliable predictions.",
    "Support vector machines find optimal hyperplanes to separate classes.",
    "Convolutional neural networks excel at processing grid-structured data like images.",
    "Recurrent neural networks maintain hidden states to process sequential data.",
    "The backpropagation algorithm computes gradients through the chain rule.",
]

# Encode all documents (this would be done offline in production)
# normalize_embeddings=True divides each vector by its L2 norm
# letting dot product to equal cosine similarity at query time
document_embeddings = model.encode(
    documents, convert_to_numpy=True, normalize_embeddings=True
)

def dense_retrieve(query, doc_embeddings, docs, model, top_k=3):
    """
    Retrieve top-k documents using dense retrieval.

    Args:
        query: The search query string
        doc_embeddings: Pre-computed and normalized document embeddings
        docs: List of document strings
        model: The sentence transformer model
        top_k: Number of documents to retrieve

    Returns:
        List of (document, score) tuples sorted by descending relevance
    """
    # Encode the query (normalize to unit length for cosine similarity)
    query_embedding = model.encode(
        query, convert_to_numpy=True, normalize_embeddings=True
    )

    # Compute cosine similarity with all documents
    # For normalized embeddings, dot product equals cosine similarity
    # Shape of result: (num_docs,)
    similarities = np.dot(doc_embeddings, query_embedding)

    # Get indices of top-k highest scores (argsort returns ascending,
    # so we reverse with [::-1] and take first top_k)
    top_indices = np.argsort(similarities)[::-1][:top_k]

    # Return documents with their similarity scores
    return [(docs[i], similarities[i]) for i in top_indices]

query = "How do neural networks learn?"
results = dense_retrieve(query, document_embeddings, documents, model, top_k=3)

def bm25_score(query, doc, avg_doc_len, k1=1.5, b=0.75):
    """Compute simplified BM25 score for a single document."""
    query_terms = query.lower().split()
    doc_terms = doc.lower().split()
    doc_len = len(doc_terms)

    term_freq = Counter(doc_terms)

    score = 0.0
    for term in query_terms:
        if term in term_freq:
            tf = term_freq[term]
            # BM25 term frequency normalization
            numerator = tf * (k1 + 1)
            denominator = tf + k1 * (1 - b + b * (doc_len / avg_doc_len))
            score += numerator / denominator

    return score


def lexical_retrieve(query, docs, top_k=3):
    """Retrieve using simplified BM25."""
    avg_len = np.mean([len(d.split()) for d in docs])
    scores = [(doc, bm25_score(query, doc, avg_len)) for doc in docs]
    scores.sort(key=lambda x: x[1], reverse=True)
    return scores[:top_k]

# Query that uses different terminology than the documents
# "deep learning optimization methods" uses vocabulary not present in most docs
query_mismatch = "deep learning optimization methods"

dense_results = dense_retrieve(
    query_mismatch, document_embeddings, documents, model, top_k=3
)
lexical_results = lexical_retrieve(query_mismatch, documents, top_k=3)

def batch_retrieve(queries, doc_embeddings, docs, model, top_k=3):
    """
    Efficiently retrieve for multiple queries using batched encoding
    and matrix multiplication.

    The key efficiency: encoding all queries simultaneously on GPU
    is much faster than encoding them sequentially.
    """
    # Encode all queries at once (enables GPU parallelism)
    query_embeddings = model.encode(
        queries, convert_to_numpy=True, normalize_embeddings=True
    )

    # Compute all query-document similarities with a single matrix multiply
    # query_embeddings shape: (num_queries, embedding_dim)
    # doc_embeddings.T shape: (embedding_dim, num_docs)
    # Result shape: (num_queries, num_docs)
    all_similarities = np.dot(query_embeddings, doc_embeddings.T)

    results = []
    for i, query in enumerate(queries):
        similarities = all_similarities[i]
        top_indices = np.argsort(similarities)[::-1][:top_k]
        query_results = [(docs[j], similarities[j]) for j in top_indices]
        results.append((query, query_results))

    return results


# Test batch retrieval with three semantically diverse queries
test_queries = [
    "attention mechanisms in deep learning",
    "ensemble methods for classification",
    "processing sequential information",
]

batch_results = batch_retrieve(
    test_queries, document_embeddings, documents, model, top_k=2
)

class DenseRetriever:
    """
    A simple dense retrieval system showing the full pipeline:
    offline indexing and online query serving.
    """

    def __init__(self, model_name="all-MiniLM-L6-v2"):
        self.model = SentenceTransformer(model_name)
        self.documents = []
        self.embeddings = None

    def index(self, documents):
        """
        Offline step: encode all documents and store embeddings.
        This is done once and can be cached to disk.
        """
        self.documents = documents
        self.embeddings = self.model.encode(
            documents,
            convert_to_numpy=True,
            normalize_embeddings=True,
            show_progress_bar=False,
        )
        return self

    def search(self, query, top_k=5):
        """
        Online step: encode query and find nearest neighbors.
        Only the query encoding happens at search time.
        """
        query_emb = self.model.encode(
            query, convert_to_numpy=True, normalize_embeddings=True
        )
        scores = np.dot(self.embeddings, query_emb)
        top_indices = np.argsort(scores)[::-1][:top_k]
        return [
            {
                "rank": i + 1,
                "score": float(scores[idx]),
                "text": self.documents[idx],
            }
            for i, idx in enumerate(top_indices)
        ]

# Build the index (offline step)
retriever = DenseRetriever()
retriever.index(documents)