Skip to content
Snippets Groups Projects
Unverified Commit a88bc070 authored by Ravi Theja's avatar Ravi Theja Committed by GitHub
Browse files

Add mixedbread reranker cookbook (#11536)

parent 53e9f523
No related branches found
No related tags found
No related merge requests found
%% Cell type:markdown id:964030f7-40e4-4398-a5ab-668aabcf3bad tags:
<a href="https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/cookbooks/mixedbread_reranker.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
%% Cell type:markdown id:360313ab-9393-430e-9647-e0d5545809b9 tags:
# mixedbread Rerank Cookbook
mixedbread.ai has released three fully open-source reranker models under the Apache 2.0 license. For more in-depth information, you can check out their detailed [blog post](https://www.mixedbread.ai/blog/mxbai-rerank-v1). The following are the three models:
1. `mxbai-rerank-xsmall-v1`
2. `mxbai-rerank-base-v1`
3. `mxbai-rerank-large-v1`
In this notebook, we'll demonstrate how to use the `mxbai-rerank-base-v1` model with the `SentenceTransformerRerank` module in LlamaIndex. This setup allows you to seamlessly swap in any reranker model of your choice using the `SentenceTransformerRerank` module to enhance your RAG pipeline.
%% Cell type:markdown id:856ecfdc-04fa-4fe9-a81c-9a5858cd4a6d tags:
### Installation
%% Cell type:code id:bfb5314f-e6c7-409c-86df-8e1a5ca59adb tags:
``` python
!pip install llama-index
!pip install sentence-transformers
```
%% Cell type:markdown id:5f5393fb-b410-4769-9380-0ef90a33b82e tags:
### Set API Keys
%% Cell type:code id:a9782acf-b0ab-4933-bb41-27cd2a02b5dd tags:
``` python
import os
os.environ["OPENAI_API_KEY"] = "YOUR OPENAI API KEY"
```
%% Cell type:code id:b7596ddf-e1de-4098-81f3-fce504d2da94 tags:
``` python
from llama_index.core import (
VectorStoreIndex,
SimpleDirectoryReader,
)
from llama_index.core.postprocessor import SentenceTransformerRerank
```
%% Cell type:markdown id:8011ff9c-2b82-47b4-983f-4fafc29e3127 tags:
### Download Data
%% Cell type:code id:6dd335cb-900b-462f-987a-d4af2aac88fa tags:
``` python
!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'
```
%% Output
--2024-03-01 09:52:09-- https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt
Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.110.133, 185.199.108.133, 185.199.109.133, ...
Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.110.133|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 75042 (73K) [text/plain]
Saving to: ‘data/paul_graham/paul_graham_essay.txt’
data/paul_graham/pa 100%[===================>] 73.28K --.-KB/s in 0.007s
2024-03-01 09:52:09 (9.86 MB/s) - ‘data/paul_graham/paul_graham_essay.txt’ saved [75042/75042]
%% Cell type:markdown id:e482b09c-a0df-4788-a75b-a33ade7001d1 tags:
### Load Documents
%% Cell type:code id:342c91b8-301f-40ed-9d09-9acdb1bbdc44 tags:
``` python
documents = SimpleDirectoryReader("./data/paul_graham/").load_data()
```
%% Cell type:markdown id:8afdfeb1-57ae-4d2b-ae73-683db205be32 tags:
### Build Index
%% Cell type:code id:47c335e9-dd4d-475c-bade-e2a588e33294 tags:
``` python
index = VectorStoreIndex.from_documents(documents=documents)
```
%% Cell type:markdown id:f1ab8157-dbcb-4588-9b3c-5bd2fc4a721e tags:
### Define postprocessor for `mxbai-rerank-base-v1` reranker
%% Cell type:code id:3fcc5590-2e58-4a7e-8b18-a7153c06d1ff tags:
``` python
from llama_index.core.postprocessor import SentenceTransformerRerank
postprocessor = SentenceTransformerRerank(
model="mixedbread-ai/mxbai-rerank-base-v1", top_n=2
)
```
%% Cell type:markdown id:c7c81b0d-0449-4092-80cb-88080e69f980 tags:
### Create Query Engine
We will first retrieve 10 relevant nodes and pick top-2 nodes using the defined postprocessor.
%% Cell type:code id:e1b23700-15ae-4f1a-9443-43eb1eecab5f tags:
``` python
query_engine = index.as_query_engine(
similarity_top_k=10,
node_postprocessors=[postprocessor],
)
```
%% Cell type:markdown id:93871f9c-8871-4f43-8ee9-b3ca4e403d86 tags:
### Test Queries
%% Cell type:code id:658d3092-7d86-4520-83a2-c3e630dc02b6 tags:
``` python
response = query_engine.query(
"Why did Sam Altman decline the offer of becoming president of Y Combinator?",
)
print(response)
```
%% Output
Sam Altman initially declined the offer of becoming president of Y Combinator because he wanted to start a startup focused on making nuclear reactors.
%% Cell type:code id:497e715e-3f7a-4140-a3ba-34356e473702 tags:
``` python
response = query_engine.query(
"Why did Paul Graham start YC?",
)
print(response)
```
%% Output
Paul Graham started YC because he and his partners wanted to create an investment firm where they could implement their own ideas and provide the kind of support to startups that they felt was lacking when they were founders themselves. They aimed to not only make seed investments but also assist startups with various aspects of setting up a company, similar to the help they had received from others in the past.
...@@ -332,4 +332,5 @@ maxdepth: 1 ...@@ -332,4 +332,5 @@ maxdepth: 1
/examples/node_postprocessor/rankGPT.ipynb /examples/node_postprocessor/rankGPT.ipynb
/examples/node_postprocessor/ColbertRerank.ipynb /examples/node_postprocessor/ColbertRerank.ipynb
/examples/node_postprocessor/JinaRerank.ipynb /examples/node_postprocessor/JinaRerank.ipynb
/cookbooks/mixedbread_reranker.ipynb
``` ```
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment