wehub-resource-sync
a0c8464e58
Build Package / build (ubuntu-latest) (push) Failing after 1s
CodeQL / Analyze (python) (push) Failing after 1s
Core Typecheck / core-typecheck (push) Failing after 1s
Linting / lint (push) Failing after 1s
llama-dev tests / test-llama-dev (push) Failing after 1s
Publish Sub-Package to PyPI if Needed / publish_subpackage_if_needed (push) Has been skipped
Sync Docs to Developer Hub / sync-docs (push) Failing after 0s
Build Package / build (windows-latest) (push) Has been cancelled
634 lines
25 KiB
Plaintext
634 lines
25 KiB
Plaintext
{
|
||
"cells": [
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "ca3bd17d-aec3-4848-ac82-def6e2d6fa18",
|
||
"metadata": {},
|
||
"source": [
|
||
"# Examples of Structured Data Extraction in LlamaIndex\n",
|
||
"\n",
|
||
"<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/structured_outputs/structured_outputs.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
|
||
"\n",
|
||
"If you haven't yet read our [structured data extraction tutorial](../../understanding/extraction/index.md), we recommend starting there. This notebook demonstrates some of the techniques introduced in the tutorial.\n",
|
||
"\n",
|
||
"We start with the simple syntax around LLMs, then move on to how to use it with higher-level modules like a query engine and agent.\n",
|
||
"\n",
|
||
"A lot of the underlying behavior around structured outputs is powered by our Pydantic Program modules. Check out our [in-depth structured outputs guide](https://docs.llamaindex.ai/en/stable/module_guides/querying/structured_outputs/) for more details."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "6625f456-bf9c-4c89-948c-488909591855",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import nest_asyncio\n",
|
||
"\n",
|
||
"nest_asyncio.apply()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "57ced509-b1fd-4413-9ac2-662fcfab1074",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from llama_index.llms.openai import OpenAI\n",
|
||
"from llama_index.embeddings.openai import OpenAIEmbedding\n",
|
||
"from llama_index.core import Settings\n",
|
||
"\n",
|
||
"llm = OpenAI(model=\"gpt-4o\")\n",
|
||
"embed_model = OpenAIEmbedding(model=\"text-embedding-3-small\")\n",
|
||
"Settings.llm = llm\n",
|
||
"Settings.embed_model = embed_model"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "72a15e8d-45dc-4aef-93a6-8d2e97604d1e",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 1. Simple Structured Extraction\n",
|
||
"\n",
|
||
"You can convert any LLM to a \"structured LLM\" by attaching an output class to it through `as_structured_llm`.\n",
|
||
"\n",
|
||
"Here we pass a simple `Album` class which contains a list of songs. We can then use the normal LLM endpoints like chat/complete.\n",
|
||
"\n",
|
||
"**NOTE**: async is supported but streaming is coming soon."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "63045118-74fb-4ec6-bb99-b66184eba017",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from typing import List\n",
|
||
"from pydantic import BaseModel, Field\n",
|
||
"\n",
|
||
"\n",
|
||
"class Song(BaseModel):\n",
|
||
" \"\"\"Data model for a song.\"\"\"\n",
|
||
"\n",
|
||
" title: str\n",
|
||
" length_seconds: int\n",
|
||
"\n",
|
||
"\n",
|
||
"class Album(BaseModel):\n",
|
||
" \"\"\"Data model for an album.\"\"\"\n",
|
||
"\n",
|
||
" name: str\n",
|
||
" artist: str\n",
|
||
" songs: List[Song]"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "d5929511-67c3-411a-925c-5155286d10cb",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from llama_index.core.llms import ChatMessage\n",
|
||
"\n",
|
||
"sllm = llm.as_structured_llm(output_cls=Album)\n",
|
||
"input_msg = ChatMessage.from_str(\"Generate an example album from The Shining\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "d8a8fd46-dbc9-4973-bf70-aa0f28a7fa2d",
|
||
"metadata": {},
|
||
"source": [
|
||
"#### Sync"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "4671ba4b-d21f-4a4d-8aa8-555a679248c8",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"output = sllm.chat([input_msg])\n",
|
||
"# get actual object\n",
|
||
"output_obj = output.raw"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "c58e29dc-b6e5-4fdb-83c4-ccb011dde2fb",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"assistant: {\"name\": \"The Shining: Original Soundtrack\", \"artist\": \"Various Artists\", \"songs\": [{\"title\": \"Main Title\", \"length_seconds\": 180}, {\"title\": \"Rocky Mountains\", \"length_seconds\": 210}, {\"title\": \"Lontano\", \"length_seconds\": 720}, {\"title\": \"Music for Strings, Percussion and Celesta\", \"length_seconds\": 540}, {\"title\": \"Utrenja (Excerpt)\", \"length_seconds\": 300}, {\"title\": \"The Awakening of Jacob\", \"length_seconds\": 480}, {\"title\": \"De Natura Sonoris No. 2\", \"length_seconds\": 540}, {\"title\": \"Home\", \"length_seconds\": 180}, {\"title\": \"Midnight, the Stars and You\", \"length_seconds\": 180}, {\"title\": \"It's All Forgotten Now\", \"length_seconds\": 150}, {\"title\": \"Masquerade\", \"length_seconds\": 180}]}\n",
|
||
"name='The Shining: Original Soundtrack' artist='Various Artists' songs=[Song(title='Main Title', length_seconds=180), Song(title='Rocky Mountains', length_seconds=210), Song(title='Lontano', length_seconds=720), Song(title='Music for Strings, Percussion and Celesta', length_seconds=540), Song(title='Utrenja (Excerpt)', length_seconds=300), Song(title='The Awakening of Jacob', length_seconds=480), Song(title='De Natura Sonoris No. 2', length_seconds=540), Song(title='Home', length_seconds=180), Song(title='Midnight, the Stars and You', length_seconds=180), Song(title=\"It's All Forgotten Now\", length_seconds=150), Song(title='Masquerade', length_seconds=180)]\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"print(str(output))\n",
|
||
"print(output_obj)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "a93195eb-b62b-43eb-acad-0652f745bc64",
|
||
"metadata": {},
|
||
"source": [
|
||
"#### Async"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "c8c4e582-5aa8-42e8-8d27-d9b18b8452d8",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"assistant: {\"name\": \"The Shining: Original Soundtrack\", \"artist\": \"Various Artists\", \"songs\": [{\"title\": \"Main Title (The Shining)\", \"length_seconds\": 180}, {\"title\": \"Rocky Mountains\", \"length_seconds\": 210}, {\"title\": \"Lontano\", \"length_seconds\": 240}, {\"title\": \"Music for Strings, Percussion and Celesta\", \"length_seconds\": 300}, {\"title\": \"Utrenja (Excerpt)\", \"length_seconds\": 180}, {\"title\": \"The Awakening of Jacob\", \"length_seconds\": 150}, {\"title\": \"De Natura Sonoris No. 2\", \"length_seconds\": 270}, {\"title\": \"Home\", \"length_seconds\": 200}, {\"title\": \"Heartbeats and Worry\", \"length_seconds\": 160}, {\"title\": \"The Overlook\", \"length_seconds\": 220}]}\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"output = await sllm.achat([input_msg])\n",
|
||
"# get actual object\n",
|
||
"output_obj = output.raw\n",
|
||
"print(str(output))"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "c6608650-c3a9-4c09-92af-9c61376d291c",
|
||
"metadata": {},
|
||
"source": [
|
||
"#### Streaming"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "403ab86e-b1ca-4f4c-bcb3-4c5173ed17f7",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"{'artist': 'Various Artists',\n",
|
||
" 'name': 'The Shining: Original Soundtrack',\n",
|
||
" 'songs': [{'length_seconds': 180, 'title': 'Main Title'},\n",
|
||
" {'length_seconds': 210, 'title': 'Rocky Mountains'},\n",
|
||
" {'length_seconds': 240, 'title': 'Lontano'},\n",
|
||
" {'length_seconds': 540,\n",
|
||
" 'title': 'Music for Strings, Percussion and Celesta'},\n",
|
||
" {'length_seconds': 300, 'title': 'Utrenja (Excerpt)'},\n",
|
||
" {'length_seconds': 360, 'title': 'The Awakening of Jacob'},\n",
|
||
" {'length_seconds': 420, 'title': 'De Natura Sonoris No. 2'},\n",
|
||
" {'length_seconds': 180, 'title': 'Home'},\n",
|
||
" {'length_seconds': 180, 'title': 'Midnight, the Stars and You'},\n",
|
||
" {'length_seconds': 150, 'title': \"It's All Forgotten Now\"},\n",
|
||
" {'length_seconds': 120, 'title': 'Masquerade'}]}\n",
|
||
"assistant: {\"name\": \"The Shining: Original Soundtrack\", \"artist\": \"Various Artists\", \"songs\": [{\"title\": \"Main Title (The Shining)\", \"length_seconds\": 180}, {\"title\": \"Rocky Mountains\", \"length_seconds\": 210}, {\"title\": \"Lontano\", \"length_seconds\": 240}, {\"title\": \"Music for Strings, Percussion and Celesta\", \"length_seconds\": 300}, {\"title\": \"Utrenja (Excerpt)\", \"length_seconds\": 180}, {\"title\": \"The Awakening of Jacob\", \"length_seconds\": 150}, {\"title\": \"De Natura Sonoris No. 2\", \"length_seconds\": 270}, {\"title\": \"Home\", \"length_seconds\": 200}, {\"title\": \"Heartbeats and Worry\", \"length_seconds\": 160}, {\"title\": \"The Overlook\", \"length_seconds\": 220}]}\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"from IPython.display import clear_output\n",
|
||
"from pprint import pprint\n",
|
||
"\n",
|
||
"stream_output = sllm.stream_chat([input_msg])\n",
|
||
"for partial_output in stream_output:\n",
|
||
" clear_output(wait=True)\n",
|
||
" pprint(partial_output.raw.dict())\n",
|
||
"\n",
|
||
"output_obj = partial_output.raw\n",
|
||
"print(str(output))"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "900037f6-3fa9-4f55-bcfc-985131517f36",
|
||
"metadata": {},
|
||
"source": [
|
||
"#### Async Streaming"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "32c0025b-b657-42c8-a873-f51929431b8c",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"{'artist': 'Various Artists',\n",
|
||
" 'name': 'The Shining: Original Soundtrack',\n",
|
||
" 'songs': [{'length_seconds': 180, 'title': 'Main Title'},\n",
|
||
" {'length_seconds': 210, 'title': 'Rocky Mountains'},\n",
|
||
" {'length_seconds': 720, 'title': 'Lontano'},\n",
|
||
" {'length_seconds': 540,\n",
|
||
" 'title': 'Music for Strings, Percussion and Celesta'},\n",
|
||
" {'length_seconds': 300, 'title': 'Utrenja (Excerpt)'},\n",
|
||
" {'length_seconds': 480, 'title': 'The Awakening of Jacob'},\n",
|
||
" {'length_seconds': 540, 'title': 'De Natura Sonoris No. 2'},\n",
|
||
" {'length_seconds': 180, 'title': 'Home'},\n",
|
||
" {'length_seconds': 180, 'title': 'Midnight, the Stars and You'},\n",
|
||
" {'length_seconds': 180, 'title': \"It's All Forgotten Now\"},\n",
|
||
" {'length_seconds': 180, 'title': 'Masquerade'}]}\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"from IPython.display import clear_output\n",
|
||
"from pprint import pprint\n",
|
||
"\n",
|
||
"stream_output = await sllm.astream_chat([input_msg])\n",
|
||
"async for partial_output in stream_output:\n",
|
||
" clear_output(wait=True)\n",
|
||
" pprint(partial_output.raw.dict())"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "9d1b4bc9-56ee-4bf7-a63c-50ce0cadeeae",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 1.b Use the `structured_predict` Function\n",
|
||
"\n",
|
||
"Instead of explicitly doing `llm.as_structured_llm(...)`, every LLM class has a `structured_predict` function which allows you to more easily call the LLM with a prompt template + template variables to return a strutured output in one line of code."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "2488b69e-c04f-48fe-afea-7224dc79efa9",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"data": {
|
||
"text/plain": [
|
||
"Album(name='Songs of Middle-earth', artist='Various Artists', songs=[Song(title='The Shire', length_seconds=180), Song(title='The Fellowship', length_seconds=240), Song(title=\"Gollum's Theme\", length_seconds=200), Song(title=\"Rohan's Call\", length_seconds=220), Song(title=\"The Battle of Helm's Deep\", length_seconds=300), Song(title='Lothlórien', length_seconds=210), Song(title='The Return of the King', length_seconds=250), Song(title='Into the West', length_seconds=260)])"
|
||
]
|
||
},
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"output_type": "execute_result"
|
||
}
|
||
],
|
||
"source": [
|
||
"# use query pipelines\n",
|
||
"from llama_index.core.prompts import ChatPromptTemplate\n",
|
||
"from llama_index.core.llms import ChatMessage\n",
|
||
"from llama_index.llms.openai import OpenAI\n",
|
||
"\n",
|
||
"chat_prompt_tmpl = ChatPromptTemplate(\n",
|
||
" message_templates=[\n",
|
||
" ChatMessage.from_str(\n",
|
||
" \"Generate an example album from {movie_name}\", role=\"user\"\n",
|
||
" )\n",
|
||
" ]\n",
|
||
")\n",
|
||
"\n",
|
||
"llm = OpenAI(model=\"gpt-4o\")\n",
|
||
"album = llm.structured_predict(\n",
|
||
" Album, chat_prompt_tmpl, movie_name=\"Lord of the Rings\"\n",
|
||
")\n",
|
||
"album"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "ae13bf7e-888a-4305-8ed8-d22f130ab6fe",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 2. Plug into RAG Pipeline\n",
|
||
"\n",
|
||
"You can also plug this into a RAG pipeline. Below we show structured extraction from an Apple 10K report."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "c72151ca-b471-41ad-a762-f483ca341d9b",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"!mkdir data\n",
|
||
"!wget \"https://s2.q4cdn.com/470004039/files/doc_financials/2021/q4/_10-K-2021-(As-Filed).pdf\" -O data/apple_2021_10k.pdf"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "f1f81633-cdcb-4e8d-a43f-43b699a0cf01",
|
||
"metadata": {},
|
||
"source": [
|
||
"#### Option 1: Use LlamaParse\n",
|
||
"\n",
|
||
"You will need an account at https://cloud.llamaindex.ai/ and an API Key to use LlamaParse, our document parser for 10K filings."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "d2745476-02be-4f2a-86ff-5386aae1260d",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Started parsing the file under job_id cac11eca-7e00-452f-93f6-19c861b4c130\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"from llama_parse import LlamaParse\n",
|
||
"\n",
|
||
"# os.environ[\"LLAMA_CLOUD_API_KEY\"] = \"llx-...\"\n",
|
||
"orig_docs = LlamaParse(result_type=\"text\").load_data(\n",
|
||
" \"./data/apple_2021_10k.pdf\"\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "9509b85e-0091-4fc0-8649-5a3d83d6dcf0",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
" UNITED STATES\n",
|
||
" SECURITIES AND EXCHANGE COMMISSION\n",
|
||
" Washington, D.C. 20549\n",
|
||
"\n",
|
||
" FORM 10-K\n",
|
||
"(Mark One)\n",
|
||
" ☒ ANNUAL REPORT PURSUANT TO SECTION 13 OR 15(d) OF THE SECURITIES EXCHANGE ACT OF 1934\n",
|
||
" For the fiscal year ended September 25, 2021\n",
|
||
" or\n",
|
||
" ☐ TRANSITION REPORT PURSUANT TO SECTION 13 OR 15(d) OF THE SECURITIES EXCHANGE ACT OF 1934\n",
|
||
" For the transition period from to .\n",
|
||
" Commission File Number: 001-36743\n",
|
||
"\n",
|
||
" Apple Inc.\n",
|
||
" (Exact name of Registrant as specified in its charter)\n",
|
||
"\n",
|
||
" California 94-2404110\n",
|
||
" (State or other jurisdiction (I.R.S. Employer Identification No.)\n",
|
||
" of incorporation or organization)\n",
|
||
"\n",
|
||
" One Apple Park Way\n",
|
||
" Cupertino, California 95014\n",
|
||
" (Address of principal executive offices) (Zip Code)\n",
|
||
" (408) 996-1010\n",
|
||
" (Registrant’s telephone number, including area code)\n",
|
||
"\n",
|
||
" Securities registered pursuant to Section 12(b) of the Act:\n",
|
||
"\n",
|
||
" Trading\n",
|
||
" Title of each class symbol(s) Name of each exchange on which registered\n",
|
||
" Common Stock, $0.00001 par value per share AAPL The Nasdaq Stock Market LLC\n",
|
||
" 1.000% Notes due 2022 — The Nasdaq Stock Market LLC\n",
|
||
" 1.375% Notes due 2024 — The Nasdaq Stock Market LLC\n",
|
||
" 0.000% Notes due 2025 — The Nasdaq Stock Market LLC\n",
|
||
" 0.875% Notes due 2025 — The Nasdaq Stock Market LLC\n",
|
||
" 1.625% Notes due 2026 — The Nasdaq Stock Market LLC\n",
|
||
" 2.000% Notes due 2027 — The Nasdaq Stock Market LLC\n",
|
||
" 1.375% Notes due 2029 — The Nasdaq Stock Market LLC\n",
|
||
" 3.050% Notes due 2029 — The Nasdaq Stock Market LLC\n",
|
||
" 0.500% Notes due 2031 — The Nasdaq Stock Market LLC\n",
|
||
" 3.600% Notes due 2042 — The Nasdaq Stock Market LLC\n",
|
||
"\n",
|
||
" Securities registered pursuant to Section 12(g) of the Act: None\n",
|
||
"\n",
|
||
"Indicate by check mark if the Registrant is a well-known seasoned issuer, as defined in Rule 405 of the Securities Act.\n",
|
||
" Yes ☒ No ☐\n",
|
||
"Indicate by check mark if the Registrant is not required to file reports pursuant to Section 13 or Section 15(d) of the Act.\n",
|
||
" Yes ☐ No ☒\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"from copy import deepcopy\n",
|
||
"from llama_index.core.schema import TextNode\n",
|
||
"\n",
|
||
"\n",
|
||
"def get_page_nodes(docs, separator=\"\\n---\\n\"):\n",
|
||
" \"\"\"Split each document into page node, by separator.\"\"\"\n",
|
||
" nodes = []\n",
|
||
" for doc in docs:\n",
|
||
" doc_chunks = doc.text.split(separator)\n",
|
||
" for doc_chunk in doc_chunks:\n",
|
||
" node = TextNode(\n",
|
||
" text=doc_chunk,\n",
|
||
" metadata=deepcopy(doc.metadata),\n",
|
||
" )\n",
|
||
" nodes.append(node)\n",
|
||
"\n",
|
||
" return nodes\n",
|
||
"\n",
|
||
"\n",
|
||
"docs = get_page_nodes(orig_docs)\n",
|
||
"print(docs[0].get_content())"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "2ae6c400-e661-4a5c-8ed0-1f31bf8325cb",
|
||
"metadata": {},
|
||
"source": [
|
||
"#### Option 2: Use SimpleDirectoryReader\n",
|
||
"\n",
|
||
"You can also choose to use the free PDF parser bundled into our `SimpleDirectoryReader`."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "ee757df9-0491-46ba-b369-d8cdf1b7d74f",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# # OPTION 2: Use SimpleDirectoryReader\n",
|
||
"# from llama_index.core import SimpleDirectoryReader\n",
|
||
"\n",
|
||
"# reader = SimpleDirectoryReader(input_files=[\"apple_2021_10k.pdf\"])\n",
|
||
"# docs = reader.load_data()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "4f3e9304-1377-46e4-9f23-939907c6198d",
|
||
"metadata": {},
|
||
"source": [
|
||
"#### Build RAG Pipeline, Define Structured Output Schema\n",
|
||
"\n",
|
||
"We build a RAG pipeline with our trusty VectorStoreIndex and reranker module. We then define the output as a Pydantic model. This allows us to create a structured LLM with the output class attached."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "1e318de5-7c7b-49a2-8de6-140f16b07d8f",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from llama_index.core import VectorStoreIndex\n",
|
||
"\n",
|
||
"# skip chunking since we're doing page-level chunking\n",
|
||
"index = VectorStoreIndex(docs)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "6df8be76-53f6-4e81-8b22-7281e0ac35eb",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from llama_index.postprocessor.flag_embedding_reranker import (\n",
|
||
" FlagEmbeddingReranker,\n",
|
||
")\n",
|
||
"\n",
|
||
"reranker = FlagEmbeddingReranker(\n",
|
||
" top_n=5,\n",
|
||
" model=\"BAAI/bge-reranker-large\",\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "d3e52dfc-31a5-41bf-972c-284c8a3263d3",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from pydantic import BaseModel, Field\n",
|
||
"from typing import List\n",
|
||
"\n",
|
||
"\n",
|
||
"class Output(BaseModel):\n",
|
||
" \"\"\"Output containing the response, page numbers, and confidence.\"\"\"\n",
|
||
"\n",
|
||
" response: str = Field(..., description=\"The answer to the question.\")\n",
|
||
" page_numbers: List[int] = Field(\n",
|
||
" ...,\n",
|
||
" description=\"The page numbers of the sources used to answer this question. Do not include a page number if the context is irrelevant.\",\n",
|
||
" )\n",
|
||
" confidence: float = Field(\n",
|
||
" ...,\n",
|
||
" description=\"Confidence value between 0-1 of the correctness of the result.\",\n",
|
||
" )\n",
|
||
" confidence_explanation: str = Field(\n",
|
||
" ..., description=\"Explanation for the confidence score\"\n",
|
||
" )\n",
|
||
"\n",
|
||
"\n",
|
||
"sllm = llm.as_structured_llm(output_cls=Output)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "1f28e413-c9ee-44dd-88c3-0f8c387341e5",
|
||
"metadata": {},
|
||
"source": [
|
||
"#### Run Queries"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "31c2d822-50cb-4eca-9c14-661c47bb70b7",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"query_engine = index.as_query_engine(\n",
|
||
" similarity_top_k=5,\n",
|
||
" node_postprocessors=[reranker],\n",
|
||
" llm=sllm,\n",
|
||
" response_mode=\"tree_summarize\", # you can also select other modes like `compact`, `refine`\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "bc211576-7101-4d1d-a874-eaf1ba5ad7a5",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"{\"response\": \"In 2021, the net sales for each product category were as follows: iPhone: $191,973 million, Mac: $35,190 million, iPad: $31,862 million, Wearables, Home and Accessories: $38,367 million, and Services: $68,425 million.\", \"page_numbers\": [21], \"confidence\": 1.0, \"confidence_explanation\": \"The figures are directly taken from the provided data, ensuring high accuracy.\"}\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"response = query_engine.query(\"Net sales for each product category in 2021\")\n",
|
||
"print(str(response))"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "340bc8a8-7460-4683-bf2e-518c35ba68af",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"data": {
|
||
"text/plain": [
|
||
"{'response': 'In 2021, the net sales for each product category were as follows: iPhone: $191,973 million, Mac: $35,190 million, iPad: $31,862 million, Wearables, Home and Accessories: $38,367 million, and Services: $68,425 million.',\n",
|
||
" 'page_numbers': [21],\n",
|
||
" 'confidence': 1.0,\n",
|
||
" 'confidence_explanation': 'The figures are directly taken from the provided data, ensuring high accuracy.'}"
|
||
]
|
||
},
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"output_type": "execute_result"
|
||
}
|
||
],
|
||
"source": [
|
||
"response.response.dict()"
|
||
]
|
||
}
|
||
],
|
||
"metadata": {
|
||
"kernelspec": {
|
||
"display_name": "llama_index_v3",
|
||
"language": "python",
|
||
"name": "llama_index_v3"
|
||
},
|
||
"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
|
||
}
|