From ee3fc01977499788226388fbceffe4142ddb3fe7 Mon Sep 17 00:00:00 2001 From: Jerry Liu Date: Mon, 7 Aug 2023 23:52:17 -0700 Subject: [PATCH] cr --- .../query_modules/retriever/modules.md | 1 + .../retrievers/ensemble_retrieval.ipynb | 1262 +++++++++++++++++ llama_index/response/notebook_utils.py | 15 +- 3 files changed, 1276 insertions(+), 2 deletions(-) create mode 100644 docs/examples/retrievers/ensemble_retrieval.ipynb diff --git a/docs/core_modules/query_modules/retriever/modules.md b/docs/core_modules/query_modules/retriever/modules.md index 3ad0528a125bd..5300a1e417a6e 100644 --- a/docs/core_modules/query_modules/retriever/modules.md +++ b/docs/core_modules/query_modules/retriever/modules.md @@ -50,4 +50,5 @@ maxdepth: 1 --- /examples/query_engine/pdf_tables/recursive_retriever.ipynb /examples/retrievers/router_retriever.ipynb +/examples/retrievers/ensemble_retrieval.ipynb ``` diff --git a/docs/examples/retrievers/ensemble_retrieval.ipynb b/docs/examples/retrievers/ensemble_retrieval.ipynb new file mode 100644 index 0000000000000..615df653155ab --- /dev/null +++ b/docs/examples/retrievers/ensemble_retrieval.ipynb @@ -0,0 +1,1262 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "5bf1de44-4047-46cf-a04c-dbf910d9e179", + "metadata": {}, + "source": [ + "# Ensemble Retrieval Guide\n", + "\n", + "Oftentimes when building a RAG applications there are many retreival parameters/strategies to decide from (from chunk size to vector vs. keyword vs. hybrid search, for instance).\n", + "\n", + "Thought: what if we could try a bunch of strategies at once, and have any AI/reranker/LLM prune the results?\n", + "\n", + "This achieves two purposes:\n", + "- Better (albeit more costly) retrieved results by pooling results from multiple strategies, assuming the reranker is good\n", + "- A way to benchmark different retrieval strategies against each other (w.r.t reranker)\n", + "\n", + "This guide showcases this over the Great Gatsby. We do ensemble retrieval over different chunk sizes and also different indices." + ] + }, + { + "cell_type": "markdown", + "id": "6e73fead-ec2c-4346-bd08-e183c13c7e29", + "metadata": {}, + "source": [ + "## Setup" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "a2d59778-4cda-47b5-8cd0-b80fee91d1e4", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# NOTE: This is ONLY necessary in jupyter notebook.\n", + "# Details: Jupyter runs an event-loop behind the scenes.\n", + "# This results in nested event-loops when we start an event-loop to make async queries.\n", + "# This is normally not allowed, we use nest_asyncio to allow it for convenience.\n", + "import nest_asyncio\n", + "\n", + "nest_asyncio.apply()" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "c628448c-573c-4eeb-a7e1-707fe8cc575c", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Note: NumExpr detected 12 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n", + "NumExpr defaulting to 8 threads.\n" + ] + } + ], + "source": [ + "import logging\n", + "import sys\n", + "\n", + "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n", + "logging.getLogger().handlers = []\n", + "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n", + "\n", + "from llama_index import (\n", + " VectorStoreIndex,\n", + " ListIndex,\n", + " SimpleDirectoryReader,\n", + " ServiceContext,\n", + " StorageContext,\n", + " SimpleKeywordTableIndex,\n", + ")\n", + "from llama_index.response.notebook_utils import display_response\n", + "from llama_index.llms import OpenAI" + ] + }, + { + "cell_type": "markdown", + "id": "787174ed-10ce-47d7-82fd-9ca7f891eea7", + "metadata": {}, + "source": [ + "## Load Data\n", + "\n", + "We first show how to convert a Document into a set of Nodes, and insert into a DocumentStore." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "dd62366b-8a24-40a7-8c47-5859851149fe", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# try loading great gatsby\n", + "\n", + "documents = SimpleDirectoryReader(input_files=[\"../../../examples/gatsby/gatsby_full.txt\"]).load_data()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7081194a-ede7-478e-bff2-23e89e23ef16", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Chunk Size: 128\n", + "Chunk Size: 256\n", + "Chunk Size: 512\n" + ] + } + ], + "source": [ + "# initialize service context (set chunk size)\n", + "llm = OpenAI(model=\"gpt-4\")\n", + "chunk_sizes = [128, 256, 512, 1024]\n", + "service_contexts = []\n", + "nodes_list = []\n", + "vector_indices = []\n", + "query_engines = []\n", + "for chunk_size in chunk_sizes:\n", + " print(f'Chunk Size: {chunk_size}')\n", + " service_context = ServiceContext.from_defaults(chunk_size=chunk_size, llm=llm)\n", + " service_contexts.append(service_context)\n", + " nodes = service_context.node_parser.get_nodes_from_documents(documents)\n", + " \n", + " # add chunk size to nodes to track later\n", + " for node in nodes:\n", + " node.metadata[\"chunk_size\"] = chunk_size\n", + " node.excluded_embed_metadata_keys = [\"chunk_size\"]\n", + " node.excluded_llm_metadata_keys = [\"chunk_size\"]\n", + " \n", + " nodes_list.append(nodes)\n", + " \n", + " # build vector index\n", + " vector_index = VectorStoreIndex(nodes)\n", + " vector_indices.append(vector_index)\n", + " \n", + " # query engines\n", + " query_engines.append(vector_index.as_query_engine())" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "fbca69b4-d8d5-4dcb-af33-f9ed4a91ec05", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# try ensemble retrieval\n", + "\n", + "from llama_index.tools import RetrieverTool\n", + "\n", + "retriever_tools = []\n", + "for chunk_size, vector_index in zip(chunk_sizes, vector_indices):\n", + " retriever_tool = RetrieverTool.from_defaults(\n", + " retriever=vector_index.as_retriever(),\n", + " description=f\"Retrieves relevant context from the Great Gatsby (chunk size {chunk_size})\",\n", + " )\n", + " retriever_tools.append(retriever_tool)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 51, + "id": "5c9eaa6f-8f11-4380-b3c6-79092f17def3", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from llama_index.selectors.pydantic_selectors import (\n", + " PydanticMultiSelector\n", + ")\n", + "from llama_index.retrievers import RouterRetriever\n", + "\n", + "\n", + "retriever = RouterRetriever(\n", + " selector=PydanticMultiSelector.from_defaults(llm=llm, max_outputs=4),\n", + " retriever_tools=retriever_tools\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "7c72c61c-d4f7-4159-bb80-1989468ab61c", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "ename": "NameError", + "evalue": "name 'retriever' is not defined", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[4], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m nodes \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m retriever\u001b[38;5;241m.\u001b[39maretrieve(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDescribe and summarize the interactions between Gatsby and Daisy\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n", + "\u001b[0;31mNameError\u001b[0m: name 'retriever' is not defined" + ] + } + ], + "source": [ + "nodes = await retriever.aretrieve(\"Describe and summarize the interactions between Gatsby and Daisy\")" + ] + }, + { + "cell_type": "code", + "execution_count": 55, + "id": "590ed8bc-83ad-4851-9ec6-bfbbdf3ff38d", + "metadata": { + "scrolled": true, + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "128\n", + "the beach that morning. Finally we came to Gatsby’s own\n", + "apartment, a bedroom and a bath, and an Adam’s study, where we sat\n", + "down and drank a glass of some Chartreuse he took from a cupboard in\n", + "the wall.\n", + "\n", + "He hadn’t once ceased looking at Daisy, and I think he revalued\n", + "everything in his house according to the measure of response it drew\n", + "from her well-loved eyes. Sometimes too, he stared around at his\n", + "possessions in a dazed\n", + "128\n", + "turn out as he had\n", + "imagined. He had intended, probably, to take what he could and go—but\n", + "now he found that he had committed himself to the following of a\n", + "grail. He knew that Daisy was extraordinary, but he didn’t realize\n", + "just how extraordinary a “nice” girl could be. She vanished into her\n", + "rich house, into her rich, full life, leaving Gatsby—nothing. He felt\n", + "married to her, that was all.\n", + "\n", + "When they met again, two days later, it\n", + "256\n", + "the\n", + "direction. In this heat every extra gesture was an affront to the\n", + "common store of life.\n", + "\n", + "The room, 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\n", + "256\n", + "In the meantime, In between time—”\n", + "\n", + "As I went over to say goodbye I saw that the expression of\n", + "bewilderment had come back into Gatsby’s face, as though a faint doubt\n", + "had occurred to him as to the quality of his present happiness. Almost\n", + "five years! There must have been moments even that afternoon when\n", + "Daisy tumbled short of his dreams—not through her own fault, but\n", + "because of the colossal vitality of his illusion. It had gone beyond\n", + "her, beyond everything. He had thrown himself into it with a creative\n", + "passion, adding to it all the time, decking it out with every bright\n", + "feather that drifted his way. No amount of fire or freshness can\n", + "challenge what a man can store up in his ghostly heart.\n", + "\n", + "As I watched him he adjusted himself a little, visibly. His hand took\n", + "hold of hers, and as she said something low in his ear he turned\n", + "toward her with a rush of emotion. I think that voice held him most,\n", + "with its fluctuating,\n", + "512\n", + "go downstairs,” interrupted Gatsby. He flipped a switch. The\n", + "grey windows disappeared as the house glowed full of light.\n", + "\n", + "In the music-room Gatsby turned on a solitary lamp beside the piano.\n", + "He lit Daisy’s cigarette from a trembling match, and sat down with her\n", + "on a couch far across the room, where there was no light save what the\n", + "gleaming floor bounced in from the hall.\n", + "\n", + "When Klipspringer had played “The Love Nest” he turned around on the\n", + "bench and searched unhappily for Gatsby in the gloom.\n", + "\n", + "“I’m all out of practice, you see. I told you I couldn’t play. I’m all\n", + "out of prac—”\n", + "\n", + "“Don’t talk so much, old sport,” commanded Gatsby. “Play!”\n", + "\n", + " “In the morning, In the evening, Ain’t we got fun—”\n", + "\n", + "Outside the wind was loud and there was a faint flow of thunder along\n", + "the Sound. All the lights were going on in West Egg now; the electric\n", + "trains, men-carrying, were plunging home through the rain from New\n", + "York. It was the hour of a profound human change, and excitement was\n", + "generating on the air.\n", + "\n", + " “One thing’s sure and nothing’s surer The rich get richer and the\n", + " poor get—children. In the meantime, In between time—”\n", + "\n", + "As I went over to say goodbye I saw that the expression of\n", + "bewilderment had come back into Gatsby’s face, as though a faint doubt\n", + "had occurred to him as to the quality of his present happiness. Almost\n", + "five years! There must have been moments even that afternoon when\n", + "Daisy tumbled short of his dreams—not through her own fault, but\n", + "because of the colossal vitality of his illusion. It had gone beyond\n", + "her, beyond everything. He had thrown himself into it with a creative\n", + "passion, adding to it all the time, decking it out\n", + "512\n", + "world complete\n", + "in itself, with its own standards and its own great figures, second to\n", + "nothing because it had no consciousness of being so, and now I was\n", + "looking at it again, through Daisy’s eyes. It is invariably saddening\n", + "to look through new eyes at things upon which you have expended your\n", + "own powers of adjustment.\n", + "\n", + "They arrived at twilight, and, as we strolled out among the sparkling\n", + "hundreds, Daisy’s voice was playing murmurous tricks in her throat.\n", + "\n", + "“These things excite me so,” she whispered. “If you want to kiss me\n", + "any time during the evening, Nick, just let me know and I’ll be glad\n", + "to arrange it for you. Just mention my name. Or present a green card.\n", + "I’m giving out green—”\n", + "\n", + "“Look around,” suggested Gatsby.\n", + "\n", + "“I’m looking around. I’m having a marvellous—”\n", + "\n", + "“You must see the faces of many people you’ve heard about.”\n", + "\n", + "Tom’s arrogant eyes roamed the crowd.\n", + "\n", + "“We don’t go around very much,” he said; “in fact, I was just thinking\n", + "I don’t know a soul here.”\n", + "\n", + "“Perhaps you know that lady.” Gatsby indicated a gorgeous, scarcely\n", + "human orchid of a woman who sat in state under a white-plum tree. Tom\n", + "and Daisy stared, with that peculiarly unreal feeling that accompanies\n", + "the recognition of a hitherto ghostly celebrity of the movies.\n", + "\n", + "“She’s lovely,” said Daisy.\n", + "\n", + "“The man bending over her is her director.”\n", + "\n", + "He took them ceremoniously from group to group:\n", + "\n", + "“Mrs. Buchanan … and Mr. Buchanan—” After an instant’s hesitation he\n", + "added: “the polo player.”\n", + "\n", + "“Oh no,” objected Tom quickly,\n", + "1024\n", + "The\n", + "grey windows disappeared as the house glowed full of light.\n", + "\n", + "In the music-room Gatsby turned on a solitary lamp beside the piano.\n", + "He lit Daisy’s cigarette from a trembling match, and sat down with her\n", + "on a couch far across the room, where there was no light save what the\n", + "gleaming floor bounced in from the hall.\n", + "\n", + "When Klipspringer had played “The Love Nest” he turned around on the\n", + "bench and searched unhappily for Gatsby in the gloom.\n", + "\n", + "“I’m all out of practice, you see. I told you I couldn’t play. I’m all\n", + "out of prac—”\n", + "\n", + "“Don’t talk so much, old sport,” commanded Gatsby. “Play!”\n", + "\n", + " “In the morning, In the evening, Ain’t we got fun—”\n", + "\n", + "Outside the wind was loud and there was a faint flow of thunder along\n", + "the Sound. All the lights were going on in West Egg now; the electric\n", + "trains, men-carrying, were plunging home through the rain from New\n", + "York. It was the hour of a profound human change, and excitement was\n", + "generating on the air.\n", + "\n", + " “One thing’s sure and nothing’s surer The rich get richer and the\n", + " poor get—children. In the meantime, In between time—”\n", + "\n", + "As I went over to say goodbye I saw that the expression of\n", + "bewilderment had come back into Gatsby’s face, as though a faint doubt\n", + "had occurred to him as to the quality of his present happiness. Almost\n", + "five years! There must have been moments even that afternoon when\n", + "Daisy tumbled short of his dreams—not through her own fault, but\n", + "because of the colossal vitality of his illusion. It had gone beyond\n", + "her, beyond everything. He had thrown himself into it with a creative\n", + "passion, adding to it all the time, decking it out with every bright\n", + "feather that drifted his way. No amount of fire or freshness can\n", + "challenge what a man can store up in his ghostly heart.\n", + "\n", + "As I watched him he adjusted himself a little, visibly. His hand took\n", + "hold of hers, and as she said something low in his ear he turned\n", + "toward her with a rush of emotion. I think that voice held him most,\n", + "with its fluctuating, feverish warmth, because it couldn’t be\n", + "over-dreamed—that voice was a deathless song.\n", + "\n", + "They had forgotten me, but Daisy glanced up and held out her hand;\n", + "Gatsby didn’t know me now at all. I looked once more at them and they\n", + "looked back at me, remotely, possessed by intense life. Then I went\n", + "out of the room and down the marble steps into the rain, leaving them\n", + "there together.\n", + "\n", + "\n", + " VI\n", + "\n", + "About this time an ambitious young reporter from New York arrived one\n", + "morning at Gatsby’s door and asked him if he had anything to say.\n", + "\n", + "“Anything to say about what?” inquired Gatsby politely.\n", + "\n", + "“Why—any statement to give out.”\n", + "\n", + "It transpired after a confused five minutes that the man had heard\n", + "Gatsby’s name around his office in a connection which he either\n", + "wouldn’t reveal or didn’t fully understand. This was his day off and\n", + "with laudable initiative he had hurried out “to see.”\n", + "\n", + "It was a random shot, and yet the reporter’s instinct was right.\n", + "Gatsby’s notoriety, spread about by the hundreds who had accepted his\n", + "hospitality and so become authorities upon his past, had increased all\n", + "summer until he fell just short of being news. Contemporary legends\n", + "such as the “underground pipeline to Canada” attached themselves to\n", + "him, and there was one persistent story that he didn’t live in a house\n", + "at all, but in a boat that looked like a house and was\n", + "1024\n", + "Daisy insistently. Gatsby’s eyes\n", + "floated toward her. “Ah,” she cried, “you look so cool.”\n", + "\n", + "Their eyes met, and they stared together at each other, alone in\n", + "space. With an effort she glanced down at the table.\n", + "\n", + "“You always look so cool,” she repeated.\n", + "\n", + "She had told him that she loved him, and Tom Buchanan saw. He was\n", + "astounded. His mouth opened a little, and he looked at Gatsby, and\n", + "then back at Daisy as if he had just recognized her as someone he knew\n", + "a long time ago.\n", + "\n", + "“You resemble the advertisement of the man,” she went on innocently.\n", + "“You know the advertisement of the man—”\n", + "\n", + "“All right,” broke in Tom quickly, “I’m perfectly willing to go to\n", + "town. Come on—we’re all going to town.”\n", + "\n", + "He got up, his eyes still flashing between Gatsby and his wife. No one\n", + "moved.\n", + "\n", + "“Come on!” His temper cracked a little. “What’s the matter, anyhow?\n", + "If we’re going to town, let’s start.”\n", + "\n", + "His hand, trembling with his effort at self-control, bore to his lips\n", + "the last of his glass of ale. Daisy’s voice got us to our feet and out\n", + "on to the blazing gravel drive.\n", + "\n", + "“Are we just going to go?” she objected. “Like this? Aren’t we going\n", + "to let anyone smoke a cigarette first?”\n", + "\n", + "“Everybody smoked all through lunch.”\n", + "\n", + "“Oh, let’s have fun,” she begged him. “It’s too hot to fuss.”\n", + "\n", + "He didn’t answer.\n", + "\n", + "“Have it your own way,” she said. “Come on, Jordan.”\n", + "\n", + "They went upstairs to get ready while we three men stood there\n", + "shuffling the hot pebbles with our feet. A silver curve of the moon\n", + "hovered already in the western sky. Gatsby started to speak, changed\n", + "his mind, but not before Tom wheeled and faced him expectantly.\n", + "\n", + "“Have you got your stables here?” asked Gatsby with an effort.\n", + "\n", + "“About a quarter of a mile down the road.”\n", + "\n", + "“Oh.”\n", + "\n", + "A pause.\n", + "\n", + "“I don’t see the idea of going to town,” broke out Tom savagely.\n", + "“Women get these notions in their heads—”\n", + "\n", + "“Shall we take anything to drink?” called Daisy from an upper window.\n", + "\n", + "“I’ll get some whisky,” answered Tom. He went inside.\n", + "\n", + "Gatsby turned to me rigidly:\n", + "\n", + "“I can’t say anything in his house, old sport.”\n", + "\n", + "“She’s got an indiscreet voice,” I remarked. “It’s full of—” I\n", + "hesitated.\n", + "\n", + "“Her voice is full of money,” he said suddenly.\n", + "\n", + "That was it. I’d never understood before. It was full of money—that\n", + "was the inexhaustible charm that rose and fell in it, the jingle of\n", + "it, the cymbals’ song of it … High in a white palace the king’s\n", + "daughter, the golden girl …\n", + "\n", + "Tom came out of the house wrapping a quart bottle in a towel, followed\n", + "by Daisy and Jordan wearing small tight hats of metallic cloth and\n", + "carrying light capes over their arms.\n", + "\n", + "“Shall we all go in my car?” suggested Gatsby. He felt the hot, green\n", + "leather of the seat. “I ought to have left it in the shade.”\n", + "\n", + "“Is it standard shift?” demanded Tom.\n", + "\n", + "“Yes.”\n", + "\n", + "“Well, you take my coupé and let me drive your car to town.”\n", + "\n", + "The suggestion was distasteful to Gatsby.\n", + "\n", + "“I\n" + ] + } + ], + "source": [ + "for node in nodes:\n", + " print(node.node.metadata[\"chunk_size\"])\n", + " print(node.node.get_text())" + ] + }, + { + "cell_type": "code", + "execution_count": 86, + "id": "1f26c527-17d2-4d4e-a6ee-8ea878ef8742", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# define reranker\n", + "from llama_index.indices.postprocessor import LLMRerank, SentenceTransformerRerank, CohereRerank\n", + "# reranker = LLMRerank()\n", + "# reranker = SentenceTransformerRerank(top_n=10)\n", + "reranker = CohereRerank(top_n=10)" + ] + }, + { + "cell_type": "code", + "execution_count": 87, + "id": "828589ef-d062-40dc-8a4b-245190769445", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# define RetrieverQueryEngine\n", + "from llama_index.query_engine import RetrieverQueryEngine\n", + "\n", + "query_engine = RetrieverQueryEngine(\n", + " retriever,\n", + " node_postprocessors=[reranker]\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 83, + "id": "53e3c341-e66d-4950-88d5-6411699d064b", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Selecting retriever 0: This choice might provide a detailed context of the interactions between Gatsby and Daisy, but it might not cover all the interactions due to the small chunk size..\n", + "Selecting retriever 1: This choice will provide a more comprehensive context of the interactions between Gatsby and Daisy than the first choice due to the larger chunk size..\n", + "Selecting retriever 2: This choice will provide an even more comprehensive context of the interactions between Gatsby and Daisy than the previous choices due to the larger chunk size..\n", + "Selecting retriever 3: This choice will provide the most comprehensive context of the interactions between Gatsby and Daisy due to the largest chunk size. However, it might also include irrelevant information..\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "0e13022c086742df9154723468b136fa", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Batches: 0%| | 0/1 [00:00 1\u001b[0m \u001b[43mdisplay_response\u001b[49m\u001b[43m(\u001b[49m\u001b[43mresponse\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mshow_source\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msource_length\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m500\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mshow_source_metadata\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n", + "\u001b[0;31mTypeError\u001b[0m: display_response() got an unexpected keyword argument 'show_source_metadata'" + ] + } + ], + "source": [ + "display_response(response, show_source=True, source_length=500, show_source_metadata=True)" + ] + }, + { + "cell_type": "code", + "execution_count": 84, + "id": "cb4f4983-cf77-4d72-a199-c5e37777bd90", + "metadata": { + "scrolled": true, + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "----\n", + "In the meantime, In between time—”\n", + "\n", + "As I went over to say goodbye I saw that the expression of\n", + "bewilderment had come back into Gatsby’s face, as though a faint doubt\n", + "had occurred to him as to the quality of his present happiness. Almost\n", + "five years! There must have been moments even that afternoon when\n", + "Daisy tumbled short of his dreams—not through her own fault, but\n", + "because of the colossal vitality of his illusion. It had gone beyond\n", + "her, beyond everything. He had thrown himself into it with a creative\n", + "passion, adding to it all the time, decking it out with every bright\n", + "feather that drifted his way. No amount of fire or freshness can\n", + "challenge what a man can store up in his ghostly heart.\n", + "\n", + "As I watched him he adjusted himself a little, visibly. His hand took\n", + "hold of hers, and as she said something low in his ear he turned\n", + "toward her with a rush of emotion. I think that voice held him most,\n", + "with its fluctuating,\n", + "256\n", + "----\n", + "go downstairs,” interrupted Gatsby. He flipped a switch. The\n", + "grey windows disappeared as the house glowed full of light.\n", + "\n", + "In the music-room Gatsby turned on a solitary lamp beside the piano.\n", + "He lit Daisy’s cigarette from a trembling match, and sat down with her\n", + "on a couch far across the room, where there was no light save what the\n", + "gleaming floor bounced in from the hall.\n", + "\n", + "When Klipspringer had played “The Love Nest” he turned around on the\n", + "bench and searched unhappily for Gatsby in the gloom.\n", + "\n", + "“I’m all out of practice, you see. I told you I couldn’t play. I’m all\n", + "out of prac—”\n", + "\n", + "“Don’t talk so much, old sport,” commanded Gatsby. “Play!”\n", + "\n", + " “In the morning, In the evening, Ain’t we got fun—”\n", + "\n", + "Outside the wind was loud and there was a faint flow of thunder along\n", + "the Sound. All the lights were going on in West Egg now; the electric\n", + "trains, men-carrying, were plunging home through the rain from New\n", + "York. It was the hour of a profound human change, and excitement was\n", + "generating on the air.\n", + "\n", + " “One thing’s sure and nothing’s surer The rich get richer and the\n", + " poor get—children. In the meantime, In between time—”\n", + "\n", + "As I went over to say goodbye I saw that the expression of\n", + "bewilderment had come back into Gatsby’s face, as though a faint doubt\n", + "had occurred to him as to the quality of his present happiness. Almost\n", + "five years! There must have been moments even that afternoon when\n", + "Daisy tumbled short of his dreams—not through her own fault, but\n", + "because of the colossal vitality of his illusion. It had gone beyond\n", + "her, beyond everything. He had thrown himself into it with a creative\n", + "passion, adding to it all the time, decking it out\n", + "512\n", + "----\n", + "the\n", + "direction. In this heat every extra gesture was an affront to the\n", + "common store of life.\n", + "\n", + "The room, 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\n", + "256\n", + "----\n", + "the beach that morning. Finally we came to Gatsby’s own\n", + "apartment, a bedroom and a bath, and an Adam’s study, where we sat\n", + "down and drank a glass of some Chartreuse he took from a cupboard in\n", + "the wall.\n", + "\n", + "He hadn’t once ceased looking at Daisy, and I think he revalued\n", + "everything in his house according to the measure of response it drew\n", + "from her well-loved eyes. Sometimes too, he stared around at his\n", + "possessions in a dazed\n", + "128\n", + "----\n", + "world complete\n", + "in itself, with its own standards and its own great figures, second to\n", + "nothing because it had no consciousness of being so, and now I was\n", + "looking at it again, through Daisy’s eyes. It is invariably saddening\n", + "to look through new eyes at things upon which you have expended your\n", + "own powers of adjustment.\n", + "\n", + "They arrived at twilight, and, as we strolled out among the sparkling\n", + "hundreds, Daisy’s voice was playing murmurous tricks in her throat.\n", + "\n", + "“These things excite me so,” she whispered. “If you want to kiss me\n", + "any time during the evening, Nick, just let me know and I’ll be glad\n", + "to arrange it for you. Just mention my name. Or present a green card.\n", + "I’m giving out green—”\n", + "\n", + "“Look around,” suggested Gatsby.\n", + "\n", + "“I’m looking around. I’m having a marvellous—”\n", + "\n", + "“You must see the faces of many people you’ve heard about.”\n", + "\n", + "Tom’s arrogant eyes roamed the crowd.\n", + "\n", + "“We don’t go around very much,” he said; “in fact, I was just thinking\n", + "I don’t know a soul here.”\n", + "\n", + "“Perhaps you know that lady.” Gatsby indicated a gorgeous, scarcely\n", + "human orchid of a woman who sat in state under a white-plum tree. Tom\n", + "and Daisy stared, with that peculiarly unreal feeling that accompanies\n", + "the recognition of a hitherto ghostly celebrity of the movies.\n", + "\n", + "“She’s lovely,” said Daisy.\n", + "\n", + "“The man bending over her is her director.”\n", + "\n", + "He took them ceremoniously from group to group:\n", + "\n", + "“Mrs. Buchanan … and Mr. Buchanan—” After an instant’s hesitation he\n", + "added: “the polo player.”\n", + "\n", + "“Oh no,” objected Tom quickly,\n", + "512\n", + "----\n", + "The\n", + "grey windows disappeared as the house glowed full of light.\n", + "\n", + "In the music-room Gatsby turned on a solitary lamp beside the piano.\n", + "He lit Daisy’s cigarette from a trembling match, and sat down with her\n", + "on a couch far across the room, where there was no light save what the\n", + "gleaming floor bounced in from the hall.\n", + "\n", + "When Klipspringer had played “The Love Nest” he turned around on the\n", + "bench and searched unhappily for Gatsby in the gloom.\n", + "\n", + "“I’m all out of practice, you see. I told you I couldn’t play. I’m all\n", + "out of prac—”\n", + "\n", + "“Don’t talk so much, old sport,” commanded Gatsby. “Play!”\n", + "\n", + " “In the morning, In the evening, Ain’t we got fun—”\n", + "\n", + "Outside the wind was loud and there was a faint flow of thunder along\n", + "the Sound. All the lights were going on in West Egg now; the electric\n", + "trains, men-carrying, were plunging home through the rain from New\n", + "York. It was the hour of a profound human change, and excitement was\n", + "generating on the air.\n", + "\n", + " “One thing’s sure and nothing’s surer The rich get richer and the\n", + " poor get—children. In the meantime, In between time—”\n", + "\n", + "As I went over to say goodbye I saw that the expression of\n", + "bewilderment had come back into Gatsby’s face, as though a faint doubt\n", + "had occurred to him as to the quality of his present happiness. Almost\n", + "five years! There must have been moments even that afternoon when\n", + "Daisy tumbled short of his dreams—not through her own fault, but\n", + "because of the colossal vitality of his illusion. It had gone beyond\n", + "her, beyond everything. He had thrown himself into it with a creative\n", + "passion, adding to it all the time, decking it out with every bright\n", + "feather that drifted his way. No amount of fire or freshness can\n", + "challenge what a man can store up in his ghostly heart.\n", + "\n", + "As I watched him he adjusted himself a little, visibly. His hand took\n", + "hold of hers, and as she said something low in his ear he turned\n", + "toward her with a rush of emotion. I think that voice held him most,\n", + "with its fluctuating, feverish warmth, because it couldn’t be\n", + "over-dreamed—that voice was a deathless song.\n", + "\n", + "They had forgotten me, but Daisy glanced up and held out her hand;\n", + "Gatsby didn’t know me now at all. I looked once more at them and they\n", + "looked back at me, remotely, possessed by intense life. Then I went\n", + "out of the room and down the marble steps into the rain, leaving them\n", + "there together.\n", + "\n", + "\n", + " VI\n", + "\n", + "About this time an ambitious young reporter from New York arrived one\n", + "morning at Gatsby’s door and asked him if he had anything to say.\n", + "\n", + "“Anything to say about what?” inquired Gatsby politely.\n", + "\n", + "“Why—any statement to give out.”\n", + "\n", + "It transpired after a confused five minutes that the man had heard\n", + "Gatsby’s name around his office in a connection which he either\n", + "wouldn’t reveal or didn’t fully understand. This was his day off and\n", + "with laudable initiative he had hurried out “to see.”\n", + "\n", + "It was a random shot, and yet the reporter’s instinct was right.\n", + "Gatsby’s notoriety, spread about by the hundreds who had accepted his\n", + "hospitality and so become authorities upon his past, had increased all\n", + "summer until he fell just short of being news. Contemporary legends\n", + "such as the “underground pipeline to Canada” attached themselves to\n", + "him, and there was one persistent story that he didn’t live in a house\n", + "at all, but in a boat that looked like a house and was\n", + "1024\n", + "----\n", + "Daisy insistently. Gatsby’s eyes\n", + "floated toward her. “Ah,” she cried, “you look so cool.”\n", + "\n", + "Their eyes met, and they stared together at each other, alone in\n", + "space. With an effort she glanced down at the table.\n", + "\n", + "“You always look so cool,” she repeated.\n", + "\n", + "She had told him that she loved him, and Tom Buchanan saw. He was\n", + "astounded. His mouth opened a little, and he looked at Gatsby, and\n", + "then back at Daisy as if he had just recognized her as someone he knew\n", + "a long time ago.\n", + "\n", + "“You resemble the advertisement of the man,” she went on innocently.\n", + "“You know the advertisement of the man—”\n", + "\n", + "“All right,” broke in Tom quickly, “I’m perfectly willing to go to\n", + "town. Come on—we’re all going to town.”\n", + "\n", + "He got up, his eyes still flashing between Gatsby and his wife. No one\n", + "moved.\n", + "\n", + "“Come on!” His temper cracked a little. “What’s the matter, anyhow?\n", + "If we’re going to town, let’s start.”\n", + "\n", + "His hand, trembling with his effort at self-control, bore to his lips\n", + "the last of his glass of ale. Daisy’s voice got us to our feet and out\n", + "on to the blazing gravel drive.\n", + "\n", + "“Are we just going to go?” she objected. “Like this? Aren’t we going\n", + "to let anyone smoke a cigarette first?”\n", + "\n", + "“Everybody smoked all through lunch.”\n", + "\n", + "“Oh, let’s have fun,” she begged him. “It’s too hot to fuss.”\n", + "\n", + "He didn’t answer.\n", + "\n", + "“Have it your own way,” she said. “Come on, Jordan.”\n", + "\n", + "They went upstairs to get ready while we three men stood there\n", + "shuffling the hot pebbles with our feet. A silver curve of the moon\n", + "hovered already in the western sky. Gatsby started to speak, changed\n", + "his mind, but not before Tom wheeled and faced him expectantly.\n", + "\n", + "“Have you got your stables here?” asked Gatsby with an effort.\n", + "\n", + "“About a quarter of a mile down the road.”\n", + "\n", + "“Oh.”\n", + "\n", + "A pause.\n", + "\n", + "“I don’t see the idea of going to town,” broke out Tom savagely.\n", + "“Women get these notions in their heads—”\n", + "\n", + "“Shall we take anything to drink?” called Daisy from an upper window.\n", + "\n", + "“I’ll get some whisky,” answered Tom. He went inside.\n", + "\n", + "Gatsby turned to me rigidly:\n", + "\n", + "“I can’t say anything in his house, old sport.”\n", + "\n", + "“She’s got an indiscreet voice,” I remarked. “It’s full of—” I\n", + "hesitated.\n", + "\n", + "“Her voice is full of money,” he said suddenly.\n", + "\n", + "That was it. I’d never understood before. It was full of money—that\n", + "was the inexhaustible charm that rose and fell in it, the jingle of\n", + "it, the cymbals’ song of it … High in a white palace the king’s\n", + "daughter, the golden girl …\n", + "\n", + "Tom came out of the house wrapping a quart bottle in a towel, followed\n", + "by Daisy and Jordan wearing small tight hats of metallic cloth and\n", + "carrying light capes over their arms.\n", + "\n", + "“Shall we all go in my car?” suggested Gatsby. He felt the hot, green\n", + "leather of the seat. “I ought to have left it in the shade.”\n", + "\n", + "“Is it standard shift?” demanded Tom.\n", + "\n", + "“Yes.”\n", + "\n", + "“Well, you take my coupé and let me drive your car to town.”\n", + "\n", + "The suggestion was distasteful to Gatsby.\n", + "\n", + "“I\n", + "1024\n", + "----\n", + "turn out as he had\n", + "imagined. He had intended, probably, to take what he could and go—but\n", + "now he found that he had committed himself to the following of a\n", + "grail. He knew that Daisy was extraordinary, but he didn’t realize\n", + "just how extraordinary a “nice” girl could be. She vanished into her\n", + "rich house, into her rich, full life, leaving Gatsby—nothing. He felt\n", + "married to her, that was all.\n", + "\n", + "When they met again, two days later, it\n", + "128\n" + ] + } + ], + "source": [ + "for node in response.source_nodes:\n", + " print('----')\n", + " print(node.node.get_content())\n", + " print(node.node.metadata[\"chunk_size\"])" + ] + }, + { + "cell_type": "code", + "execution_count": 123, + "id": "0a7a8303-be94-45c5-8bc5-13ec8c7f1694", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# compute the average precision for each chunk size based on positioning in combined ranking\n", + "from collections import defaultdict\n", + "import pandas as pd\n", + "\n", + "def mrr_all(metadata_values, metadata_key, source_nodes):\n", + " # source nodes is a ranked list\n", + " # go through each value, find out positioning in source_nodes\n", + " value_to_mrr_dict = {}\n", + " for metadata_value in metadata_values:\n", + " mrr = 0\n", + " for idx, source_node in enumerate(source_nodes):\n", + " if source_node.node.metadata[metadata_key] == metadata_value:\n", + " mrr = 1 / (idx+1)\n", + " break\n", + " else:\n", + " continue\n", + " \n", + " # normalize AP, set in dict\n", + " value_to_mrr_dict[metadata_value] = mrr\n", + " \n", + " df = pd.DataFrame(value_to_mrr_dict, index=[\"MRR\"])\n", + " df.style.set_caption(\"Mean Reciprocal Rank\")\n", + " return df\n", + " " + ] + }, + { + "cell_type": "code", + "execution_count": 125, + "id": "adebbb82-764e-4b45-933e-84bf4ad64d40", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Mean Reciprocal Rank for each Chunk Size\n" + ] + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
1282565121024
MRR0.251.00.50.166667
\n", + "
" + ], + "text/plain": [ + " 128 256 512 1024\n", + "MRR 0.25 1.0 0.5 0.166667" + ] + }, + "execution_count": 125, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Compute the Mean Reciprocal Rank for each chunk size (higher is better)\n", + "# we can see that chunk size of 256 has the highest ranked results.\n", + "print('Mean Reciprocal Rank for each Chunk Size')\n", + "mrr_all(chunk_sizes, \"chunk_size\", response.source_nodes)" + ] + }, + { + "cell_type": "markdown", + "id": "b27a2f3c-55ce-4fa6-a15a-be539723a967", + "metadata": {}, + "source": [ + "## Compare Against Baseline\n", + "\n", + "Compare against a baseline of chunk size 1024 (k=2)" + ] + }, + { + "cell_type": "code", + "execution_count": 132, + "id": "a4d66b14-4f38-4b61-809c-f603d7e09ef9", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "query_engine_1024 = query_engines[-1]" + ] + }, + { + "cell_type": "code", + "execution_count": 134, + "id": "43f3e441-f372-4df2-ae21-71fa7968e606", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "response_1024 = query_engine_1024.query(\"Describe and summarize the interactions between Gatsby and Daisy\")" + ] + }, + { + "cell_type": "code", + "execution_count": 144, + "id": "4b161e69-da17-4e4e-b8c0-b9c846ce723f", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "data": { + "text/markdown": [ + "**`Final Response:`** Gatsby and Daisy have a passionate and intense relationship. They are deeply in love, and Daisy is drawn to Gatsby's mysterious and glamorous lifestyle. Gatsby is devoted to Daisy, and he is willing to do anything to make her happy. In the scene described, Daisy and Gatsby are alone in the music room, and Gatsby has lit Daisy's cigarette from a trembling match. Klipspringer is playing the piano, and Gatsby tells him to stop talking and start playing. Daisy compliments Gatsby on his coolness, and they share a moment of intense eye contact. Tom Buchanan then interrupts them, and Daisy and Gatsby are forced to part. Gatsby is left with a look of bewilderment on his face, as if he is questioning the quality of his happiness. Daisy and Gatsby then part ways, with Daisy holding out her hand to Gatsby and Gatsby not recognizing the reporter who is watching them. The reporter is left with the impression that Gatsby and Daisy are deeply in love and that their relationship is full of emotion and passion." + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "---" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "**`Source Node 1/2`**" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "**Node ID:** 9594f1c0-2662-420c-b7e9-daead80647a4
**Similarity:** 0.8610318997979541
**Text:** The\n", + "grey windows disappeared as the house glowed full of light.\n", + "\n", + "In the music-room Gatsby turned on a solitary lamp beside the piano.\n", + "He lit Daisy’s cigarette from a trembling match, and sat down with her\n", + "on a couch far across the room, where there was no light save what the\n", + "gleaming floor bounced in from the hall.\n", + "\n", + "When Klipspringer had played “The Love Nest” he turned around on the\n", + "bench and searched unhappily for Gatsby in the gloom.\n", + "\n", + "“I’m all out of practice, you see. I told you I couldn’...
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "---" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "**`Source Node 2/2`**" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/markdown": [ + "**Node ID:** 7949c539-aa6a-46b6-b064-a68763433f77
**Similarity:** 0.8589377236179734
**Text:** Daisy insistently. Gatsby’s eyes\n", + "floated toward her. “Ah,” she cried, “you look so cool.”\n", + "\n", + "Their eyes met, and they stared together at each other, alone in\n", + "space. With an effort she glanced down at the table.\n", + "\n", + "“You always look so cool,” she repeated.\n", + "\n", + "She had told him that she loved him, and Tom Buchanan saw. He was\n", + "astounded. His mouth opened a little, and he looked at Gatsby, and\n", + "then back at Daisy as if he had just recognized her as someone he knew\n", + "a long time ago.\n", + "\n", + "“You resemble the adver...
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "display_response(response_1024, show_source=True, source_length=500)" + ] + } + ], + "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", + "version": "3.10.10" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/llama_index/response/notebook_utils.py b/llama_index/response/notebook_utils.py index 5810ed3cd3876..777415fa72fed 100644 --- a/llama_index/response/notebook_utils.py +++ b/llama_index/response/notebook_utils.py @@ -18,7 +18,11 @@ def display_image(img_str: str, size: Tuple[int, int] = DEFAULT_THUMBNAIL_SIZE) display(img) -def display_source_node(source_node: NodeWithScore, source_length: int = 100) -> None: +def display_source_node( + source_node: NodeWithScore, + source_length: int = 100, + show_source_metadata: bool = False, +) -> None: """Display source node for jupyter notebook.""" source_text_fmt = truncate_text( source_node.node.get_content().strip(), source_length @@ -28,6 +32,8 @@ def display_source_node(source_node: NodeWithScore, source_length: int = 100) -> f"**Similarity:** {source_node.score}
" f"**Text:** {source_text_fmt}
" ) + if show_source_metadata: + text_md += f"**Metadata:** {source_node.node.metadata}
" if isinstance(source_node.node, ImageNode): text_md += "**Image:**" @@ -46,6 +52,7 @@ def display_response( source_length: int = 100, show_source: bool = False, show_metadata: bool = False, + show_source_metadata: bool = False, ) -> None: """Display response for jupyter notebook.""" if response.response is None: @@ -60,7 +67,11 @@ def display_response( display( Markdown(f"**`Source Node {ind + 1}/{len(response.source_nodes)}`**") ) - display_source_node(source_node, source_length=source_length) + display_source_node( + source_node, + source_length=source_length, + show_source_metadata=show_source_metadata, + ) if show_metadata: if response.metadata is not None: display_metadata(response.metadata)