microsoft--graphrag/docs/examples_notebooks/index_migration_to_v3.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) 2024 Microsoft Corporation.\n",
    "# Licensed under the MIT License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Index Migration (v2 to v3)\n",
    "\n",
    "This notebook is used to maintain data model parity with older indexes for version 3.0 of GraphRAG. If you have a pre-3.0 index and need to migrate without re-running the entire pipeline, you can use this notebook to only update the pieces necessary for alignment. If you have a pre-2.0 index, please run the v2 migration notebook first!\n",
    "\n",
    "NOTE: we recommend regenerating your settings.yml with the latest version of GraphRAG using `graphrag init`. Copy your LLM settings into it before running this notebook. This ensures your config is aligned with the latest version for the migration. The config changes from v2 to v3 are significant in places!\n",
    "\n",
    "WARNING: This will overwrite your parquet files, you may want to make a backup!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is the directory that has your settings.yaml\n",
    "PROJECT_DIRECTORY = \"<your project directory>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "from graphrag.config.models.graph_rag_config import GraphRagConfig\n",
    "from graphrag_common.config import load_config\n",
    "from graphrag_storage.storage_factory import create_storage\n",
    "\n",
    "config = load_config(GraphRagConfig, config_path=Path(PROJECT_DIRECTORY))\n",
    "storage = create_storage(config.output_storage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def remove_columns(df, columns):\n",
    "    \"\"\"Remove columns from a DataFrame, suppressing errors.\"\"\"\n",
    "    df.drop(labels=columns, axis=1, errors=\"ignore\", inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from graphrag_storage.tables.parquet_table_provider import ParquetTableProvider\n",
    "\n",
    "# Create table provider from storage\n",
    "table_provider = ParquetTableProvider(storage)\n",
    "\n",
    "text_units = await table_provider.read_dataframe(\"text_units\")\n",
    "\n",
    "text_units[\"document_id\"] = text_units[\"document_ids\"].apply(lambda ids: ids[0])\n",
    "remove_columns(text_units, [\"document_ids\"])\n",
    "\n",
    "await table_provider.write_dataframe(\"text_units\", text_units)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Update settings.yaml\n",
    "If you have left the default settings for your vector store schema, you may need to set explicit values that map each embedding type to a vector schema name. If you have already customized your vector store schema it may not be necessary.\n",
    "\n",
    "Old default index names:\n",
    "- default-text_unit-text\n",
    "- default-entity-description\n",
    "- default-community-full_content\n",
    "\n",
    "(if you left all of the defaults, check your output/lancedb folder to confirm the above)\n",
    "\n",
    "v3 versions are:\n",
    "- text_unit_text\n",
    "- entity_description\n",
    "- community_full_content\n",
    "\n",
    "Therefore, with a v2 index need to explicitly set the old index names so it connects correctly. We no longer support the \"prefix\" - you can just set an explicit index_name for each embedding.\n",
    "\n",
    "NOTE: we are also setting the default vector_size for each index below, under the assumption that you are using a prior default with 1536 dimensions. Our new default of text-embedding-3-large has 3072 dimensions, which will be populated as the default if unset. Again, if you have a more complicated situation you may want to manually configure this.\n",
    "\n",
    "Here is an example of the new vector store config block that you may need in your settings.yaml:\n",
    "\n",
    "```yaml\n",
    "vector_store:\n",
    "  type: lancedb\n",
    "  db_uri: output/lancedb\n",
    "  index_schema:\n",
    "    text_unit_text:\n",
    "      index_name: default-text_unit-text\n",
    "      vector_size: 1536\n",
    "    entity_description:\n",
    "      index_name: default-entity-description\n",
    "      vector_size: 1536\n",
    "    community_full_content:\n",
    "      index_name: default-community-full_content\n",
    "      vector_size: 1536\n",
    "```\n"
   ]
  }
 ],
 "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}