{ "cells": [ { "attachments": {}, "cell_type": "markdown", "id": "f1b23ffc-22d8-4414-b2e4-66d45a03523d", "metadata": {}, "source": [ "# LLM Reranker Demonstration (Great Gatsby)\n", "\n", "This tutorial showcases how to do a two-stage pass for retrieval. Use embedding-based retrieval with a high top-k value\n", "in order to maximize recall and get a large set of candidate items. Then, use LLM-based retrieval\n", "to dynamically select the nodes that are actually relevant to the query." ] }, { "cell_type": "code", "execution_count": null, "id": "1bd392e1", "metadata": {}, "outputs": [], "source": [ "%pip install llama-index-llms-openai" ] }, { "cell_type": "code", "execution_count": null, "id": "ab68a3ee-f4c1-4830-856a-9c37f7a42364", "metadata": {}, "outputs": [], "source": [ "import nest_asyncio\n", "\n", "nest_asyncio.apply()" ] }, { "cell_type": "code", "execution_count": null, "id": "8c16404b-db81-4ee6-9c0f-cc4cdf3cb147", "metadata": {}, "outputs": [], "source": [ "import logging\n", "import sys\n", "\n", "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n", "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n", "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader\n", "from llama_index.core.postprocessor import LLMRerank\n", "from llama_index.llms.openai import OpenAI\n", "from IPython.display import Markdown, display" ] }, { "attachments": {}, "cell_type": "markdown", "id": "8a5b0ac6-ec9c-40e1-9120-d20e33c37f80", "metadata": {}, "source": [ "## Load Data, Build Index" ] }, { "cell_type": "code", "execution_count": null, "id": "b785b75a-cf4d-4799-8fd2-970e5c39727c", "metadata": {}, "outputs": [], "source": [ "from llama_index.core import Settings\n", "\n", "# LLM (gpt-3.5-turbo)\n", "Settings.llm = OpenAI(temperature=0, model=\"gpt-3.5-turbo\")\n", "Settings.chunk_size = 512" ] }, { "cell_type": "code", "execution_count": null, "id": "9b689e92-fc55-48b5-859f-2a2c1394c872", "metadata": {}, "outputs": [], "source": [ "# load documents\n", "documents = SimpleDirectoryReader(\"../../../examples/gatsby/data\").load_data()" ] }, { "cell_type": "code", "execution_count": null, "id": "1a01180d-67c8-4664-a4ee-ceab1367cb78", "metadata": {}, "outputs": [], "source": [ "documents" ] }, { "cell_type": "code", "execution_count": null, "id": "d8555b6d-5289-4803-bced-6b2f85dea3da", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens\n", "> [build_index_from_nodes] Total LLM token usage: 0 tokens\n", "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 49266 tokens\n", "> [build_index_from_nodes] Total embedding token usage: 49266 tokens\n" ] } ], "source": [ "index = VectorStoreIndex.from_documents(\n", " documents,\n", ")" ] }, { "attachments": {}, "cell_type": "markdown", "id": "805847d9-a2c1-4930-98a9-98126e730000", "metadata": {}, "source": [ "## Retrieval" ] }, { "cell_type": "code", "execution_count": null, "id": "ace86835-832e-4964-a025-428891aa2c8b", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/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.\n", " pd.set_option('display.max_colwidth', -1)\n" ] } ], "source": [ "from llama_index.core.retrievers import VectorIndexRetriever\n", "from llama_index.core import QueryBundle\n", "import pandas as pd\n", "from IPython.display import display, HTML\n", "\n", "\n", "pd.set_option(\"display.max_colwidth\", -1)\n", "\n", "\n", "def get_retrieved_nodes(\n", " query_str, vector_top_k=10, reranker_top_n=3, with_reranker=False\n", "):\n", " query_bundle = QueryBundle(query_str)\n", " # configure retriever\n", " retriever = VectorIndexRetriever(\n", " index=index,\n", " similarity_top_k=vector_top_k,\n", " )\n", " retrieved_nodes = retriever.retrieve(query_bundle)\n", "\n", " if with_reranker:\n", " # configure reranker\n", " reranker = LLMRerank(\n", " choice_batch_size=5,\n", " top_n=reranker_top_n,\n", " )\n", " retrieved_nodes = reranker.postprocess_nodes(\n", " retrieved_nodes, query_bundle\n", " )\n", "\n", " return retrieved_nodes\n", "\n", "\n", "def pretty_print(df):\n", " return display(HTML(df.to_html().replace(\"\\\\n\", \"
\")))\n", "\n", "\n", "def visualize_retrieved_nodes(nodes) -> None:\n", " result_dicts = []\n", " for node in nodes:\n", " result_dict = {\"Score\": node.score, \"Text\": node.node.get_text()}\n", " result_dicts.append(result_dict)\n", "\n", " pretty_print(pd.DataFrame(result_dicts))" ] }, { "cell_type": "code", "execution_count": null, "id": "20993200-b725-403e-8a79-e2571dab2ebc", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "INFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens\n", "> [retrieve] Total LLM token usage: 0 tokens\n", "INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 10 tokens\n", "> [retrieve] Total embedding token usage: 10 tokens\n" ] } ], "source": [ "new_nodes = get_retrieved_nodes(\n", " \"Who was driving the car that hit Myrtle?\",\n", " vector_top_k=3,\n", " with_reranker=False,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "9dd3bc25-6ef1-46b7-869e-247c83af9d4d", "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
ScoreText
00.828844and 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
10.827754she 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
20.826390went 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
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "visualize_retrieved_nodes(new_nodes)" ] }, { "cell_type": "code", "execution_count": null, "id": "0ad7d62b-dba9-45d5-896f-4baa9a40edba", "metadata": {}, "outputs": [], "source": [ "new_nodes = get_retrieved_nodes(\n", " \"Who was driving the car that hit Myrtle?\",\n", " vector_top_k=10,\n", " reranker_top_n=3,\n", " with_reranker=True,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "ae22d297-dd3b-4a71-a890-faceafab4e1a", "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
ScoreText
010.0went 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
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "visualize_retrieved_nodes(new_nodes)" ] }, { "cell_type": "code", "execution_count": null, "id": "4e540fb6-189f-4a54-98d6-852e561f39ee", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "INFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens\n", "> [retrieve] Total LLM token usage: 0 tokens\n", "INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 14 tokens\n", "> [retrieve] Total embedding token usage: 14 tokens\n" ] } ], "source": [ "new_nodes = get_retrieved_nodes(\n", " \"What did Gatsby want Daisy to do in front of Tom?\",\n", " vector_top_k=3,\n", " with_reranker=False,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "3b82380c-e29f-4e9e-a46d-038ca1469904", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "\n", "****Score****: 0.8647796939111776\n", "****Node text****\n", ": got to make your house into a pigsty in order to have any\n", "friends—in the modern world.”\n", "\n", "Angry as I was, as we all were, I was tempted to laugh whenever he\n", "opened his mouth. The transition from libertine to prig was so\n", "complete.\n", "\n", "“I’ve got something to tell you, old sport—” began Gatsby. But Daisy\n", "guessed at his intention.\n", "\n", "“Please don’t!” she interrupted helplessly. “Please let’s all go\n", "home. Why don’t we all go home?”\n", "\n", "“That’s a good idea,” I got up. “Come on, Tom. Nobody wants a drink.”\n", "\n", "“I want to know what Mr. Gatsby has to tell me.”\n", "\n", "“Your wife doesn’t love you,” said Gatsby. “She’s never loved you.\n", "She loves me.”\n", "\n", "“You must be crazy!” exclaimed Tom automatically.\n", "\n", "Gatsby sprang to his feet, vivid with excitement.\n", "\n", "“She never loved you, do you hear?” he cried. “She only married you\n", "because I was poor and she was tired of waiting for me. It was a\n", "terrible mistake, but in her heart she never loved anyone except me!”\n", "\n", "At this point Jordan and I tried to go, but Tom and Gatsby insisted\n", "with competitive firmness that we remain—as though neither of them had\n", "anything to conceal and it would be a privilege to partake vicariously\n", "of their emotions.\n", "\n", "“Sit down, Daisy,” Tom’s voice groped unsuccessfully for the paternal\n", "note. “What’s been going on? I want to hear all about it.”\n", "\n", "“I told you what’s been going on,” said Gatsby. “Going on for five\n", "years—and you didn’t know.”\n", "\n", "Tom turned to Daisy\n", "\n", "\n", "****Score****: 0.8609230717744326\n", "****Node text****\n", ": to keep your\n", "shoes dry?” There was a husky tenderness in his tone … “Daisy?”\n", "\n", "“Please don’t.” Her voice was cold, but the rancour was gone from it.\n", "She looked at Gatsby. “There, Jay,” she said—but her hand as she tried\n", "to light a cigarette was trembling. Suddenly she threw the cigarette\n", "and the burning match on the carpet.\n", "\n", "“Oh, you want too much!” she cried to Gatsby. “I love you now—isn’t\n", "that enough? I can’t help what’s past.” She began to sob\n", "helplessly. “I did love him once—but I loved you too.”\n", "\n", "Gatsby’s eyes opened and closed.\n", "\n", "“You loved me too?” he repeated.\n", "\n", "“Even that’s a lie,” said Tom savagely. “She didn’t know you were\n", "alive. Why—there’s things between Daisy and me that you’ll never know,\n", "things that neither of us can ever forget.”\n", "\n", "The words seemed to bite physically into Gatsby.\n", "\n", "“I want to speak to Daisy alone,” he insisted. “She’s all excited\n", "now—”\n", "\n", "“Even alone I can’t say I never loved Tom,” she admitted in a pitiful\n", "voice. “It wouldn’t be true.”\n", "\n", "“Of course it wouldn’t,” agreed Tom.\n", "\n", "She turned to her husband.\n", "\n", "“As if it mattered to you,” she said.\n", "\n", "“Of course it matters. I’m going to take better care of you from now\n", "on.”\n", "\n", "“You don’t understand,” said Gatsby, with a touch of panic. “You’re\n", "not going to take care of her any more.”\n", "\n", "“I’m not?” Tom opened his eyes wide and\n", "\n", "\n", "****Score****: 0.8555028907426916\n", "****Node text****\n", ": shadowed well with awnings, was dark and cool. Daisy and\n", "Jordan lay upon an enormous couch, like silver idols weighing down\n", "their own white dresses against the singing breeze of the fans.\n", "\n", "“We can’t move,” they said together.\n", "\n", "Jordan’s fingers, powdered white over their tan, rested for a moment\n", "in mine.\n", "\n", "“And Mr. Thomas Buchanan, the athlete?” I inquired.\n", "\n", "Simultaneously I heard his voice, gruff, muffled, husky, at the hall\n", "telephone.\n", "\n", "Gatsby stood in the centre of the crimson carpet and gazed around with\n", "fascinated eyes. Daisy watched him and laughed, her sweet, exciting\n", "laugh; a tiny gust of powder rose from her bosom into the air.\n", "\n", "“The rumour is,” whispered Jordan, “that that’s Tom’s girl on the\n", "telephone.”\n", "\n", "We were silent. The voice in the hall rose high with annoyance: “Very\n", "well, then, I won’t sell you the car at all … I’m under no obligations\n", "to you at all … and as for your bothering me about it at lunch time, I\n", "won’t stand that at all!”\n", "\n", "“Holding down the receiver,” said Daisy cynically.\n", "\n", "“No, he’s not,” I assured her. “It’s a bona-fide deal. I happen to\n", "know about it.”\n", "\n", "Tom flung open the door, blocked out its space for a moment with his\n", "thick body, and hurried into the room.\n", "\n", "“Mr. Gatsby!” He put out his broad, flat hand with well-concealed\n", "dislike. “I’m glad to see you, sir … Nick …”\n", "\n", "“Make us a cold drink,” cried Daisy.\n", "\n", "As he left the room again she got up and went over to Gatsby and\n", "pulled his face\n" ] } ], "source": [ "visualize_retrieved_nodes(new_nodes)" ] }, { "cell_type": "code", "execution_count": null, "id": "aec75734-52e4-4b7e-9b6d-d0b7306ae5d8", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "INFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens\n", "> [retrieve] Total LLM token usage: 0 tokens\n", "INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 14 tokens\n", "> [retrieve] Total embedding token usage: 14 tokens\n", "Doc: 2, Relevance: 10\n", "No relevant documents found. Please provide a different question.\n" ] } ], "source": [ "new_nodes = get_retrieved_nodes(\n", " \"What did Gatsby want Daisy to do in front of Tom?\",\n", " vector_top_k=10,\n", " reranker_top_n=3,\n", " with_reranker=True,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "89e5408e-bf4b-483f-9144-fc5ac4fb48e8", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "\n", "****Score****: 10.0\n", "****Node text****\n", ": to keep your\n", "shoes dry?” There was a husky tenderness in his tone … “Daisy?”\n", "\n", "“Please don’t.” Her voice was cold, but the rancour was gone from it.\n", "She looked at Gatsby. “There, Jay,” she said—but her hand as she tried\n", "to light a cigarette was trembling. Suddenly she threw the cigarette\n", "and the burning match on the carpet.\n", "\n", "“Oh, you want too much!” she cried to Gatsby. “I love you now—isn’t\n", "that enough? I can’t help what’s past.” She began to sob\n", "helplessly. “I did love him once—but I loved you too.”\n", "\n", "Gatsby’s eyes opened and closed.\n", "\n", "“You loved me too?” he repeated.\n", "\n", "“Even that’s a lie,” said Tom savagely. “She didn’t know you were\n", "alive. Why—there’s things between Daisy and me that you’ll never know,\n", "things that neither of us can ever forget.”\n", "\n", "The words seemed to bite physically into Gatsby.\n", "\n", "“I want to speak to Daisy alone,” he insisted. “She’s all excited\n", "now—”\n", "\n", "“Even alone I can’t say I never loved Tom,” she admitted in a pitiful\n", "voice. “It wouldn’t be true.”\n", "\n", "“Of course it wouldn’t,” agreed Tom.\n", "\n", "She turned to her husband.\n", "\n", "“As if it mattered to you,” she said.\n", "\n", "“Of course it matters. I’m going to take better care of you from now\n", "on.”\n", "\n", "“You don’t understand,” said Gatsby, with a touch of panic. “You’re\n", "not going to take care of her any more.”\n", "\n", "“I’m not?” Tom opened his eyes wide and\n" ] } ], "source": [ "visualize_retrieved_nodes(new_nodes)" ] }, { "attachments": {}, "cell_type": "markdown", "id": "23dcdac6-f4dd-469e-9f47-d030f27bacda", "metadata": {}, "source": [ "## Query Engine" ] }, { "cell_type": "code", "execution_count": null, "id": "94f6531d-953a-41c6-9bcc-9c293550f8c4", "metadata": {}, "outputs": [], "source": [ "query_engine = index.as_query_engine(\n", " similarity_top_k=10,\n", " node_postprocessors=[\n", " LLMRerank(\n", " choice_batch_size=5,\n", " top_n=2,\n", " )\n", " ],\n", " response_mode=\"tree_summarize\",\n", ")\n", "response = query_engine.query(\n", " \"What did the author do during his time at Y Combinator?\",\n", ")" ] } ], "metadata": { "kernelspec": { "display_name": "llama_index_v2", "language": "python", "name": "llama_index_v2" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3" } }, "nbformat": 4, "nbformat_minor": 5 }