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:

from sentence_transformers import SentenceTransformer
import numpy as np

# Load a pre-trained sentence transformer model
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# Sample query and documents (vectorize the query and the documents)
query = "How do I fix a broken landing gear?"
documents = ["Report 1 on landing gear failure", "Report 2 on engine problems"]

# Get embeddings for query and documents
query_embedding = model.encode(query)
document_embeddings = model.encode(documents)

# Calculate cosine similarity between query and documents
similarities = np.dot(document_embeddings, query_embedding)

# Retrieve top-k most relevant documents
top_k = np.argsort(similarities)[-5:]
print("Top 5 documents:", [documents[i] for i in top_k])

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:

from sentence_transformers import CrossEncoder

# Load the cross-encoder model
cross_encoder = CrossEncoder("cross-encoder/ms-marco-TinyBERT-L-2-v2")

# Use the cross-encoder to rerank top-k retrieved documents
query_document_pairs = [(query, doc) for doc in documents]
scores = cross_encoder.predict(query_document_pairs)

# Rank documents based on the new scores
top_k_reranked = np.argsort(scores)[-5:]
print("Top 5 reranked documents:", [documents[i] for i in top_k_reranked])

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:

import argparse
import numpy as np
import pandas as pd
from typing import Any
from openvino.runtime import Core
from transformers import AutoTokenizer


def load_openvino_model(model_path: str) -> Core:
    core = Core()
    model = core.read_model(model_path + ".xml")
    compiled_model = core.compile_model(model, "CPU")
    return compiled_model


def rerank(
    compiled_model: Core,
    query: str,
    results: list[str],
    tokenizer: AutoTokenizer,
    batch_size: int,
) -> np.ndarray[np.float32, Any]:
    max_length = 512
    all_logits = []

    # Split results into batches
    for i in range(0, len(results), batch_size):
        batch_results = results[i : i + batch_size]
        inputs = tokenizer(
            [(query, item) for item in batch_results],
            padding=True,
            truncation="longest_first",
            max_length=max_length,
            return_tensors="np",
        )

        # Extract input tensors (convert to NumPy arrays)
        input_ids = inputs["input_ids"].astype(np.int32)
        attention_mask = inputs["attention_mask"].astype(np.int32)
        token_type_ids = inputs.get("token_type_ids", np.zeros_like(input_ids)).astype(
            np.int32
        )

        infer_request = compiled_model.create_infer_request()
        output = infer_request.infer(
            {
                "input_ids": input_ids,
                "attention_mask": attention_mask,
                "token_type_ids": token_type_ids,
            }
        )

        logits = output["logits"]
        all_logits.append(logits)

    all_logits = np.concatenate(all_logits, axis=0)
    return all_logits


def fetch_search_data(search_text: str) -> pd.DataFrame:
    # Usually you would fetch the data from a database
    df = pd.read_csv("cnbc_headlines.csv")
    df = df[~df["Headlines"].isnull()]

    texts = df["Headlines"].tolist()

    # Load the model and rerank
    openvino_model = load_openvino_model("cross-encoder-openvino-model/model")
    tokenizer = AutoTokenizer.from_pretrained("cross-encoder/ms-marco-TinyBERT-L-2-v2")
    rerank_scores = rerank(openvino_model, search_text, texts, tokenizer, batch_size=16)

    # Add the rerank scores to the DataFrame and sort by the new scores
    df["rerank_score"] = rerank_scores
    df = df.sort_values(by="rerank_score", ascending=False)

    return df


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Fetch search results with reranking using OpenVINO"
    )

    parser.add_argument(
        "--search_text",
        type=str,
        required=True,
        help="The search text to use for reranking",
    )

    args = parser.parse_args()

    df = fetch_search_data(args.search_text)
    print(df)

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!