LLM Reranker Demonstration (Great Gatsby)

This tutorial showcases how to do a two-stage pass for retrieval. Use embedding-based retrieval with a high top-k value in order to maximize recall and get a large set of candidate items. Then, use LLM-based retrieval to dynamically select the nodes that are actually relevant to the query.

import nest_asyncio

nest_asyncio.apply()

import logging
import sys

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

Load Data, Build Index

# LLM Predictor (gpt-3.5-turbo) + service context
llm = OpenAI(temperature=0, model="gpt-3.5-turbo")
service_context = ServiceContext.from_defaults(llm=llm, chunk_size=512)

WARNING:llama_index.llm_predictor.base:Unknown max input size for gpt-3.5-turbo, using defaults.
Unknown max input size for gpt-3.5-turbo, using defaults.

# load documents
documents = SimpleDirectoryReader("../../../examples/gatsby/data").load_data()

documents

index = VectorStoreIndex.from_documents(
    documents, service_context=service_context
)

INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens
> [build_index_from_nodes] Total LLM token usage: 0 tokens
INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 49266 tokens
> [build_index_from_nodes] Total embedding token usage: 49266 tokens

Retrieval

from llama_index.retrievers import VectorIndexRetriever
from llama_index.schema import QueryBundle
import pandas as pd
from IPython.display import display, HTML


pd.set_option("display.max_colwidth", -1)


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 = LLMRerank(
            choice_batch_size=5,
            top_n=reranker_top_n,
            service_context=service_context,
        )
        retrieved_nodes = reranker.postprocess_nodes(
            retrieved_nodes, query_bundle
        )

    return retrieved_nodes


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


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))

/var/folders/1r/c3h91d9s49xblwfvz79s78_c0000gn/T/ipykernel_44297/3519340226.py:7: FutureWarning: Passing a negative integer is deprecated in version 1.0 and will not be supported in future version. Instead, use None to not limit the column width.
  pd.set_option('display.max_colwidth', -1)

new_nodes = get_retrieved_nodes(
    "Who was driving the car that hit Myrtle?",
    vector_top_k=3,
    with_reranker=False,
)

INFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens
> [retrieve] Total LLM token usage: 0 tokens
INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 10 tokens
> [retrieve] Total embedding token usage: 10 tokens

visualize_retrieved_nodes(new_nodes)

	Score	Text
0	0.828844	and some garrulous man telling over and over what had happened, until it became less and less real even to him and he could tell it no longer, and Myrtle Wilson’s tragic achievement was forgotten. Now I want to go back a little and tell what happened at the garage after we left there the night before. They had difficulty in locating the sister, Catherine. She must have broken her rule against drinking that night, for when she arrived she was stupid with liquor and unable to understand that the ambulance had already gone to Flushing. When they convinced her of this, she immediately fainted, as if that was the intolerable part of the affair. Someone, kind or curious, took her in his car and drove her in the wake of her sister’s body. Until long after midnight a changing crowd lapped up against the front of the garage, while George Wilson rocked himself back and forth on the couch inside. For a while the door of the office was open, and everyone who came into the garage glanced irresistibly through it. Finally someone said it was a shame, and closed the door. Michaelis and several other men were with him; first, four or five men, later two or three men. Still later Michaelis had to ask the last stranger to wait there fifteen minutes longer, while he went back to his own place and made a pot of coffee. After that, he stayed there alone with Wilson until dawn. About three o’clock the quality of Wilson’s incoherent muttering changed—he grew quieter and began to talk about the yellow car. He announced that he had a way of finding out whom the yellow car belonged to, and then he blurted out that a couple of months ago his wife had come from the city with her face bruised and her nose swollen. But when he heard himself say this, he flinched and began to cry “Oh, my God!” again in his groaning voice. Michaelis made a clumsy
1	0.827754	she rushed out into the dusk, waving her hands and shouting—before he could move from his door the business was over. The “death car” as the newspapers called it, didn’t stop; it came out of the gathering darkness, wavered tragically for a moment, and then disappeared around the next bend. Mavro Michaelis wasn’t even sure of its colour—he told the first policeman that it was light green. The other car, the one going toward New York, came to rest a hundred yards beyond, and its driver hurried back to where Myrtle Wilson, her life violently extinguished, knelt in the road and mingled her thick dark blood with the dust. Michaelis and this man reached her first, but when they had torn open her shirtwaist, still damp with perspiration, they saw that her left breast was swinging loose like a flap, and there was no need to listen for the heart beneath. The mouth was wide open and ripped a little at the corners, as though she had choked a little in giving up the tremendous vitality she had stored so long. ------------------------------------------------------------------------ We saw the three or four automobiles and the crowd when we were still some distance away. “Wreck!” said Tom. “That’s good. Wilson’ll have a little business at last.” He slowed down, but still without any intention of stopping, until, as we came nearer, the hushed, intent faces of the people at the garage door made him automatically put on the brakes. “We’ll take a look,” he said doubtfully, “just a look.” I became aware now of a hollow, wailing sound which issued incessantly from the garage, a sound which as we got out of the coupé and walked toward the door resolved itself into the words “Oh, my God!” uttered over and over in a gasping
2	0.826390	went on, “and left the car in my garage. I don’t think anybody saw us, but of course I can’t be sure.” I disliked him so much by this time that I didn’t find it necessary to tell him he was wrong. “Who was the woman?” he inquired. “Her name was Wilson. Her husband owns the garage. How the devil did it happen?” “Well, I tried to swing the wheel—” He broke off, and suddenly I guessed at the truth. “Was Daisy driving?” “Yes,” he said after a moment, “but of course I’ll say I was. You see, when we left New York she was very nervous and she thought it would steady her to drive—and this woman rushed out at us just as we were passing a car coming the other way. It all happened in a minute, but it seemed to me that she wanted to speak to us, thought we were somebody she knew. Well, first Daisy turned away from the woman toward the other car, and then she lost her nerve and turned back. The second my hand reached the wheel I felt the shock—it must have killed her instantly.” “It ripped her open—” “Don’t tell me, old sport.” He winced. “Anyhow—Daisy stepped on it. I tried to make her stop, but she couldn’t, so I pulled on the emergency brake. Then she fell over into my lap and I drove on. “She’ll be all right tomorrow,” he said presently. “I’m just going to wait here and see if he tries to bother her about that unpleasantness this afternoon. She’s locked herself into her room, and if he tries any brutality she’s going to turn the light out and on again.” “He won’t touch

