Semantic Search Implementation

Conceptualizing Vector Embeddings

At its core, semantic search relies on the ability to map unstructured data, such as sentences or paragraphs, into a high-dimensional numerical space known as a vector space. An embedding model transforms text into a fixed-length list of floating-point numbers, or 'embeddings'. The fundamental principle is that text snippets with similar meanings are projected closer together in this geometric space, while unrelated texts are pushed further apart. This is achieved by training neural networks on massive datasets to capture complex linguistic patterns, synonyms, and contextual nuances. By representing text as coordinates in a multi-dimensional space, we move beyond the binary constraints of keyword matching. We are not just indexing words; we are indexing the 'concept' behind the words. Understanding this transformation is crucial because it dictates that the quality of your search is entirely dependent on the embedding model's capacity to generalize meaning correctly across the specific domain of your dataset.

from sentence_transformers import SentenceTransformer

# Initialize a pre-trained embedding model
# This model converts text into a 384-dimensional vector
model = SentenceTransformer('all-MiniLM-L6-v2')

# Transform a sample text into a vector representation
text = "The quick brown fox jumps over the lazy dog"
embedding = model.encode(text)

# Output the first 5 dimensions of the embedding vector
print(f"Embedding shape: {embedding.shape}")
print(f"First 5 dimensions: {embedding[:5]}")

Calculating Semantic Similarity

Once text is converted into vectors, we need a mathematical method to determine how close two points are to each other in that high-dimensional space. The most common metric for this is Cosine Similarity, which measures the cosine of the angle between two vectors. Unlike Euclidean distance, which measures the straight-line distance, Cosine Similarity focuses on the orientation of the vectors, effectively ignoring their magnitude. This is particularly useful in natural language processing because text length varies; we want to know if two documents talk about the same topic regardless of whether one is a short sentence and the other is a long paragraph. If two vectors point in the same direction, the cosine of the angle is 1; if they are orthogonal, it is 0. By calculating these similarities, we can quantitatively rank how closely a query matches any given document in our collection, enabling precise retrieval without explicit lexical overlap.

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

# Define two different vectors representing two sentences
vec1 = np.array([[0.1, 0.2, 0.3]])
vec2 = np.array([[0.1, 0.2, 0.35]])

# Calculate cosine similarity between the two vectors
similarity = cosine_similarity(vec1, vec2)

# Higher values indicate greater semantic proximity
print(f"Similarity score: {similarity[0][0]:.4f}")

Building the Vector Index

In a real-world scenario, you will rarely perform a manual comparison between a query and every single document in your database; such an approach is computationally prohibitive as your data grows. To scale effectively, we use a vector index, which organizes embeddings in a structure that allows for 'Approximate Nearest Neighbor' (ANN) search. Instead of an exhaustive scan, an index uses algorithms like HNSW (Hierarchical Navigable Small Worlds) to partition the space, drastically reducing the search surface. When a query comes in, the index traverses these partitions to find candidates that are mathematically likely to be the closest, providing sub-millisecond response times even over millions of entries. Choosing the right indexing strategy involves balancing search speed against recall accuracy. By offloading this organization to a specialized structure, we ensure that our semantic search remains performant and responsive as our collection of documents scales to massive proportions while maintaining the integrity of our semantic mapping.

import faiss
import numpy as np

# Define the dimensions of our vectors
d = 384
# Create a flat index using L2 distance (or Inner Product)
index = faiss.IndexFlatL2(d)

# Add dummy data (100 vectors)
data = np.random.random((100, d)).astype('float32')
index.add(data)

# Perform a search for the top 5 nearest neighbors
query = np.random.random((1, d)).astype('float32')
D, I = index.search(query, 5)
print(f"Nearest neighbor indices: {I}")

Refining Retrieval with Metadata

Raw semantic search is powerful but occasionally returns results that are conceptually relevant yet contextually inappropriate, such as outdated information or private documents. To mitigate this, we implement metadata filtering. Metadata allows us to attach specific attributes to our vector embeddings—such as timestamps, categories, or user access permissions. During the retrieval process, we apply a 'pre-filter' or 'post-filter' to the index. For example, if a user performs a search, we filter the index to only consider vectors where the 'status' is 'public' or the 'date' is within the last year. This hybrid approach combines the semantic richness of vector search with the precision of deterministic databases. By integrating these filters, you ensure that the system respects business logic and data governance, resulting in a more robust and reliable search experience that users can actually trust for professional workflows.

# This logic illustrates how metadata filters prune results
# metadata_store acts as a lookup for vector IDs
metadata_store = {0: {'type': 'public'}, 1: {'type': 'private'}}

def search_filtered(query_vec, target_type):
    # Perform search as usual
    _, indices = index.search(query_vec, 10)
    
    # Filter indices based on metadata match
    filtered_results = [i for i in indices[0] if metadata_store.get(i, {}).get('type') == target_type]
    return filtered_results

print(f"Filtered results: {search_filtered(query, 'public')}")

Optimizing with Re-ranking

The final stage of a sophisticated semantic search pipeline is re-ranking. While the vector index is excellent at narrowing down a million documents to the top 50, its approximation methods may sacrifice precision for speed. A re-ranker is a secondary, more computationally expensive model, often a Cross-Encoder, that examines the query and the top results simultaneously. Unlike vector search, which encodes the query and document separately, a re-ranker analyzes the interaction between the two, assessing how well they fit together as a pair. This process provides a significant boost in precision, as the model can evaluate nuanced relationships that the initial embedding phase missed. By using a two-stage approach—retrieval via vector search, followed by re-ranking via a heavy model—you achieve a perfect balance between massive scale and surgical accuracy, ensuring the most relevant content always reaches the top of the list.

from sentence_transformers import CrossEncoder

# Use a CrossEncoder to perform deep comparison
# CrossEncoders score query-document pairs directly
reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

pairs = [("How does a cell divide?", "Mitosis is the process of cell division."),
         ("How does a cell divide?", "The weather is sunny today.")]

scores = reranker.predict(pairs)
print(f"Re-ranking scores: {scores}")

Interview questions

All Generative AI interview questions →

Check yourself

Take the full Generative AI quiz →