Photo by real-napster on Pixabay

In one of my recent projects, I had to build a semantic search system that could scale with high performance and deliver real-time responses for report searches. We used PostgreSQL with pgvector on AWS RDS, paired with AWS Lambda, to achieve this. The challenge was to allow users to search using natural language queries instead of relying on rigid keywords, all while ensuring responses were under 1-2 seconds or even below and could only leverage CPU resources.

In this post, I will walk through the steps I took to build this search system, from retrieval to reranking, and the optimizations made using OpenVINO and intelligent batching for tokenization.

Overview of Semantic Search: Retrieval and Reranking

Modern state-of-the-art search systems usually consist of two main steps: retrieval and reranking.

1) Retrieval: The first step involves retrieving a subset of relevant documents based on the user query. This can be done using pre-trained embeddings models, such as OpenAI's small and large embeddings, Cohere's Embed models, or Mixbread’s mxbai embeddings. Retrieval focuses on narrowing down the pool of documents by measuring their similarity to the query.

Here's a simplified example using Huggingface's sentence-transformers library for retrieval which is one of my favorite libraries for this:

2) Reranking: Once the most relevant documents have been retrieved, we further improve the ranking of these documents using a cross-encoder model. This step re-evaluates each document in relation to the query more accurately, focusing on deeper contextual understanding.
Reranking is beneficial because it adds an additional layer of refinement by scoring the relevance of each document more precisely.

Here's a code example for reranking using cross-encoder/ms-marco-TinyBERT-L-2-v2, a lightweight cross-encoder:

Identifying Bottlenecks: The Cost of Tokenization and Prediction

During the development, I found that the tokenization and prediction stages were taking quite long when handling 1,000 reports with default settings for sentence-transformers. This created a performance bottleneck, especially since we aimed for real-time responses.

Below I profiled my code using SnakeViz to visualize the performances:

As you can see, the tokenization and prediction steps are disproportionately slow, leading to significant delays in serving search results. Overall it took like 4-5 seconds on average. This is due to the fact that there are blocking operations between the tokenization and prediction steps. If we also add up other operations like database call, filtering etc, we easily ended up with 8-9 seconds in total.

Optimizing Performance with OpenVINO

The question I faced was: Can we make it faster? The answer is yes, by leveraging OpenVINO, an optimized backend for CPU inference. OpenVINO helps accelerate deep learning model inference on Intel hardware, which we use on AWS Lambda.

Code Example for OpenVINO Optimization
Here’s how I integrated OpenVINO into the search system to speed up inference:

With this approach we could get a 2-3x speedup reducing the original 4-5 seconds to 1-2 seconds. The full working code is on Github.

Fine-Tuning for Speed: Batch Size and Tokenization

Another critical factor in improving performance was optimizing the tokenization process and adjusting the batch size and token length. By increasing the batch size (batch_size=16) and reducing the token length (max_length=512), we could parallelize the tokenization and reduce the overhead of repetitive operations. In our experiments, we found that a batch_size between 16 and 64 worked well, with anything larger degrading performance. Similarly, we settled on a max_length of 128, which is viable if the average length of your reports is relatively short. With these changes, we achieved an overall 8x speed-up, reducing the reranking time to under 1 second, even on CPU.

In practice, this meant experimenting with different batch sizes and token lengths to find the right balance between speed and accuracy for your data. By doing so, we saw significant improvements in response times, making the search system scalable even with 1,000 reports.

Conclusion

By using OpenVINO and optimizing tokenization and batching, we were able to build a high-performance semantic search system that meets real-time requirements on a CPU-only setup. In fact, we experienced a 8x speedup overall. The combination of retrieval using sentence-transformers and reranking with a cross-encoder model creates a powerful, user-friendly search experience.

If you’re building similar systems with constraints on response time and computational resources, I highly recommend exploring OpenVINO and intelligent batching to unlock better performance.

Hopefully, you enjoyed this article. If you found this article useful, give me a like so others can find it too, and share it with your friends. Follow me on Linkedin to stay up-to-date with my work. Thanks for reading!

The above is the detailed content of Building a Fast and Efficient Semantic Search System Using OpenVINO and Postgres. For more information, please follow other related articles on the PHP Chinese website!