mirror of
https://github.com/microsoft/graphrag.git
synced 2025-06-26 23:19:58 +00:00
c02ab0984a
* Remove create_final_nodes * Rename final entity output to "entities" * Remove duplicate code from graph extraction * Rename create_final_relationships output to "relationships" * Rename create_final_communities output to "communities" * Combine compute_communities and create_final_communities * Rename create_final_covariates output to "covariates" * Rename create_final_community_reports output to "community_reports" * Rename create_final_text_units output to "text_units" * Rename create_final_documents output to "documents" * Remove transient snapshots config * Move create_final_entities to finalize_entities operation * Move create_final_relationships flow to finalize_relationships operation * Reuse some community report functions * Collapse most of graph and text unit-based report generation * Unify schemas files * Move community reports extractor * Move NLP report prompt to prompts folder * Fix a few pandas warnings * Rename embeddings config to embed_text * Rename claim_extraction config to extract_claims * Remove nltk from standard graph extraction * Fix verb tests * Fix extract graph config naming * Fix moved file reference * Create v1-to-v2 migration notebook * Semver * Fix smoke test artifact count * Raise tpm/rpm on smoke tests * Update drift settings for smoke tests * Reuse project directory var in api notebook * Format * Format
262 lines
11 KiB
Plaintext
262 lines
11 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 62,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Copyright (c) 2024 Microsoft Corporation.\n",
|
|
"# Licensed under the MIT License."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Index Migration (pre-v1 to v1)\n",
|
|
"\n",
|
|
"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",
|
|
"\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. This also ensures that you have default vector store config, which is now required or indexing will fail.\n",
|
|
"\n",
|
|
"WARNING: This will overwrite your parquet files, you may want to make a backup!"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 63,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# This is the directory that has your settings.yaml\n",
|
|
"# NOTE: much older indexes may have been output with a timestamped directory\n",
|
|
"# if this is the case, you will need to make sure the storage.base_dir in settings.yaml points to it correctly\n",
|
|
"PROJECT_DIRECTORY = \"<your project directory\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 64,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from pathlib import Path\n",
|
|
"\n",
|
|
"from graphrag.config.load_config import load_config\n",
|
|
"from graphrag.storage.factory import StorageFactory\n",
|
|
"\n",
|
|
"config = load_config(Path(PROJECT_DIRECTORY))\n",
|
|
"storage_config = config.output.model_dump()\n",
|
|
"storage = StorageFactory().create_storage(\n",
|
|
" storage_type=storage_config[\"type\"],\n",
|
|
" kwargs=storage_config,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 65,
|
|
"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": 66,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def get_community_parent(nodes):\n",
|
|
" \"\"\"Compute the parent community using the node membership as a lookup.\"\"\"\n",
|
|
" parent_mapping = nodes.loc[:, [\"level\", \"community\", \"title\"]]\n",
|
|
" nodes = nodes.loc[:, [\"level\", \"community\", \"title\"]]\n",
|
|
"\n",
|
|
" # Create a parent mapping by adding 1 to the level column\n",
|
|
" parent_mapping[\"level\"] += 1 # Shift levels for parent relationship\n",
|
|
" parent_mapping.rename(columns={\"community\": \"parent\"}, inplace=True)\n",
|
|
"\n",
|
|
" # Merge the parent information back into the base DataFrame\n",
|
|
" nodes = nodes.merge(parent_mapping, on=[\"level\", \"title\"], how=\"left\")\n",
|
|
"\n",
|
|
" # Fill missing parents with -1 (default value)\n",
|
|
" nodes[\"parent\"] = nodes[\"parent\"].fillna(-1).astype(int)\n",
|
|
"\n",
|
|
" join = (\n",
|
|
" nodes.groupby([\"community\", \"level\", \"parent\"])\n",
|
|
" .agg({\"title\": list})\n",
|
|
" .reset_index()\n",
|
|
" )\n",
|
|
" return join[join[\"community\"] > -1].loc[:, [\"community\", \"parent\"]]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 67,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from uuid import uuid4\n",
|
|
"\n",
|
|
"from graphrag.utils.storage import load_table_from_storage, write_table_to_storage\n",
|
|
"\n",
|
|
"# First we'll go through any parquet files that had model changes and update them\n",
|
|
"# The new data model may have removed excess columns as well, but we will only make the minimal changes required for compatibility\n",
|
|
"\n",
|
|
"final_documents = await load_table_from_storage(\"create_final_documents\", storage)\n",
|
|
"final_text_units = await load_table_from_storage(\"create_final_text_units\", storage)\n",
|
|
"final_entities = await load_table_from_storage(\"create_final_entities\", storage)\n",
|
|
"final_nodes = await load_table_from_storage(\"create_final_nodes\", storage)\n",
|
|
"final_relationships = await load_table_from_storage(\n",
|
|
" \"create_final_relationships\", storage\n",
|
|
")\n",
|
|
"final_communities = await load_table_from_storage(\"create_final_communities\", storage)\n",
|
|
"final_community_reports = await load_table_from_storage(\n",
|
|
" \"create_final_community_reports\", storage\n",
|
|
")\n",
|
|
"\n",
|
|
"\n",
|
|
"# Documents renames raw_content for consistency\n",
|
|
"if \"raw_content\" in final_documents.columns:\n",
|
|
" final_documents.rename(columns={\"raw_content\": \"text\"}, inplace=True)\n",
|
|
"final_documents[\"human_readable_id\"] = final_documents.index + 1\n",
|
|
"\n",
|
|
"# Text units just get a human_readable_id or consistency\n",
|
|
"final_text_units[\"human_readable_id\"] = final_text_units.index + 1\n",
|
|
"\n",
|
|
"# We renamed \"name\" to \"title\" for consistency with the rest of the tables\n",
|
|
"if \"name\" in final_entities.columns:\n",
|
|
" final_entities.rename(columns={\"name\": \"title\"}, inplace=True)\n",
|
|
"remove_columns(\n",
|
|
" final_entities, [\"name_embedding\", \"graph_embedding\", \"description_embedding\"]\n",
|
|
")\n",
|
|
"\n",
|
|
"# Final nodes uses community for joins, which is now an int everywhere\n",
|
|
"final_nodes[\"community\"] = final_nodes[\"community\"].fillna(-1)\n",
|
|
"final_nodes[\"community\"] = final_nodes[\"community\"].astype(int)\n",
|
|
"remove_columns(\n",
|
|
" final_nodes,\n",
|
|
" [\n",
|
|
" \"type\",\n",
|
|
" \"description\",\n",
|
|
" \"source_id\",\n",
|
|
" \"graph_embedding\",\n",
|
|
" \"entity_type\",\n",
|
|
" \"top_level_node_id\",\n",
|
|
" \"size\",\n",
|
|
" ],\n",
|
|
")\n",
|
|
"\n",
|
|
"# Relationships renames \"rank\" to \"combined_degree\" to be clear what the default ranking is\n",
|
|
"if \"rank\" in final_relationships.columns:\n",
|
|
" final_relationships.rename(columns={\"rank\": \"combined_degree\"}, inplace=True)\n",
|
|
"\n",
|
|
"\n",
|
|
"# Compute the parents for each community, to add to communities and reports\n",
|
|
"parent_df = get_community_parent(final_nodes)\n",
|
|
"\n",
|
|
"# Communities previously used the \"id\" field for the Leiden id, but we've moved this to the community field and use a uuid for id like the others\n",
|
|
"if \"community\" not in final_communities.columns:\n",
|
|
" final_communities[\"community\"] = final_communities[\"id\"].astype(int)\n",
|
|
" final_communities[\"human_readable_id\"] = final_communities[\"community\"]\n",
|
|
" final_communities[\"id\"] = [str(uuid4()) for _ in range(len(final_communities))]\n",
|
|
"if \"parent\" not in final_communities.columns:\n",
|
|
" final_communities = final_communities.merge(parent_df, on=\"community\", how=\"left\")\n",
|
|
"if \"entity_ids\" not in final_communities.columns:\n",
|
|
" node_mapping = (\n",
|
|
" final_nodes.loc[:, [\"community\", \"id\"]]\n",
|
|
" .groupby(\"community\")\n",
|
|
" .agg(entity_ids=(\"id\", list))\n",
|
|
" )\n",
|
|
" final_communities = final_communities.merge(\n",
|
|
" node_mapping, on=\"community\", how=\"left\"\n",
|
|
" )\n",
|
|
"remove_columns(final_communities, [\"raw_community\"])\n",
|
|
"\n",
|
|
"# We need int for community and the human_readable_id copy for consistency\n",
|
|
"final_community_reports[\"community\"] = final_community_reports[\"community\"].astype(int)\n",
|
|
"final_community_reports[\"human_readable_id\"] = final_community_reports[\"community\"]\n",
|
|
"if \"parent\" not in final_community_reports.columns:\n",
|
|
" final_community_reports = final_community_reports.merge(\n",
|
|
" parent_df, on=\"community\", how=\"left\"\n",
|
|
" )\n",
|
|
"\n",
|
|
"await write_table_to_storage(final_documents, \"create_final_documents\", storage)\n",
|
|
"await write_table_to_storage(final_text_units, \"create_final_text_units\", storage)\n",
|
|
"await write_table_to_storage(final_entities, \"create_final_entities\", storage)\n",
|
|
"await write_table_to_storage(final_nodes, \"create_final_nodes\", storage)\n",
|
|
"await write_table_to_storage(final_relationships, \"create_final_relationships\", storage)\n",
|
|
"await write_table_to_storage(final_communities, \"create_final_communities\", storage)\n",
|
|
"await write_table_to_storage(\n",
|
|
" final_community_reports, \"create_final_community_reports\", storage\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from graphrag.cache.factory import CacheFactory\n",
|
|
"from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks\n",
|
|
"from graphrag.config.embeddings import get_embedded_fields, get_embedding_settings\n",
|
|
"from graphrag.index.flows.generate_text_embeddings import generate_text_embeddings\n",
|
|
"\n",
|
|
"# We only need to re-run the embeddings workflow, to ensure that embeddings for all required search fields are in place\n",
|
|
"# We'll construct the context and run this function flow directly to avoid everything else\n",
|
|
"\n",
|
|
"\n",
|
|
"embedded_fields = get_embedded_fields(config)\n",
|
|
"text_embed = get_embedding_settings(config)\n",
|
|
"callbacks = NoopWorkflowCallbacks()\n",
|
|
"cache_config = config.cache.model_dump() # type: ignore\n",
|
|
"cache = CacheFactory().create_cache(\n",
|
|
" cache_type=cache_config[\"type\"], # type: ignore\n",
|
|
" root_dir=PROJECT_DIRECTORY,\n",
|
|
" kwargs=cache_config,\n",
|
|
")\n",
|
|
"\n",
|
|
"await generate_text_embeddings(\n",
|
|
" final_documents=None,\n",
|
|
" final_relationships=None,\n",
|
|
" final_text_units=final_text_units,\n",
|
|
" final_entities=final_entities,\n",
|
|
" final_community_reports=final_community_reports,\n",
|
|
" callbacks=callbacks,\n",
|
|
" cache=cache,\n",
|
|
" storage=storage,\n",
|
|
" text_embed_config=text_embed,\n",
|
|
" embedded_fields=embedded_fields,\n",
|
|
" snapshot_embeddings_enabled=False,\n",
|
|
")"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": ".venv",
|
|
"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.11.9"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 2
|
|
}
|