This commit is contained in:
Nathan Evans 2025-09-10 14:21:23 -07:00
parent e1662fb2e1
commit db36524a18

View File

@ -108,7 +108,11 @@
"from graphrag.config.models.embed_graph_config import EmbedGraphConfig\n",
"from graphrag.index.operations.embed_graph.embed_graph import embed_graph\n",
"\n",
"start = time.time()\n",
"pipeline_embeddings = embed_graph(graph, entities, communities, EmbedGraphConfig())\n",
"end = time.time()\n",
"print(\"gee time:\", end - start)\n",
"\n",
"p_umap = run_umap(graph, pipeline_embeddings, lambda x: x)\n",
"\n",
"p_umap_list = [{\"title\": p.label, \"x_gee_p\": p.x, \"y_gee_p\": p.y} for p in p_umap]\n",
@ -129,8 +133,12 @@
"merged_entities = merged_entities.merge(\n",
" p_df, left_on=\"title\", right_on=\"title\", how=\"left\"\n",
")\n",
"community_labels = communities.explode(\"entity_ids\")[[\"community\", \"entity_ids\", \"level\"]]\n",
"merged_entities = merged_entities.merge(community_labels, left_on=\"id\", right_on=\"entity_ids\", how=\"left\")\n",
"community_labels = communities.explode(\"entity_ids\")[\n",
" [\"community\", \"entity_ids\", \"level\"]\n",
"]\n",
"merged_entities = merged_entities.merge(\n",
" community_labels, left_on=\"id\", right_on=\"entity_ids\", how=\"left\"\n",
")\n",
"merged_entities = merged_entities[merged_entities[\"level\"] == 0]"
]
},