new_nodes = get_retrieved_nodes(
    "Who was driving the car that hit Myrtle?",
    vector_top_k=10,
    reranker_top_n=3,
    with_reranker=True,
)

visualize_retrieved_nodes(new_nodes)

	Score	Text
0	10.0	went on, “and left the car in my garage. I don’t think anybody saw us, but of course I can’t be sure.” I disliked him so much by this time that I didn’t find it necessary to tell him he was wrong. “Who was the woman?” he inquired. “Her name was Wilson. Her husband owns the garage. How the devil did it happen?” “Well, I tried to swing the wheel—” He broke off, and suddenly I guessed at the truth. “Was Daisy driving?” “Yes,” he said after a moment, “but of course I’ll say I was. You see, when we left New York she was very nervous and she thought it would steady her to drive—and this woman rushed out at us just as we were passing a car coming the other way. It all happened in a minute, but it seemed to me that she wanted to speak to us, thought we were somebody she knew. Well, first Daisy turned away from the woman toward the other car, and then she lost her nerve and turned back. The second my hand reached the wheel I felt the shock—it must have killed her instantly.” “It ripped her open—” “Don’t tell me, old sport.” He winced. “Anyhow—Daisy stepped on it. I tried to make her stop, but she couldn’t, so I pulled on the emergency brake. Then she fell over into my lap and I drove on. “She’ll be all right tomorrow,” he said presently. “I’m just going to wait here and see if he tries to bother her about that unpleasantness this afternoon. She’s locked herself into her room, and if he tries any brutality she’s going to turn the light out and on again.” “He won’t touch

new_nodes = get_retrieved_nodes(
    "What did Gatsby want Daisy to do in front of Tom?",
    vector_top_k=3,
    with_reranker=False,
)

INFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens
> [retrieve] Total LLM token usage: 0 tokens
INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 14 tokens
> [retrieve] Total embedding token usage: 14 tokens

visualize_retrieved_nodes(new_nodes)

****Score****: 0.8647796939111776
****Node text****
: got to make your house into a pigsty in order to have any
friends—in the modern world.”

Angry as I was, as we all were, I was tempted to laugh whenever he
opened his mouth. The transition from libertine to prig was so
complete.

“I’ve got something to tell you, old sport—” began Gatsby. But Daisy
guessed at his intention.

“Please don’t!” she interrupted helplessly. “Please let’s all go
home. Why don’t we all go home?”

“That’s a good idea,” I got up. “Come on, Tom. Nobody wants a drink.”

“I want to know what Mr. Gatsby has to tell me.”

“Your wife doesn’t love you,” said Gatsby. “She’s never loved you.
She loves me.”

“You must be crazy!” exclaimed Tom automatically.

Gatsby sprang to his feet, vivid with excitement.

“She never loved you, do you hear?” he cried. “She only married you
because I was poor and she was tired of waiting for me. It was a
terrible mistake, but in her heart she never loved anyone except me!”

At this point Jordan and I tried to go, but Tom and Gatsby insisted
with competitive firmness that we remain—as though neither of them had
anything to conceal and it would be a privilege to partake vicariously
of their emotions.

“Sit down, Daisy,” Tom’s voice groped unsuccessfully for the paternal
note. “What’s been going on? I want to hear all about it.”

“I told you what’s been going on,” said Gatsby. “Going on for five
years—and you didn’t know.”

Tom turned to Daisy


****Score****: 0.8609230717744326
****Node text****
: to keep your
shoes dry?” There was a husky tenderness in his tone … “Daisy?”

