Back to Llama Index

RankGPT Reranker Demonstration (Van Gogh Wiki)

docs/examples/node_postprocessor/rankGPT.ipynb

0.14.216.7 KB
Original Source

<a href="https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/node_postprocessor/rankGPT.ipynb" target="_parent"></a>

RankGPT Reranker Demonstration (Van Gogh Wiki)

This demo integrates RankGPT into LlamaIndex as a reranker.

Paper: Is ChatGPT Good at Search? Investigating Large Language Models as Re-Ranking Agents

the idea of RankGPT:

  • it is a zero-shot listwise passage reranking using LLM (ChatGPT or GPT-4 or other LLMs)
  • it applies permutation generation approach and sliding window strategy to rerank passages efficiently.

In this example, we use Van Gogh's wikipedia as an example to compare the Retrieval results with/without RankGPT reranking. we showcase two models for RankGPT:

  • OpenAI GPT3.5
  • Mistral model.
python
%pip install llama-index-postprocessor-rankgpt-rerank
%pip install llama-index-llms-huggingface
%pip install llama-index-llms-huggingface-api
%pip install llama-index-llms-openai
%pip install llama-index-llms-ollama
python
import nest_asyncio

nest_asyncio.apply()
python
import logging
import sys

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.postprocessor import LLMRerank
from llama_index.llms.openai import OpenAI
from IPython.display import Markdown, display
python
import os

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

Load Data, Build Index

python
from llama_index.core import Settings

Settings.llm = OpenAI(temperature=0, model="gpt-3.5-turbo")
Settings.chunk_size = 512

Download Van Gogh wiki from Wikipedia

python
from pathlib import Path
import requests

wiki_titles = [
    "Vincent van Gogh",
]


data_path = Path("data_wiki")

for title in wiki_titles:
    response = requests.get(
        "https://en.wikipedia.org/w/api.php",
        params={
            "action": "query",
            "format": "json",
            "titles": title,
            "prop": "extracts",
            "explaintext": True,
        },
    ).json()
    page = next(iter(response["query"]["pages"].values()))
    wiki_text = page["extract"]

    if not data_path.exists():
        Path.mkdir(data_path)

    with open(data_path / f"{title}.txt", "w") as fp:
        fp.write(wiki_text)
python
# load documents
documents = SimpleDirectoryReader("./data_wiki/").load_data()

Build vector store index for this Wikipedia page

python
index = VectorStoreIndex.from_documents(
    documents,
)

Retrieval + RankGPT reranking

Steps:

  1. Setting up retriever and reranker (as an option)
  2. Retrieve results given a search query without reranking
  3. Retrieve results given a search query with RankGPT reranking enabled
  4. Comparing the results with and without reranking
python
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core import QueryBundle
from llama_index.postprocessor.rankgpt_rerank import RankGPTRerank

import pandas as pd
from IPython.display import display, HTML


def get_retrieved_nodes(
    query_str, vector_top_k=10, reranker_top_n=3, with_reranker=False
):
    query_bundle = QueryBundle(query_str)
    # configure retriever
    retriever = VectorIndexRetriever(
        index=index,
        similarity_top_k=vector_top_k,
    )
    retrieved_nodes = retriever.retrieve(query_bundle)

    if with_reranker:
        # configure reranker
        reranker = RankGPTRerank(
            llm=OpenAI(
                model="gpt-3.5-turbo-16k",
                temperature=0.0,
                api_key=OPENAI_API_KEY,
            ),
            top_n=reranker_top_n,
            verbose=True,
        )
        retrieved_nodes = reranker.postprocess_nodes(
            retrieved_nodes, query_bundle
        )

    return retrieved_nodes


def pretty_print(df):
    return display(HTML(df.to_html().replace("\\n", "
")))


def visualize_retrieved_nodes(nodes) -> None:
    result_dicts = []
    for node in nodes:
        result_dict = {"Score": node.score, "Text": node.node.get_text()}
        result_dicts.append(result_dict)

    pretty_print(pd.DataFrame(result_dicts))

Retrieval top 3 results without Reranking

python
new_nodes = get_retrieved_nodes(
    "Which date did Paul Gauguin arrive in Arles?",
    vector_top_k=3,
    with_reranker=False,
)

Expected result is:

After much pleading from Van Gogh, Gauguin arrived in Arles on 23 October and, in November, the two painted together. Gauguin depicted Van Gogh in his The Painter of Sunflowers;

python
visualize_retrieved_nodes(new_nodes)

Finding: the right result is ranked at 2nd without reranking

Retrieve and Reranking top 10 results using RankGPT and return top 3

python
new_nodes = get_retrieved_nodes(
    "Which date did Paul Gauguin arrive in Arles ?",
    vector_top_k=10,
    reranker_top_n=3,
    with_reranker=True,
)
python
visualize_retrieved_nodes(new_nodes)

Finding: After RankGPT reranking, the top 1st result is the right text containing the answer

Using other LLM for RankGPT reranking

Using Ollama for serving local Mistral models

python
from llama_index.llms.ollama import Ollama

llm = Ollama(
    model="mistral",
    request_timeout=30.0,
    # Manually set the context window to limit memory usage
    context_window=8000,
)
python
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core import QueryBundle
import pandas as pd
from IPython.display import display, HTML
from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
from llama_index.llms.huggingface import HuggingFaceLLM

from llama_index.postprocessor.rankgpt_rerank import RankGPTRerank


def get_retrieved_nodes(
    query_str, vector_top_k=5, reranker_top_n=3, with_reranker=False
):
    query_bundle = QueryBundle(query_str)
    # configure retriever
    retriever = VectorIndexRetriever(
        index=index,
        similarity_top_k=vector_top_k,
    )
    retrieved_nodes = retriever.retrieve(query_bundle)

    if with_reranker:
        # configure reranker
        reranker = RankGPTRerank(
            llm=llm,
            top_n=reranker_top_n,
            verbose=True,
        )
        retrieved_nodes = reranker.postprocess_nodes(
            retrieved_nodes, query_bundle
        )

    return retrieved_nodes
python
new_nodes = get_retrieved_nodes(
    "Which date did Paul Gauguin arrive in Arles ?",
    vector_top_k=10,
    reranker_top_n=3,
    with_reranker=True,
)
python
visualize_retrieved_nodes(new_nodes)