How to Use Rerankers in KDB.AI: Cohere, Jina AI, and Voyage AI

This page provides details on how to improve your search results with Cohere, Jina AI, and Voyage AI rerankers in KDB.AI.

KDB.AI allows you to enhance the relevance and personalization of search results through the integration of Cohere, Voyage AI, and Jina AI rerankers. Integrating rerankers in KDB.AI involves a few steps to set up your environment, prepare your data, and configure the reranking process. Here’s a general guide to follow:

Get started

Before you begin, make sure you have the following:

• Installed Python 3 (versions 3.8 to 3.11), Pip, and Git

• Active KDB.AI Cloud or Server license

• Valid API key for KDB.AI Cloud

• Valid API keys for Cohere, Jina AI, or/and Voyage AI.

• Basic knowledge of working with vector databases and embedding models

• Necessary configurations for interacting with either KDB.AI Cloud or Server

Choose a reranker

Use one reranker at a time (or multiple within the same notebook), depending on your use case. Check out the comparison table below to decide which reranker is suitable for you:

Choose the reranking model

Different rerankers provide different models. Each reranker has a default model. model determines the rerank model. With Cohere you can even create your own model.

reranker = CohereReranker(api_key='COHERE_API_KEY', model='rerank-english-v3.0')

reranker = JinaAIReranker(api_key='JINAAI_API_KEY', model='jina-reranker-v2-base-multilingual')

reranker = VoyageAIReranker(api_key='VOYAGEAI_API_KEY', model='rerank-2')

Reranker set up

Install dependencies

Use OpenAI to embed the data and queries, and Langchain to chunk the data.

!pip install langchain
!pip install openai

Import dependencies

import os
import numpy as np
import pandas as pd
import langchain
import openai
from langchain_text_splitters import RecursiveCharacterTextSplitter
from openai import OpenAI
import kdbai_client as kdbai 

os.environ["OPENAI_API_KEY"] = "OPENAI_API_KEY"

Set up your KDB.AI database

session = kdbai.Session(endpoint='http://localhost:8082')

schema = [
    {'name': 'id', 'type': 'int16'},
    {'name': 'content', 'type': 'str'},
    {'name': 'embeddings', 'type': 'float64s'}
]

indexes = [
    {'type': 'flat', 'name': 'flat', 'column': 'embeddings',  'params': {'dims': 1536}},
    {'type': 'hnsw', 'name': 'fast_hnsw', 'column': 'embeddings', 'params': {'dims': 1536,'M': 8, 'efConstruction': 8}},
    {'type': 'hnsw', 'name': 'accurate_hnsw','column': 'embeddings', 'params': {'dims': 1536,'M': 64, 'efConstruction':256}} 
]

database = session.database('rerank')

table = database.create_table('rerank', schema=schema, indexes=indexes)

Prepare your data

data refers to any type of data you want to use. Below is an example of how to prepare data before inserting it to your KDB.AI database.

client = OpenAI()

def get_data_chunks(data: str) -> list[str]:
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=384,
        chunk_overlap=20,
        length_function=len,
        is_separator_regex=False,
    )
    texts = text_splitter.create_documents([data])
    return texts

def get_embedding(text, model):
    text = text.replace("\n", " ")
    return client.embeddings.create(input=[text], model=model).data[0].embedding

df = pd.DataFrame(columns = ["id", "content", "embeddings"])

data_chunks = get_data_chunks(data)

for num, data_chunk in enumerate(data_chunks):
    text = data_chunk.page_content
    text_embeddings = get_embedding(text, model='text-embedding-3-small')
    df.loc[len(df), :] = (num, text, text_embeddings)

table.insert(df)

Execute search and rerank

query1 = 'What are the key things president delivered in the speech?'
query2 = 'What did president say about inflation?'

query1_embeddings = get_embedding(query1, model='text-embedding-3-small')
query2_embeddings = get_embedding(query2, model='text-embedding-3-small')

from kdbai_client.rerankers import CohereReranker

reranker = CohereReranker(api_key='COHERE_API_KEY')

table.search_and_rerank(vectors={"flat": [query1_embeddings, query2_embeddings]}, n=3, reranker=reranker, queries=[query1, query2], text_column='content')

Fine tune the reranker

Choose the number of documents to rerank

The overfetch_factor value determines how many documents are processed by the reranker. The overfetch_factor is multiplied by n (nearest neighbors). By default, the overfetch_factor is 2, so if n is 3, the amount of documents processed by the reranker is 6. You only get n results returned. To further fine tune the results, adjust the overfetch_factor.

reranker = CohereReranker(api_key='COHERE_API_KEY', overfetch_factor=4)

reranker = JinaAIReranker(api_key='JINAAI_API_KEY', overfetch_factor=4)

reranker = VoyageAIReranker(api_key='VOYAGEAI_API_KEY', overfetch_factor=4)

Next steps

Now that you're familiar with reranking, try the following:

• Head to the Reranking with Cohere docs.

• Read our Jina AI-related blog article The end of high dimensions to discover how Matryoshka learning is revolutionizing AI search forever.