examples/notebooks/02-Separate-Extract-and-Build.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "434fea4e",
   "metadata": {},
   "source": [
    "# 02 - Separate Extract and Build"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9fb5cff",
   "metadata": {},
   "source": [
    "## Install toolkit\n",
    "\n",
    "Run the command below to install the graphrag-toolkit. If you've already installed the toolkit, you don't need to install it again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a59f1e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install https://github.com/awslabs/graphrag-toolkit/archive/refs/tags/v1.1.3.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dfff49b",
   "metadata": {},
   "source": [
    "## Extract\n",
    "\n",
    "See [Run the extract and build stages separately](https://github.com/awslabs/graphrag-toolkit/blob/main/docs/indexing.md#run-the-extract-and-build-stages-separately)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79b6c7bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext dotenv\n",
    "%dotenv\n",
    "\n",
    "import os\n",
    "\n",
    "from graphrag_toolkit import LexicalGraphIndex\n",
    "from graphrag_toolkit.storage import GraphStoreFactory\n",
    "from graphrag_toolkit.storage import VectorStoreFactory\n",
    "from graphrag_toolkit.indexing.load import FileBasedChunks\n",
    "from graphrag_toolkit.indexing.build import Checkpoint\n",
    "\n",
    "from llama_index.readers.web import SimpleWebPageReader\n",
    "\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()\n",
    "\n",
    "chunks = FileBasedChunks(\n",
    "    chunks_directory='extracted'\n",
    ")\n",
    "\n",
    "checkpoint = Checkpoint('extraction-checkpoint')\n",
    "\n",
    "graph_store = GraphStoreFactory.for_graph_store(os.environ['GRAPH_STORE'])\n",
    "vector_store = VectorStoreFactory.for_vector_store(os.environ['VECTOR_STORE'])\n",
    "\n",
    "graph_index = LexicalGraphIndex(\n",
    "    graph_store, \n",
    "    vector_store\n",
    ")\n",
    "\n",
    "doc_urls = [\n",
    "    'https://docs.aws.amazon.com/neptune/latest/userguide/intro.html',\n",
    "    'https://docs.aws.amazon.com/neptune-analytics/latest/userguide/what-is-neptune-analytics.html',\n",
    "    'https://docs.aws.amazon.com/neptune-analytics/latest/userguide/neptune-analytics-features.html',\n",
    "    'https://docs.aws.amazon.com/neptune-analytics/latest/userguide/neptune-analytics-vs-neptune-database.html'\n",
    "]\n",
    "\n",
    "docs = SimpleWebPageReader(\n",
    "    html_to_text=True,\n",
    "    metadata_fn=lambda url:{'url': url}\n",
    ").load_data(doc_urls)\n",
    "\n",
    "graph_index.extract(docs, handler=chunks, checkpoint=checkpoint, show_progress=True)\n",
    "\n",
    "collection_id = chunks.collection_id\n",
    "\n",
    "print('Extraction complete')\n",
    "print(f'collection_id: {collection_id}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40c3f5e1",
   "metadata": {},
   "source": [
    "## Build"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaa952bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext dotenv\n",
    "%dotenv\n",
    "\n",
    "import os\n",
    "\n",
    "from graphrag_toolkit import LexicalGraphIndex\n",
    "from graphrag_toolkit.storage import GraphStoreFactory\n",
    "from graphrag_toolkit.storage import VectorStoreFactory\n",
    "from graphrag_toolkit.indexing.load import FileBasedChunks\n",
    "from graphrag_toolkit.indexing.build import Checkpoint\n",
    "\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()\n",
    "\n",
    "chunks = FileBasedChunks(\n",
    "    chunks_directory='extracted',\n",
    "    collection_id=collection_id\n",
    ")\n",
    "checkpoint = Checkpoint('build-checkpoint')\n",
    "\n",
    "graph_store = GraphStoreFactory.for_graph_store(os.environ['GRAPH_STORE'])\n",
    "vector_store = VectorStoreFactory.for_vector_store(os.environ['VECTOR_STORE'])\n",
    "\n",
    "graph_index = LexicalGraphIndex(\n",
    "    graph_store, \n",
    "    vector_store\n",
    ")\n",
    "\n",
    "graph_index.build(chunks, checkpoint=checkpoint, show_progress=True)\n",
    "\n",
    "print('Build complete')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}