Building a Scalable Hybrid Search System Lessons from Papers with Code

The Challenge of Searching Academic Research

Searching for AI research papers is not the same as searching for regular web content. A robust research search engine must handle exact queries like a specific arXiv ID, but it also needs to understand conceptual queries. For instance, a user searching for "small language models for code generation" expects results even if those exact words don't appear in a paper's text. It must also recognize navigational requests like "the original BERT paper" and be tolerant of typos or incomplete titles.

To meet these demands, the team behind Papers with Code didn't just rely on a single method. They built a sophisticated hybrid search system. A key part of this system relies on various Python tools and best practices for managing data and compute.

This post delves into the architecture they built, focusing on the design decisions that make it powerful and fast.

A Tale of Two Retrieval Methods: Keyword and Vector Search

The core insight is that keyword search and vector search are complementary. Keyword search (like PostgreSQL's full-text search) is great for finding exact matches, identifiers, and rare names. On the other hand, vector search uses dense embeddings to find semantically similar terms, capturing the 'fuzzy' meaning behind a query.

The Architecture: Offline vs. Online

The system is deliberately split into an offline corpus build and an online search service. This separation is key to balancing cost, performance, and reliability.

1. Offline Corpus Build (Jobs): The expensive, throughput-oriented work of embedding the entire paper corpus (over 110,000 papers) is done as a batch job on a GPU. This is powered by Hugging Face Jobs, which provides burstable compute.
2. Durable Storage (Buckets): The output of these jobs (the vector embeddings) is stored in a Hugging Face Storage Bucket. This acts as the connective tissue, ensuring a durable handoff between the ephemeral compute and the production database.
3. Online Search Service (Inference Endpoints): Only the small, latency-sensitive task of embedding a live user query is on the request path. This is handled by a Hugging Face Inference Endpoint, which is designed for low-latency responses.

Here’s a simplified look at the query embedding step:

# Pseudocode for the online query embedding process
import requests

# 1. Send the user's query to the Inference Endpoint
response = requests.post(
    "YOUR_INFERENCE_ENDPOINT_URL",
    headers={"Authorization": "Bearer YOUR_TOKEN"},
    json={"inputs": query_text}
)

# 2. Process the response to get the query vector
query_vector = response.json()["embedding"]

# 3. Use this vector to perform a cosine distance search in PostgreSQL (pgvector)
# SELECT paper_id, embedding <=> CAST(:query_vector AS halfvec(256)) AS distance
# FROM paper_embeddings
# WHERE generation_id = :active_generation
# ORDER BY embedding <=> CAST(:query_vector AS halfvec(256))
# LIMIT 50;

Key Design Patterns

• Strict Embedding Contract: They treat the embedding format as a versioned API. This includes the model revision, output dimensions, input format, and normalization method. This prevents subtle failures when models are updated or code is refactored.
• Resumable and Reproducible Jobs: By organizing artifacts under immutable run prefixes and using checksums, they can safely retry failed jobs or trace any result back to its exact source data and model revision.
• Graceful Degradation: The system is designed for cold starts. If the Inference Endpoint is scaling up or times out, the query client immediately falls back to lexical search. This ensures users always get results, even if they aren't as semantically rich.

Lessons Learned and the Road Ahead

The team's experience offers valuable lessons for anyone building similar AI-powered search features.

1. Separate Throughput from Latency: Batch embedding and query embedding are different infrastructure problems. Optimize for cost and throughput in the former, and for availability and speed in the latter.
2. Make Storage the Contract: Use a storage bucket as an explicit, auditable boundary between your compute and production systems.
3. Pin More Than the Model Name: The revision, dimension, and prompt all affect the final retrieval quality. Store and validate them all.
4. Design for Cold Starts: If you use scale-to-zero, you must have a fast fallback. Hybrid search provides this naturally.
5. Smaller Vectors Can Be a Feature: Matryoshka Representation Learning (MRL) allows you to trade off quality for speed and storage. In their pilot, 256 dimensions preserved recall while using a fraction of the storage.
6. Activation Should Be Boring: New embedding generations should be imported and validated alongside the current one before being activated atomically. This makes rollback a simple configuration change.

Limitations and Considerations

While this architecture is powerful, it's important to note that hybrid search isn't always the best solution. The team recommends starting with keyword search as a cheap baseline and only adding semantic and/or hybrid search when you see a reasonable boost in retrieval quality. Adding a reranker can also improve results but introduces additional latency and overhead.

Conclusion

The hybrid search system powering Papers with Code is a masterclass in pragmatic system design. By separating offline and online workloads, using a storage bucket as a contract, and designing for graceful degradation, they've built a system that is both powerful and reliable. The use of Hugging Face Jobs, Buckets, and Inference Endpoints provides a blueprint for building scalable AI infrastructure.

For those looking to further optimize their retrieval systems, exploring the capabilities of newer models is a great next step. You can also read about other exciting developments in the Python ecosystem, like the new features in the upcoming Python 3.15 release.

함께 보면 좋은 글

• Scale Python to a Cloud Cluster with Ray: A Practical AWS Tutorial
• What's Coming in Python 3.15? A Deep Dive into Alpha 5 Features