Skip to content

Commit c030acf

Browse files
natoversenatoverse
committed
Create v1-to-v2 migration notebook
1 parent fce1a13 commit c030acf

File tree

2 files changed

+165
-18
lines changed

2 files changed

+165
-18
lines changed

docs/examples_notebooks/index_migration.ipynb renamed to docs/examples_notebooks/index_migration_to_v1.ipynb

Lines changed: 25 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
"cells": [
33
{
44
"cell_type": "code",
5-
"execution_count": 66,
5+
"execution_count": 62,
66
"metadata": {},
77
"outputs": [],
88
"source": [
@@ -14,9 +14,9 @@
1414
"cell_type": "markdown",
1515
"metadata": {},
1616
"source": [
17-
"## Index Migration\n",
17+
"## Index Migration (pre-v1 to v1)\n",
1818
"\n",
19-
"This notebook is used to maintain data model parity with older indexes for the latest versions of GraphRAG. If you have a pre-1.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.\n",
19+
"This notebook is used to maintain data model parity with older indexes for version 1.0 of GraphRAG. If you have a pre-1.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.\n",
2020
"\n",
2121
"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. This also ensures that you have default vector store config, which is now required or indexing will fail.\n",
2222
"\n",
@@ -25,40 +25,38 @@
2525
},
2626
{
2727
"cell_type": "code",
28-
"execution_count": 67,
28+
"execution_count": 63,
2929
"metadata": {},
3030
"outputs": [],
3131
"source": [
32-
"# This is the directory that has your settings.yml\n",
32+
"# This is the directory that has your settings.yaml\n",
3333
"# NOTE: much older indexes may have been output with a timestamped directory\n",
34-
"# if this is the case, you will need to make sure the storage.base_dir in settings.yml points to it correctly\n",
35-
"PROJECT_DIRECTORY = \"<your project directory>\""
34+
"# if this is the case, you will need to make sure the storage.base_dir in settings.yaml points to it correctly\n",
35+
"PROJECT_DIRECTORY = \"<your project directory\""
3636
]
3737
},
3838
{
3939
"cell_type": "code",
40-
"execution_count": null,
40+
"execution_count": 64,
4141
"metadata": {},
4242
"outputs": [],
4343
"source": [
4444
"from pathlib import Path\n",
4545
"\n",
4646
"from graphrag.config.load_config import load_config\n",
47-
"from graphrag.config.resolve_path import resolve_paths\n",
4847
"from graphrag.storage.factory import StorageFactory\n",
4948
"\n",
50-
"# This first block does some config loading, path resolution, and translation that is normally done by the CLI/API when running a full workflow\n",
5149
"config = load_config(Path(PROJECT_DIRECTORY))\n",
52-
"resolve_paths(config)\n",
53-
"storage_config = config.storage.model_dump() # type: ignore\n",
50+
"storage_config = config.output.model_dump()\n",
5451
"storage = StorageFactory().create_storage(\n",
55-
" storage_type=storage_config[\"type\"], kwargs=storage_config\n",
52+
" storage_type=storage_config[\"type\"],\n",
53+
" kwargs=storage_config,\n",
5654
")"
5755
]
5856
},
5957
{
6058
"cell_type": "code",
61-
"execution_count": 69,
59+
"execution_count": 65,
6260
"metadata": {},
6361
"outputs": [],
6462
"source": [
@@ -69,7 +67,7 @@
6967
},
7068
{
7169
"cell_type": "code",
72-
"execution_count": 70,
70+
"execution_count": 66,
7371
"metadata": {},
7472
"outputs": [],
7573
"source": [
@@ -98,7 +96,7 @@
9896
},
9997
{
10098
"cell_type": "code",
101-
"execution_count": 71,
99+
"execution_count": 67,
102100
"metadata": {},
103101
"outputs": [],
104102
"source": [
@@ -134,7 +132,7 @@
134132
"if \"name\" in final_entities.columns:\n",
135133
" final_entities.rename(columns={\"name\": \"title\"}, inplace=True)\n",
136134
"remove_columns(\n",
137-
" final_entities, [\"mname_embedding\", \"graph_embedding\", \"description_embedding\"]\n",
135+
" final_entities, [\"name_embedding\", \"graph_embedding\", \"description_embedding\"]\n",
138136
")\n",
139137
"\n",
140138
"# Final nodes uses community for joins, which is now an int everywhere\n",
@@ -168,6 +166,15 @@
168166
" final_communities[\"id\"] = [str(uuid4()) for _ in range(len(final_communities))]\n",
169167
"if \"parent\" not in final_communities.columns:\n",
170168
" final_communities = final_communities.merge(parent_df, on=\"community\", how=\"left\")\n",
169+
"if \"entity_ids\" not in final_communities.columns:\n",
170+
" node_mapping = (\n",
171+
" final_nodes.loc[:, [\"community\", \"id\"]]\n",
172+
" .groupby(\"community\")\n",
173+
" .agg(entity_ids=(\"id\", list))\n",
174+
" )\n",
175+
" final_communities = final_communities.merge(\n",
176+
" node_mapping, on=\"community\", how=\"left\"\n",
177+
" )\n",
171178
"remove_columns(final_communities, [\"raw_community\"])\n",
172179
"\n",
173180
"# We need int for community and the human_readable_id copy for consistency\n",
@@ -197,7 +204,7 @@
197204
"source": [
198205
"from graphrag.cache.factory import CacheFactory\n",
199206
"from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks\n",
200-
"from graphrag.index.config.embeddings import get_embedded_fields, get_embedding_settings\n",
207+
"from graphrag.config.embeddings import get_embedded_fields, get_embedding_settings\n",
201208
"from graphrag.index.flows.generate_text_embeddings import generate_text_embeddings\n",
202209
"\n",
203210
"# We only need to re-run the embeddings workflow, to ensure that embeddings for all required search fields are in place\n",

‎docs/examples_notebooks/index_migration_to_v2.ipynb‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": 41,
6+
"metadata": {},
7+
"outputs": [],
8+
"source": [
9+
"# Copyright (c) 2024 Microsoft Corporation.\n",
10+
"# Licensed under the MIT License."
11+
]
12+
},
13+
{
14+
"cell_type": "markdown",
15+
"metadata": {},
16+
"source": [
17+
"## Index Migration (v1 to v2)\n",
18+
"\n",
19+
"This notebook is used to maintain data model parity with older indexes for version 2.0 of GraphRAG. If you have a pre-2.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-1.0 index, please run the v1 migration notebook first!\n",
20+
"\n",
21+
"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. This also ensures that you have default vector store config, which is now required or indexing will fail.\n",
22+
"\n",
23+
"WARNING: This will overwrite your parquet files, you may want to make a backup!"
24+
]
25+
},
26+
{
27+
"cell_type": "code",
28+
"execution_count": 42,
29+
"metadata": {},
30+
"outputs": [],
31+
"source": [
32+
"# This is the directory that has your settings.yaml\n",
33+
"PROJECT_DIRECTORY = \"<your project directory\""
34+
]
35+
},
36+
{
37+
"cell_type": "code",
38+
"execution_count": 43,
39+
"metadata": {},
40+
"outputs": [],
41+
"source": [
42+
"from pathlib import Path\n",
43+
"\n",
44+
"from graphrag.config.load_config import load_config\n",
45+
"from graphrag.storage.factory import StorageFactory\n",
46+
"\n",
47+
"config = load_config(Path(PROJECT_DIRECTORY))\n",
48+
"storage_config = config.output.model_dump()\n",
49+
"storage = StorageFactory().create_storage(\n",
50+
" storage_type=storage_config[\"type\"],\n",
51+
" kwargs=storage_config,\n",
52+
")"
53+
]
54+
},
55+
{
56+
"cell_type": "code",
57+
"execution_count": 44,
58+
"metadata": {},
59+
"outputs": [],
60+
"source": [
61+
"def remove_columns(df, columns):\n",
62+
" \"\"\"Remove columns from a DataFrame, suppressing errors.\"\"\"\n",
63+
" df.drop(labels=columns, axis=1, errors=\"ignore\", inplace=True)"
64+
]
65+
},
66+
{
67+
"cell_type": "code",
68+
"execution_count": 45,
69+
"metadata": {},
70+
"outputs": [],
71+
"source": [
72+
"from graphrag.utils.storage import (\n",
73+
" delete_table_from_storage,\n",
74+
" load_table_from_storage,\n",
75+
" write_table_to_storage,\n",
76+
")\n",
77+
"\n",
78+
"final_documents = await load_table_from_storage(\"create_final_documents\", storage)\n",
79+
"final_text_units = await load_table_from_storage(\"create_final_text_units\", storage)\n",
80+
"final_entities = await load_table_from_storage(\"create_final_entities\", storage)\n",
81+
"final_nodes = await load_table_from_storage(\"create_final_nodes\", storage)\n",
82+
"final_relationships = await load_table_from_storage(\n",
83+
" \"create_final_relationships\", storage\n",
84+
")\n",
85+
"final_communities = await load_table_from_storage(\"create_final_communities\", storage)\n",
86+
"final_community_reports = await load_table_from_storage(\n",
87+
" \"create_final_community_reports\", storage\n",
88+
")\n",
89+
"\n",
90+
"# we've renamed document attributes as metadata\n",
91+
"if \"attributes\" in final_documents.columns:\n",
92+
" final_documents.rename(columns={\"attributes\": \"metadata\"}, inplace=True)\n",
93+
"\n",
94+
"# we're removing the nodes table, so we need to copy the graph columns into entities\n",
95+
"graph_props = (\n",
96+
" final_nodes.loc[:, [\"id\", \"degree\", \"x\", \"y\"]].groupby(\"id\").first().reset_index()\n",
97+
")\n",
98+
"final_entities = final_entities.merge(graph_props, on=\"id\", how=\"left\")\n",
99+
"\n",
100+
"# we renamed all the output files for better clarity now that we don't have workflow naming constraints from DataShaper\n",
101+
"await write_table_to_storage(final_documents, \"documents\", storage)\n",
102+
"await write_table_to_storage(final_text_units, \"text_units\", storage)\n",
103+
"await write_table_to_storage(final_entities, \"entities\", storage)\n",
104+
"await write_table_to_storage(final_relationships, \"relationships\", storage)\n",
105+
"await write_table_to_storage(final_communities, \"communities\", storage)\n",
106+
"await write_table_to_storage(final_community_reports, \"community_reports\", storage)\n",
107+
"\n",
108+
"# delete all the old versions\n",
109+
"await delete_table_from_storage(\"create_final_documents\", storage)\n",
110+
"await delete_table_from_storage(\"create_final_text_units\", storage)\n",
111+
"await delete_table_from_storage(\"create_final_entities\", storage)\n",
112+
"await delete_table_from_storage(\"create_final_nodes\", storage)\n",
113+
"await delete_table_from_storage(\"create_final_relationships\", storage)\n",
114+
"await delete_table_from_storage(\"create_final_communities\", storage)\n",
115+
"await delete_table_from_storage(\"create_final_community_reports\", storage)"
116+
]
117+
}
118+
],
119+
"metadata": {
120+
"kernelspec": {
121+
"display_name": ".venv",
122+
"language": "python",
123+
"name": "python3"
124+
},
125+
"language_info": {
126+
"codemirror_mode": {
127+
"name": "ipython",
128+
"version": 3
129+
},
130+
"file_extension": ".py",
131+
"mimetype": "text/x-python",
132+
"name": "python",
133+
"nbconvert_exporter": "python",
134+
"pygments_lexer": "ipython3",
135+
"version": "3.11.9"
136+
}
137+
},
138+
"nbformat": 4,
139+
"nbformat_minor": 2
140+
}

0 commit comments

Comments
 (0)