Last active
September 23, 2024 12:07
-
-
Save alonsosilvaallende/ebb6fb864be3fb29d3d6b21a222bfdb6 to your computer and use it in GitHub Desktop.
LlamaIndex upsert nodes and relationships
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| { | |
| "cells": [ | |
| { | |
| "cell_type": "code", | |
| "execution_count": 1, | |
| "id": "a259711f-8c92-4c33-bb0b-febcf69f3804", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:15.840536Z", | |
| "iopub.status.busy": "2024-09-23T12:11:15.839441Z", | |
| "iopub.status.idle": "2024-09-23T12:11:15.873392Z", | |
| "shell.execute_reply": "2024-09-23T12:11:15.872166Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:15.840470Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "%load_ext autoreload\n", | |
| "%autoreload 2" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 2, | |
| "id": "44160830-387e-4b58-a3e1-b1f0585cb67b", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:15.874700Z", | |
| "iopub.status.busy": "2024-09-23T12:11:15.874408Z", | |
| "iopub.status.idle": "2024-09-23T12:11:15.901718Z", | |
| "shell.execute_reply": "2024-09-23T12:11:15.900561Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:15.874672Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "from dotenv import find_dotenv, load_dotenv\n", | |
| "\n", | |
| "_ = load_dotenv(find_dotenv())" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 3, | |
| "id": "6a2413d1-0235-4382-a107-4218db45782b", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:15.903775Z", | |
| "iopub.status.busy": "2024-09-23T12:11:15.903272Z", | |
| "iopub.status.idle": "2024-09-23T12:11:15.921732Z", | |
| "shell.execute_reply": "2024-09-23T12:11:15.920614Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:15.903726Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "import os\n", | |
| "\n", | |
| "openai_api_key = os.getenv(\"OPENAI_API_KEY\")" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 4, | |
| "id": "d8805d3b-db0a-40d3-a1aa-39e3062b59a2", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:15.923847Z", | |
| "iopub.status.busy": "2024-09-23T12:11:15.922928Z", | |
| "iopub.status.idle": "2024-09-23T12:11:15.941532Z", | |
| "shell.execute_reply": "2024-09-23T12:11:15.940519Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:15.923782Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "import nest_asyncio\n", | |
| "\n", | |
| "nest_asyncio.apply()" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 5, | |
| "id": "f5510434-fb17-4ef2-a24f-adef4ab17184", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:15.943081Z", | |
| "iopub.status.busy": "2024-09-23T12:11:15.942737Z", | |
| "iopub.status.idle": "2024-09-23T12:11:18.656659Z", | |
| "shell.execute_reply": "2024-09-23T12:11:18.655664Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:15.943048Z" | |
| } | |
| }, | |
| "outputs": [ | |
| { | |
| "data": { | |
| "text/html": [ | |
| "<div><style>\n", | |
| ".dataframe > thead > tr,\n", | |
| ".dataframe > tbody > tr {\n", | |
| " text-align: right;\n", | |
| " white-space: pre-wrap;\n", | |
| "}\n", | |
| "</style>\n", | |
| "<small>shape: (5, 4)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>Song</th><th>Lyrics</th><th>Album</th><th>Artist</th></tr><tr><td>str</td><td>str</td><td>str</td><td>str</td></tr></thead><tbody><tr><td>"DON'T TREAD ON ME"</td><td>"Liberty or death\n", | |
| "what we so pr…</td><td>"THE BLACK ALBUM"</td><td>"Metallica"</td></tr><tr><td>"NOTHING ELSE MATTERS"</td><td>"So close no matter how far\n", | |
| "cou…</td><td>"THE BLACK ALBUM"</td><td>"Metallica"</td></tr><tr><td>"THROUGH THE NEVER"</td><td>"All that is\n", | |
| "was and will be\n", | |
| "un…</td><td>"THE BLACK ALBUM"</td><td>"Metallica"</td></tr><tr><td>"MY FRIEND OF MISERY"</td><td>"You just stood there screaming…</td><td>"THE BLACK ALBUM"</td><td>"Metallica"</td></tr><tr><td>"WHEREVER I MAY ROAM"</td><td>"...and the road becomes my bri…</td><td>"THE BLACK ALBUM"</td><td>"Metallica"</td></tr></tbody></table></div>" | |
| ], | |
| "text/plain": [ | |
| "shape: (5, 4)\n", | |
| "┌──────────────────────┬─────────────────────────────────┬─────────────────┬───────────┐\n", | |
| "│ Song ┆ Lyrics ┆ Album ┆ Artist │\n", | |
| "│ --- ┆ --- ┆ --- ┆ --- │\n", | |
| "│ str ┆ str ┆ str ┆ str │\n", | |
| "╞══════════════════════╪═════════════════════════════════╪═════════════════╪═══════════╡\n", | |
| "│ DON'T TREAD ON ME ┆ Liberty or death ┆ THE BLACK ALBUM ┆ Metallica │\n", | |
| "│ ┆ what we so pr… ┆ ┆ │\n", | |
| "│ NOTHING ELSE MATTERS ┆ So close no matter how far ┆ THE BLACK ALBUM ┆ Metallica │\n", | |
| "│ ┆ cou… ┆ ┆ │\n", | |
| "│ THROUGH THE NEVER ┆ All that is ┆ THE BLACK ALBUM ┆ Metallica │\n", | |
| "│ ┆ was and will be ┆ ┆ │\n", | |
| "│ ┆ un… ┆ ┆ │\n", | |
| "│ MY FRIEND OF MISERY ┆ You just stood there screaming… ┆ THE BLACK ALBUM ┆ Metallica │\n", | |
| "│ WHEREVER I MAY ROAM ┆ ...and the road becomes my bri… ┆ THE BLACK ALBUM ┆ Metallica │\n", | |
| "└──────────────────────┴─────────────────────────────────┴─────────────────┴───────────┘" | |
| ] | |
| }, | |
| "execution_count": 5, | |
| "metadata": {}, | |
| "output_type": "execute_result" | |
| } | |
| ], | |
| "source": [ | |
| "import polars as pl\n", | |
| "\n", | |
| "df = pl.read_csv(\n", | |
| " \"https://drive.google.com/uc?export=download&id=1uD3h7xYxr9EoZ0Ggoh99JtQXa3AxtxyU\"\n", | |
| ")\n", | |
| "\n", | |
| "df.head()" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 6, | |
| "id": "2a9265be-cca8-486b-8835-196c745f46df", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:18.660358Z", | |
| "iopub.status.busy": "2024-09-23T12:11:18.659798Z", | |
| "iopub.status.idle": "2024-09-23T12:11:18.677306Z", | |
| "shell.execute_reply": "2024-09-23T12:11:18.676814Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:18.660320Z" | |
| } | |
| }, | |
| "outputs": [ | |
| { | |
| "data": { | |
| "text/html": [ | |
| "<div><style>\n", | |
| ".dataframe > thead > tr,\n", | |
| ".dataframe > tbody > tr {\n", | |
| " text-align: right;\n", | |
| " white-space: pre-wrap;\n", | |
| "}\n", | |
| "</style>\n", | |
| "<small>shape: (5, 4)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>Song</th><th>Lyrics</th><th>Album</th><th>Artist</th></tr><tr><td>str</td><td>str</td><td>str</td><td>str</td></tr></thead><tbody><tr><td>"Don't Tread On Me"</td><td>"Liberty or death\n", | |
| "what we so pr…</td><td>"The Black Album"</td><td>"Metallica"</td></tr><tr><td>"Nothing Else Matters"</td><td>"So close no matter how far\n", | |
| "cou…</td><td>"The Black Album"</td><td>"Metallica"</td></tr><tr><td>"Through The Never"</td><td>"All that is\n", | |
| "was and will be\n", | |
| "un…</td><td>"The Black Album"</td><td>"Metallica"</td></tr><tr><td>"My Friend Of Misery"</td><td>"You just stood there screaming…</td><td>"The Black Album"</td><td>"Metallica"</td></tr><tr><td>"Wherever I May Roam"</td><td>"...and the road becomes my bri…</td><td>"The Black Album"</td><td>"Metallica"</td></tr></tbody></table></div>" | |
| ], | |
| "text/plain": [ | |
| "shape: (5, 4)\n", | |
| "┌──────────────────────┬─────────────────────────────────┬─────────────────┬───────────┐\n", | |
| "│ Song ┆ Lyrics ┆ Album ┆ Artist │\n", | |
| "│ --- ┆ --- ┆ --- ┆ --- │\n", | |
| "│ str ┆ str ┆ str ┆ str │\n", | |
| "╞══════════════════════╪═════════════════════════════════╪═════════════════╪═══════════╡\n", | |
| "│ Don't Tread On Me ┆ Liberty or death ┆ The Black Album ┆ Metallica │\n", | |
| "│ ┆ what we so pr… ┆ ┆ │\n", | |
| "│ Nothing Else Matters ┆ So close no matter how far ┆ The Black Album ┆ Metallica │\n", | |
| "│ ┆ cou… ┆ ┆ │\n", | |
| "│ Through The Never ┆ All that is ┆ The Black Album ┆ Metallica │\n", | |
| "│ ┆ was and will be ┆ ┆ │\n", | |
| "│ ┆ un… ┆ ┆ │\n", | |
| "│ My Friend Of Misery ┆ You just stood there screaming… ┆ The Black Album ┆ Metallica │\n", | |
| "│ Wherever I May Roam ┆ ...and the road becomes my bri… ┆ The Black Album ┆ Metallica │\n", | |
| "└──────────────────────┴─────────────────────────────────┴─────────────────┴───────────┘" | |
| ] | |
| }, | |
| "execution_count": 6, | |
| "metadata": {}, | |
| "output_type": "execute_result" | |
| } | |
| ], | |
| "source": [ | |
| "import string\n", | |
| "\n", | |
| "df = df.with_columns(\n", | |
| " pl.Series(\"Album\", [string.capwords(album) for album in df[\"Album\"]])\n", | |
| ")\n", | |
| "df = df.with_columns(pl.Series(\"Song\", [string.capwords(song) for song in df[\"Song\"]]))\n", | |
| "df = df.with_columns(pl.col(\"Lyrics\").fill_null(\"None\"))\n", | |
| "df.head()" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 7, | |
| "id": "84f7ab9a-dffb-4124-81e0-99e29b5c13d7", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:18.678201Z", | |
| "iopub.status.busy": "2024-09-23T12:11:18.678001Z", | |
| "iopub.status.idle": "2024-09-23T12:11:18.701567Z", | |
| "shell.execute_reply": "2024-09-23T12:11:18.701017Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:18.678182Z" | |
| } | |
| }, | |
| "outputs": [ | |
| { | |
| "data": { | |
| "text/html": [ | |
| "<div><style>\n", | |
| ".dataframe > thead > tr,\n", | |
| ".dataframe > tbody > tr {\n", | |
| " text-align: right;\n", | |
| " white-space: pre-wrap;\n", | |
| "}\n", | |
| "</style>\n", | |
| "<small>shape: (5, 5)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>Song</th><th>Lyrics</th><th>Album</th><th>Artist</th><th>text</th></tr><tr><td>str</td><td>str</td><td>str</td><td>str</td><td>str</td></tr></thead><tbody><tr><td>"Don't Tread On Me"</td><td>"Liberty or death\n", | |
| "what we so pr…</td><td>"The Black Album"</td><td>"Metallica"</td><td>"# The Black Album: Don't Tread…</td></tr><tr><td>"Nothing Else Matters"</td><td>"So close no matter how far\n", | |
| "cou…</td><td>"The Black Album"</td><td>"Metallica"</td><td>"# The Black Album: Nothing Els…</td></tr><tr><td>"Through The Never"</td><td>"All that is\n", | |
| "was and will be\n", | |
| "un…</td><td>"The Black Album"</td><td>"Metallica"</td><td>"# The Black Album: Through The…</td></tr><tr><td>"My Friend Of Misery"</td><td>"You just stood there screaming…</td><td>"The Black Album"</td><td>"Metallica"</td><td>"# The Black Album: My Friend O…</td></tr><tr><td>"Wherever I May Roam"</td><td>"...and the road becomes my bri…</td><td>"The Black Album"</td><td>"Metallica"</td><td>"# The Black Album: Wherever I …</td></tr></tbody></table></div>" | |
| ], | |
| "text/plain": [ | |
| "shape: (5, 5)\n", | |
| "┌──────────────────────┬──────────────────────┬─────────────────┬───────────┬──────────────────────┐\n", | |
| "│ Song ┆ Lyrics ┆ Album ┆ Artist ┆ text │\n", | |
| "│ --- ┆ --- ┆ --- ┆ --- ┆ --- │\n", | |
| "│ str ┆ str ┆ str ┆ str ┆ str │\n", | |
| "╞══════════════════════╪══════════════════════╪═════════════════╪═══════════╪══════════════════════╡\n", | |
| "│ Don't Tread On Me ┆ Liberty or death ┆ The Black Album ┆ Metallica ┆ # The Black Album: │\n", | |
| "│ ┆ what we so pr… ┆ ┆ ┆ Don't Tread… │\n", | |
| "│ Nothing Else Matters ┆ So close no matter ┆ The Black Album ┆ Metallica ┆ # The Black Album: │\n", | |
| "│ ┆ how far ┆ ┆ ┆ Nothing Els… │\n", | |
| "│ ┆ cou… ┆ ┆ ┆ │\n", | |
| "│ Through The Never ┆ All that is ┆ The Black Album ┆ Metallica ┆ # The Black Album: │\n", | |
| "│ ┆ was and will be ┆ ┆ ┆ Through The… │\n", | |
| "│ ┆ un… ┆ ┆ ┆ │\n", | |
| "│ My Friend Of Misery ┆ You just stood there ┆ The Black Album ┆ Metallica ┆ # The Black Album: │\n", | |
| "│ ┆ screaming… ┆ ┆ ┆ My Friend O… │\n", | |
| "│ Wherever I May Roam ┆ ...and the road ┆ The Black Album ┆ Metallica ┆ # The Black Album: │\n", | |
| "│ ┆ becomes my bri… ┆ ┆ ┆ Wherever I … │\n", | |
| "└──────────────────────┴──────────────────────┴─────────────────┴───────────┴──────────────────────┘" | |
| ] | |
| }, | |
| "execution_count": 7, | |
| "metadata": {}, | |
| "output_type": "execute_result" | |
| } | |
| ], | |
| "source": [ | |
| "df = df.with_columns(\n", | |
| " text=pl.lit(\"# \")\n", | |
| " + pl.col(\"Album\")\n", | |
| " + pl.lit(\": \")\n", | |
| " + pl.col(\"Song\")\n", | |
| " + pl.lit(\"\\n\\n\")\n", | |
| " + pl.col(\"Lyrics\")\n", | |
| ")\n", | |
| "\n", | |
| "df.head()" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 8, | |
| "id": "fde0b271-d7bf-459f-9a6d-7edd70b5c117", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:18.702482Z", | |
| "iopub.status.busy": "2024-09-23T12:11:18.702282Z", | |
| "iopub.status.idle": "2024-09-23T12:11:18.841640Z", | |
| "shell.execute_reply": "2024-09-23T12:11:18.841115Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:18.702463Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "# Initialize database\n", | |
| "import shutil\n", | |
| "import kuzu\n", | |
| "\n", | |
| "shutil.rmtree(\"test_kuzudb\", ignore_errors=True) # clean up\n", | |
| "db = kuzu.Database(\"test_kuzudb\")\n", | |
| "conn = kuzu.Connection(db)" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 9, | |
| "id": "783c4e2d-7607-4e68-baac-a90e90085288", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:18.842530Z", | |
| "iopub.status.busy": "2024-09-23T12:11:18.842342Z", | |
| "iopub.status.idle": "2024-09-23T12:11:20.108946Z", | |
| "shell.execute_reply": "2024-09-23T12:11:20.107894Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:18.842513Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "from llama_index.graph_stores.kuzu import KuzuPropertyGraphStore\n", | |
| "\n", | |
| "graph_store = KuzuPropertyGraphStore(db)" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 10, | |
| "id": "73184688-8922-4a23-a5bb-3a7628c2d051", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:20.110553Z", | |
| "iopub.status.busy": "2024-09-23T12:11:20.110023Z", | |
| "iopub.status.idle": "2024-09-23T12:11:23.108495Z", | |
| "shell.execute_reply": "2024-09-23T12:11:23.106801Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:20.110523Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "from llama_index.core.graph_stores.types import Relation, EntityNode\n", | |
| "\n", | |
| "# Insert nodes\n", | |
| "graph_store.upsert_nodes([\n", | |
| " EntityNode(label=\"SONG\", name=song)\n", | |
| " for song in df[\"Song\"].unique()\n", | |
| "])\n", | |
| "\n", | |
| "graph_store.upsert_nodes([\n", | |
| " EntityNode(label=\"ALBUM\", name=album)\n", | |
| " for album in df[\"Album\"].unique()\n", | |
| "])\n", | |
| "\n", | |
| "graph_store.upsert_nodes([\n", | |
| " EntityNode(label=\"ARTIST\", name=album)\n", | |
| " for album in df[\"Artist\"].unique()\n", | |
| "])" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 11, | |
| "id": "8bba6f8f-7fb6-4c92-85c9-b4c77e2699d6", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:23.111836Z", | |
| "iopub.status.busy": "2024-09-23T12:11:23.110864Z", | |
| "iopub.status.idle": "2024-09-23T12:11:31.775313Z", | |
| "shell.execute_reply": "2024-09-23T12:11:31.773607Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:23.111764Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "graph_store.upsert_relations(\n", | |
| " [\n", | |
| " Relation(\n", | |
| " label=\"PERFORMED_BY\",\n", | |
| " source_id=album,\n", | |
| " target_id=artist,\n", | |
| " )\n", | |
| " for album, artist in df.select([\"Album\", \"Artist\"]).rows()\n", | |
| " ]\n", | |
| ")" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 12, | |
| "id": "e178b87e-3333-400b-b35f-4941e4119097", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:31.777814Z", | |
| "iopub.status.busy": "2024-09-23T12:11:31.777277Z", | |
| "iopub.status.idle": "2024-09-23T12:11:40.284659Z", | |
| "shell.execute_reply": "2024-09-23T12:11:40.282932Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:31.777763Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "graph_store.upsert_relations(\n", | |
| " [\n", | |
| " Relation(\n", | |
| " label=\"CONTAINS_SONG\",\n", | |
| " source_id=album,\n", | |
| " target_id=song,\n", | |
| " )\n", | |
| " for song, album in df.select([\"Song\", \"Album\"]).rows()\n", | |
| " ]\n", | |
| ")" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 13, | |
| "id": "956c6f16-c2a9-428a-befe-88fc96cb6339", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:40.287136Z", | |
| "iopub.status.busy": "2024-09-23T12:11:40.286587Z", | |
| "iopub.status.idle": "2024-09-23T12:11:41.161261Z", | |
| "shell.execute_reply": "2024-09-23T12:11:41.159924Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:40.287083Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "from llama_index.embeddings.openai import OpenAIEmbedding\n", | |
| "from llama_index.core import Settings\n", | |
| "Settings.embed_model = OpenAIEmbedding(model_name=\"text-embedding-3-small\")\n", | |
| "embed_model = OpenAIEmbedding(model_name=\"text-embedding-3-small\")" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 14, | |
| "id": "c3c741cb-3ac0-4ef1-8492-21345602ec59", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:41.162255Z", | |
| "iopub.status.busy": "2024-09-23T12:11:41.162015Z", | |
| "iopub.status.idle": "2024-09-23T12:11:41.482646Z", | |
| "shell.execute_reply": "2024-09-23T12:11:41.481360Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:41.162230Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "from llama_index.core import PropertyGraphIndex\n", | |
| "\n", | |
| "kg_index = PropertyGraphIndex.from_existing(\n", | |
| " property_graph_store=graph_store,\n", | |
| " embed_model=embed_model,\n", | |
| " show_progress=True,\n", | |
| ")" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 15, | |
| "id": "fad7991b-3439-44a5-bcc9-a4a4540d5e65", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:41.483648Z", | |
| "iopub.status.busy": "2024-09-23T12:11:41.483438Z", | |
| "iopub.status.idle": "2024-09-23T12:11:41.505361Z", | |
| "shell.execute_reply": "2024-09-23T12:11:41.504670Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:41.483624Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "kg_retriever = kg_index.as_retriever()\n", | |
| "kg_query_engine = kg_index.as_query_engine(include_text=False)" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 16, | |
| "id": "9642eae8-29d5-4a68-b659-b6fcadb9d272", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:41.505953Z", | |
| "iopub.status.busy": "2024-09-23T12:11:41.505808Z", | |
| "iopub.status.idle": "2024-09-23T12:11:43.775088Z", | |
| "shell.execute_reply": "2024-09-23T12:11:43.773436Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:41.505938Z" | |
| } | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "response = kg_query_engine.query(\"How many songs does The Black Album have?\")" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 17, | |
| "id": "e8dad2f4-bc09-4f83-892c-e15053b3e30c", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:43.778041Z", | |
| "iopub.status.busy": "2024-09-23T12:11:43.777449Z", | |
| "iopub.status.idle": "2024-09-23T12:11:43.820279Z", | |
| "shell.execute_reply": "2024-09-23T12:11:43.819442Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:43.777986Z" | |
| } | |
| }, | |
| "outputs": [ | |
| { | |
| "data": { | |
| "text/plain": [ | |
| "'The Black Album contains 8 songs.'" | |
| ] | |
| }, | |
| "execution_count": 17, | |
| "metadata": {}, | |
| "output_type": "execute_result" | |
| } | |
| ], | |
| "source": [ | |
| "response.response" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": 18, | |
| "id": "0a1c86ef-6063-422c-9ea7-9e8757cb500f", | |
| "metadata": { | |
| "execution": { | |
| "iopub.execute_input": "2024-09-23T12:11:43.821528Z", | |
| "iopub.status.busy": "2024-09-23T12:11:43.821165Z", | |
| "iopub.status.idle": "2024-09-23T12:11:43.862722Z", | |
| "shell.execute_reply": "2024-09-23T12:11:43.861258Z", | |
| "shell.execute_reply.started": "2024-09-23T12:11:43.821506Z" | |
| } | |
| }, | |
| "outputs": [ | |
| { | |
| "name": "stdout", | |
| "output_type": "stream", | |
| "text": [ | |
| "Master Of Puppets ({'id': 'Master Of Puppets', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> PERFORMED_BY -> Metallica ({'id': 'Master Of Puppets', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "Master Of Puppets ({'id': 'Master Of Puppets', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Battery ({'id': 'Master Of Puppets', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "Master Of Puppets ({'id': 'Master Of Puppets', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Disposable Heroes ({'id': 'Master Of Puppets', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "Master Of Puppets ({'id': 'Master Of Puppets', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Master Of Puppets ({'id': 'Master Of Puppets', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "Master Of Puppets ({'id': 'Master Of Puppets', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> The Thing That Should Not Be ({'id': 'Master Of Puppets', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "Master Of Puppets ({'id': 'Master Of Puppets', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Damage Inc. ({'id': 'Master Of Puppets', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "Master Of Puppets ({'id': 'Master Of Puppets', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Welcome Home ({'id': 'Master Of Puppets', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "Master Of Puppets ({'id': 'Master Of Puppets', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Leper Messiah ({'id': 'Master Of Puppets', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "St. Anger ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> PERFORMED_BY -> Metallica ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "St. Anger ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Shoot Me Again ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "St. Anger ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Frantic ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "St. Anger ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Dirty Window ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "St. Anger ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> The Unnamed Feeling ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "St. Anger ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Invisible Kid ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "St. Anger ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Purify ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "St. Anger ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Sweet Amber ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "St. Anger ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> All Within My Hands ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "St. Anger ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> My World ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "St. Anger ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> St. Anger ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "St. Anger ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Some Kind Of Monster ({'id': 'St. Anger', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "The Black Album ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> PERFORMED_BY -> Metallica ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "The Black Album ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Don't Tread On Me ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "The Black Album ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Nothing Else Matters ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "The Black Album ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Through The Never ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "The Black Album ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> My Friend Of Misery ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "The Black Album ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Wherever I May Roam ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "The Black Album ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> The Unforgiven ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "The Black Album ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> The Struggle Within ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "The Black Album ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> Of Wolf And Man ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n", | |
| "The Black Album ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None}) -> CONTAINS_SONG -> The God That Failed ({'id': 'The Black Album', 'text': None, 'label': 'ALBUM', 'embedding': None, 'creation_date': None, 'last_modified_date': None, 'file_name': None, 'file_path': None, 'file_size': None, 'file_type': None, 'ref_doc_id': None, 'triplet_source_id': None})\n", | |
| "\n" | |
| ] | |
| } | |
| ], | |
| "source": [ | |
| "for i in range(len(response.source_nodes)):\n", | |
| " print(response.source_nodes[i].node.text)\n", | |
| " print()" | |
| ] | |
| } | |
| ], | |
| "metadata": { | |
| "kernelspec": { | |
| "display_name": "Python 3 (ipykernel)", | |
| "language": "python", | |
| "name": "python3" | |
| }, | |
| "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.14" | |
| } | |
| }, | |
| "nbformat": 4, | |
| "nbformat_minor": 5 | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment