Building Semantic Search Systems: From Retrieval to Ranking
Traditional search engines rely heavily on keyword matching. Algorithms like BM25 look for exact word overlapping between the user's query and indexed documents. While highly efficient, keyword search fails when users search for concepts using synonyms, parallel phrasing, or natural conversational questions.
Semantic search resolves this by understanding the intent and contextual meaning behind queries. By leveraging deep learning models, semantic search projects both queries and documents into a shared vector space, retrieving documents based on conceptual similarity rather than character matching.
However, building an production-grade semantic search system is not as simple as querying a vector database. To achieve optimal search relevance, modern architectures separate search into a two-stage process: Retrieval (filtering down candidates quickly) and Ranking (sorting candidates with highly precise but computationally heavy models).
The Two-Stage Search Architecture
The two-stage search paradigm strikes a balance between computational throughput and ranking accuracy:
- Stage 1: Retrieval (Bi-Encoder): Queries and documents are encoded independently into vectors. A fast approximate nearest neighbor (ANN) search retrieves the top
Mcandidates (e.g.,M = 100) from millions of documents. - Stage 2: Re-ranking (Cross-Encoder): The retrieved
Mdocuments are paired with the query and analyzed together by a specialized transformer that models attention across both inputs. This yields highly precise similarity scores for final sorting.
Semantic Search Systems: Semantic search pipeline showing document indexing and real-time query vector search
Bi-Encoders vs. Cross-Encoders
- Bi-Encoders: Process query
qand documentdseparately:s(q, d) = \cos(\mathbf{E}(q), \mathbf{E}(d)). Because document embeddings can be pre-calculated and stored in a vector database, search is incredibly fast. However, it cannot capture fine-grained interactions between specific words in the query and document. - Cross-Encoders: Process query and document simultaneously by concatenating them:
\mathbf{E}(\text{"[CLS] "} + q + \text{" [SEP] "} + d). Self-attention layers compute cross-token alignments between the query and text. This is highly accurate but computationally expensive, as it requires running the transformer network in real-time for every candidate document.
Python Walkthrough: Building the Pipeline
Below is a complete Python implementation of a semantic search engine with Cross-Encoder re-ranking, using the popular sentence-transformers library.
from sentence_transformers import SentenceTransformer, CrossEncoder, util
import numpy as np
# 1. Initialize our models
# Bi-Encoder: Used to generate dense vectors for quick retrieval
bi_encoder = SentenceTransformer("all-MiniLM-L6-v2")
# Cross-Encoder: Used to precisely score the query-document pairs
cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
# 2. Our Document Corpus (Mock Knowledge Base)
corpus = [
"A computer network consists of two or more computers connected together to share resources.",
"Python is an interpreted, high-level, general-purpose programming language designed by Guido van Rossum.",
"Web browsers are software applications used to locate, retrieve, and display content on the World Wide Web.",
"Deep learning is a subset of machine learning based on artificial neural networks with representation learning.",
"To compile source code, a software developer uses tools like GCC or Clang to produce machine instructions.",
"Fruit-bearing trees like apples, pears, and oranges thrive in well-drained, organic soil environments."
]
# Pre-calculate and cache embeddings for the corpus (offline index phase)
print("Encoding corpus documents...")
corpus_embeddings = bi_encoder.encode(corpus, convert_to_tensor=True)
def search(query: str, top_k_retrieve: int = 4, top_k_rank: int = 2):
print(f"\nUser Query: '{query}'")
# ----------------------------------------------------
# Stage 1: Retrieval (Bi-Encoder)
# ----------------------------------------------------
# Encode the user query
query_embedding = bi_encoder.encode(query, convert_to_tensor=True)
# Compute cosine similarity between query and all documents
cos_scores = util.cos_sim(query_embedding, corpus_embeddings)[0]
# Get top retrieved candidates
top_results_idx = np.argsort(cos_scores.cpu().numpy())[::-1][:top_k_retrieve]
print("\n--- [Stage 1] Top Retrieved Candidates ---")
retrieved_pairs = []
for idx in top_results_idx:
score = cos_scores[idx].item()
doc = corpus[idx]
print(f"Score: {score:.4f} | Doc: {doc}")
retrieved_pairs.append((doc, idx))
# ----------------------------------------------------
# Stage 2: Re-ranking (Cross-Encoder)
# ----------------------------------------------------
# Prepare query-document pairs for the cross-encoder
cross_inputs = [[query, pair[0]] for pair in retrieved_pairs]
# Predict similarity scores
cross_scores = cross_encoder.predict(cross_inputs)
# Sort candidates by their cross-encoder scores
ranked_indices = np.argsort(cross_scores)[::-1]
print("\n--- [Stage 2] Re-ranked Results (Final Output) ---")
for rank_idx in ranked_indices[:top_k_rank]:
score = cross_scores[rank_idx]
doc = retrieved_pairs[rank_idx][0]
print(f"Rank Score: {score:.4f} | Doc: {doc}")
if __name__ == "__main__":
# Test query targeting programming concept (should retrieve Python & compilation)
search(query="how do I build software from code files?", top_k_retrieve=3, top_k_rank=2)
Optimizing the Retrieval Pipeline
To scale this pipeline to millions of pages:
- Chunking Strategies: Do not embed entire documents. Break texts into 150-250 word semantic chunks. Implement overlapping sliding windows (e.g., 20-50 tokens overlap) to avoid loss of context at chunk boundaries.
- Hybrid Search: Combine BM25 scores with Vector search scores using Reciprocal Rank Fusion (RRF):
\text{RRF\_Score}(d \in D) = \sum_{m \in M} \frac{1}{k + r_m(d)}
where r_m(d) is the rank of document d in search system m, and k is a constant (typically 60). RRF allows you to easily combine keyword matching (perfect for exact names, IDs, serial numbers) with semantic matching (perfect for abstract conceptual queries).
3. Metadata Filtering: Pre-filter vector spaces by adding metadata constraints (e.g., user_id, category, tenant_id) directly to the vector query, preventing the index from spending computing cycles comparing embeddings of unreachable documents.