Semantic Search

Two-stage retrieval pipeline with SentenceTransformers

from sentence_transformers import SentenceTransformer, CrossEncoder
import numpy as np

# Stage 1: Bi-encoder retrieval
bi_encoder = SentenceTransformer('intfloat/e5-large-v2')
docs = ["heart attack symptoms", "myocardial infarction treatment", "chest pain causes", ...]

# Pre-compute document embeddings (done once, cached)
doc_embeddings = bi_encoder.encode(["passage: " + d for d in docs], normalize_embeddings=True)

query = "query: what causes heart attacks"
query_emb = bi_encoder.encode([query], normalize_embeddings=True)

# Retrieve top-100 by cosine similarity
scores = (query_emb @ doc_embeddings.T)[0]
top100_indices = np.argsort(scores)[::-1][:100]
top100_docs = [docs[i] for i in top100_indices]

# Stage 2: Cross-encoder reranking
cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-12-v2')
pairs = [("what causes heart attacks", doc) for doc in top100_docs]
rerank_scores = cross_encoder.predict(pairs)

# Final top-10
top10 = sorted(zip(rerank_scores, top100_docs), reverse=True)[:10]
for score, doc in top10:
    print(f"{score:.3f}  {doc}")

DPR contrastive loss: maximize similarity between (query, positive passage) pairs while minimizing similarity to negative passages. sim(q, p) = dot product of BERT encodings. In-batch negatives: other passages in the batch serve as easy negatives; BM25 hard negatives added for harder training signal.

End-to-end semantic search with pgvector (PostgreSQL)

import psycopg2
from sentence_transformers import SentenceTransformer

model = SentenceTransformer('intfloat/e5-large-v2')

# ── INDEXING PHASE (run once / on updates) ─────────────────────────
conn = psycopg2.connect("postgresql://user:pass@localhost/mydb")
cur = conn.cursor()
cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
cur.execute(
    "CREATE TABLE IF NOT EXISTS documents "
    "(id SERIAL PRIMARY KEY, content TEXT, embedding vector(1024))"
)
cur.execute(
    "CREATE INDEX IF NOT EXISTS doc_idx ON documents "
    "USING hnsw (embedding vector_cosine_ops) WITH (m=16, ef_construction=64)"
)

docs = ["Introduction to neural networks...", "Attention mechanism explained..."]
embeddings = model.encode(["passage: " + d for d in docs], normalize_embeddings=True)
for doc, emb in zip(docs, embeddings):
    cur.execute("INSERT INTO documents (content, embedding) VALUES (%s, %s)",
                (doc, emb.tolist()))
conn.commit()

# ── RETRIEVAL PHASE (per query) ─────────────────────────────────────
query = "query: how does self-attention work"
q_emb = model.encode([query], normalize_embeddings=True)[0]
cur.execute(
    "SELECT content, 1 - (embedding <=> %s::vector) AS sim "
    "FROM documents ORDER BY embedding <=> %s::vector LIMIT 10",
    (q_emb.tolist(), q_emb.tolist())
)
for text, sim in cur.fetchall():
    print(f"{sim:.3f}  {text[:80]}")

ColBERT MaxSim scoring: for each query token embedding, find the most similar document token embedding (MaxSim). Sum these per-query-token max scores to get the final relevance score. Documents are pre-encoded to token vectors and stored in compressed form.

1. What is the difference between semantic search and lexical search in terms of query handling? (Answer: Lexical search (BM25): tokenises query and document, computes term frequency and inverse document frequency scores. Handles: exact term matching, rare technical terms, product codes, proper nouns. Fails on: paraphrase, synonym, cross-language queries. Semantic search (bi-encoder): converts query and document to dense embeddings, retrieves by vector similarity. Handles: paraphrase ('cheap car' matches 'affordable vehicle'), intent ('how to fix' matches troubleshooting docs). Fails on: exact code/ID matching, very rare domain terms not in training data.)
2. What is a bi-encoder vs cross-encoder for semantic search and when do you use each? (Answer: Bi-encoder: encode query and document INDEPENDENTLY → embeddings stored offline. Query at search time: encode query, ANN search. Very fast (O(1) per query). Good recall but not optimal precision. Cross-encoder: jointly encode (query, document) pair → single relevance score. Much more accurate (attends across both). Cannot pre-encode — must run for every (query, document) pair at query time. Too slow for first-stage retrieval. Architecture: bi-encoder for retrieval (recall), cross-encoder for re-ranking top-k results (precision).)
3. What is BEIR (Benchmarking IR) and what has it revealed about semantic search generalisation? (Answer: BEIR (Thakur et al. 2021): 18 diverse information retrieval benchmarks (MSMARCO, TREC-COVID, NQ, ArguAna, etc.) covering different domains and query types. Key finding: models fine-tuned on one retrieval dataset (MSMARCO) significantly underperform on other domains — generalisation is poor. Dense retrieval often underperforms BM25 on out-of-domain data. Hybrid BM25 + dense retrieval consistently outperforms either alone. Conclusion: domain-specific fine-tuning or robust generalisation training is essential for production semantic search.)
4. What is late interaction semantic search (ColBERT) and when is it preferred? (Answer: ColBERT: encode query and document separately into per-token embeddings (not pooled sentence vectors). Scoring: sum of max cosine similarities between each query token and its best-matching document token (MaxSim). More accurate than bi-encoder (richer interaction) while allowing offline document indexing (unlike cross-encoder). Storage cost: all document token embeddings stored (much larger than bi-encoder). Use when: bi-encoder recall is insufficient, cross-encoder is too slow, and storage cost is acceptable.)
5. What is approximate nearest neighbour (ANN) indexing and which algorithm does FAISS use by default? (Answer: ANN trades exact nearest neighbour accuracy for massive speed gains. FAISS (Facebook AI Similarity Search) supports multiple index types: IndexFlatL2/IP (exact, brute force — baseline). IndexIVFFlat: inverted file index with cluster-based search — 10–100× faster, ~95% recall with nprobe=64. IndexHNSW: graph-based, excellent recall and speed balance, memory-intensive. IndexIVFPQ: inverted file + product quantisation for memory compression. Production recommendation: IVF_HNSW for high-recall/low-latency; IVFPQ for memory-constrained deployments.)