“Please don’t.” Her voice was cold, but the rancour was gone from it.
She looked at Gatsby. “There, Jay,” she said—but her hand as she tried
to light a cigarette was trembling. Suddenly she threw the cigarette
and the burning match on the carpet.

“Oh, you want too much!” she cried to Gatsby. “I love you now—isn’t
that enough? I can’t help what’s past.” She began to sob
helplessly. “I did love him once—but I loved you too.”

Gatsby’s eyes opened and closed.

“You loved me too?” he repeated.

“Even that’s a lie,” said Tom savagely. “She didn’t know you were
alive. Why—there’s things between Daisy and me that you’ll never know,
things that neither of us can ever forget.”

The words seemed to bite physically into Gatsby.

“I want to speak to Daisy alone,” he insisted. “She’s all excited
now—”

“Even alone I can’t say I never loved Tom,” she admitted in a pitiful
voice. “It wouldn’t be true.”

“Of course it wouldn’t,” agreed Tom.

She turned to her husband.

“As if it mattered to you,” she said.

“Of course it matters. I’m going to take better care of you from now
on.”

“You don’t understand,” said Gatsby, with a touch of panic. “You’re
not going to take care of her any more.”

“I’m not?” Tom opened his eyes wide and


****Score****: 0.8555028907426916
****Node text****
: shadowed well with awnings, was dark and cool. Daisy and
Jordan lay upon an enormous couch, like silver idols weighing down
their own white dresses against the singing breeze of the fans.

“We can’t move,” they said together.

Jordan’s fingers, powdered white over their tan, rested for a moment
in mine.

“And Mr. Thomas Buchanan, the athlete?” I inquired.

Simultaneously I heard his voice, gruff, muffled, husky, at the hall
telephone.

Gatsby stood in the centre of the crimson carpet and gazed around with
fascinated eyes. Daisy watched him and laughed, her sweet, exciting
laugh; a tiny gust of powder rose from her bosom into the air.

“The rumour is,” whispered Jordan, “that that’s Tom’s girl on the
telephone.”

We were silent. The voice in the hall rose high with annoyance: “Very
well, then, I won’t sell you the car at all … I’m under no obligations
to you at all … and as for your bothering me about it at lunch time, I
won’t stand that at all!”

“Holding down the receiver,” said Daisy cynically.

“No, he’s not,” I assured her. “It’s a bona-fide deal. I happen to
know about it.”

Tom flung open the door, blocked out its space for a moment with his
thick body, and hurried into the room.

“Mr. Gatsby!” He put out his broad, flat hand with well-concealed
dislike. “I’m glad to see you, sir … Nick …”

“Make us a cold drink,” cried Daisy.

As he left the room again she got up and went over to Gatsby and
pulled his face

new_nodes = get_retrieved_nodes(
    "What did Gatsby want Daisy to do in front of Tom?",
    vector_top_k=10,
    reranker_top_n=3,
    with_reranker=True,
)

INFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens
> [retrieve] Total LLM token usage: 0 tokens
INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 14 tokens
> [retrieve] Total embedding token usage: 14 tokens
Doc: 2, Relevance: 10
No relevant documents found. Please provide a different question.

visualize_retrieved_nodes(new_nodes)

****Score****: 10.0
****Node text****
: to keep your
shoes dry?” There was a husky tenderness in his tone … “Daisy?”

“Please don’t.” Her voice was cold, but the rancour was gone from it.
She looked at Gatsby. “There, Jay,” she said—but her hand as she tried
to light a cigarette was trembling. Suddenly she threw the cigarette
and the burning match on the carpet.

“Oh, you want too much!” she cried to Gatsby. “I love you now—isn’t
that enough? I can’t help what’s past.” She began to sob
helplessly. “I did love him once—but I loved you too.”

Gatsby’s eyes opened and closed.

“You loved me too?” he repeated.

“Even that’s a lie,” said Tom savagely. “She didn’t know you were
alive. Why—there’s things between Daisy and me that you’ll never know,
things that neither of us can ever forget.”

The words seemed to bite physically into Gatsby.

“I want to speak to Daisy alone,” he insisted. “She’s all excited
now—”

“Even alone I can’t say I never loved Tom,” she admitted in a pitiful
voice. “It wouldn’t be true.”

“Of course it wouldn’t,” agreed Tom.

She turned to her husband.

“As if it mattered to you,” she said.

“Of course it matters. I’m going to take better care of you from now
on.”

“You don’t understand,” said Gatsby, with a touch of panic. “You’re
not going to take care of her any more.”

“I’m not?” Tom opened his eyes wide and

Query Engine

query_engine = index.as_query_engine(
    similarity_top_k=10,
    node_postprocessors=[reranker],
    response_mode="tree_summarize",
)
response = query_engine.query(
    "What did the author do during his time at Y Combinator?",
)

query_engine = index.as_query_engine(
    similarity_top_k=3, response_mode="tree_summarize"
)
response = query_engine.query(
    "What did the author do during his time at Y Combinator?",
)

retrieval =