Implementing reranking in ZenML

Learn how to implement reranking in ZenML.

We already have a working RAG pipeline, so inserting a reranker into the pipeline is relatively straightforward. The reranker will take the retrieved documents from the initial retrieval step and reorder them in terms of the query that was used to retrieve them.

How and where to add reranking

We'll use the rerankers package to handle the reranking process in our RAG inference pipeline. It's a relatively low-cost (in terms of technical debt and complexity) and lightweight dependency to add into our pipeline. It offers an interface to most of the model types that are commonly used for reranking and means we don't have to worry about the specifics of each model.

This package provides a Reranker abstract class that you can use to define your own reranker. You can also use the provided implementations to add reranking to your pipeline. The reranker takes the query and a list of retrieved documents as input and outputs a reordered list of documents based on the reranking scores. Here's a toy example:

from rerankers import Reranker

ranker = Reranker('cross-encoder')

texts = [

"I like to play soccer",

"I like to play football",

"War and Peace is a great book"

"I love dogs",

"Ginger cats aren't very smart",

"I like to play basketball",

results = ranker.rank(query="What's your favorite sport?", docs=texts)

And results will look something like this:

RankedResults(

results=[

Result(doc_id=5, text='I like to play basketball', score=-0.46533203125, rank=1),

Result(doc_id=0, text='I like to play soccer', score=-0.7353515625, rank=2),

Result(doc_id=1, text='I like to play football', score=-0.9677734375, rank=3),

Result(doc_id=2, text='War and Peace is a great book', score=-5.40234375, rank=4),

Result(doc_id=3, text='I love dogs', score=-5.5859375, rank=5),

Result(doc_id=4, text="Ginger cats aren't very smart", score=-5.94921875, rank=6)

],

query="What's your favorite sport?",