mirror of
https://github.com/microsoft/graphrag.git
synced 2025-06-26 23:19:58 +00:00
321d479ab6
* Update API overview * Fix global search example * Fix local search example * Fix global dynamic example * Fix drift example * Update multi-index example * Semver
296 lines
11 KiB
Plaintext
296 lines
11 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Copyright (c) 2024 Microsoft Corporation.\n",
|
|
"# Licensed under the MIT License."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import os\n",
|
|
"\n",
|
|
"import pandas as pd\n",
|
|
"import tiktoken\n",
|
|
"\n",
|
|
"from graphrag.config.enums import ModelType\n",
|
|
"from graphrag.config.models.language_model_config import LanguageModelConfig\n",
|
|
"from graphrag.language_model.manager import ModelManager\n",
|
|
"from graphrag.query.indexer_adapters import (\n",
|
|
" read_indexer_communities,\n",
|
|
" read_indexer_entities,\n",
|
|
" read_indexer_reports,\n",
|
|
")\n",
|
|
"from graphrag.query.structured_search.global_search.community_context import (\n",
|
|
" GlobalCommunityContext,\n",
|
|
")\n",
|
|
"from graphrag.query.structured_search.global_search.search import GlobalSearch"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Global Search example\n",
|
|
"\n",
|
|
"Global search method generates answers by searching over all AI-generated community reports in a map-reduce fashion. This is a resource-intensive method, but often gives good responses for questions that require an understanding of the dataset as a whole (e.g. What are the most significant values of the herbs mentioned in this notebook?)."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### LLM setup"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"api_key = os.environ[\"GRAPHRAG_API_KEY\"]\n",
|
|
"llm_model = os.environ[\"GRAPHRAG_LLM_MODEL\"]\n",
|
|
"\n",
|
|
"config = LanguageModelConfig(\n",
|
|
" api_key=api_key,\n",
|
|
" type=ModelType.OpenAIChat,\n",
|
|
" model=llm_model,\n",
|
|
" max_retries=20,\n",
|
|
")\n",
|
|
"model = ModelManager().get_or_create_chat_model(\n",
|
|
" name=\"global_search\",\n",
|
|
" model_type=ModelType.OpenAIChat,\n",
|
|
" config=config,\n",
|
|
")\n",
|
|
"\n",
|
|
"token_encoder = tiktoken.encoding_for_model(llm_model)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Load community reports as context for global search\n",
|
|
"\n",
|
|
"- Load all community reports in the `community_reports` table from the indexing engine, to be used as context data for global search.\n",
|
|
"- Load entities from the `entities` tables from the indexing engine, to be used for calculating community weights for context ranking. Note that this is optional (if no entities are provided, we will not calculate community weights and only use the rank attribute in the community reports table for context ranking)\n",
|
|
"- Load all communities in the `communities` table from the indexing engine, to be used to reconstruct the community graph hierarchy for dynamic community selection."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# parquet files generated from indexing pipeline\n",
|
|
"INPUT_DIR = \"./inputs/operation dulce\"\n",
|
|
"COMMUNITY_TABLE = \"communities\"\n",
|
|
"COMMUNITY_REPORT_TABLE = \"community_reports\"\n",
|
|
"ENTITY_TABLE = \"entities\"\n",
|
|
"\n",
|
|
"# we don't fix a specific community level but instead use an agent to dynamicially\n",
|
|
"# search through all the community reports to check if they are relevant.\n",
|
|
"COMMUNITY_LEVEL = None"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"community_df = pd.read_parquet(f\"{INPUT_DIR}/{COMMUNITY_TABLE}.parquet\")\n",
|
|
"entity_df = pd.read_parquet(f\"{INPUT_DIR}/{ENTITY_TABLE}.parquet\")\n",
|
|
"report_df = pd.read_parquet(f\"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet\")\n",
|
|
"\n",
|
|
"communities = read_indexer_communities(community_df, report_df)\n",
|
|
"reports = read_indexer_reports(\n",
|
|
" report_df,\n",
|
|
" community_df,\n",
|
|
" community_level=COMMUNITY_LEVEL,\n",
|
|
" dynamic_community_selection=True,\n",
|
|
")\n",
|
|
"entities = read_indexer_entities(\n",
|
|
" entity_df, community_df, community_level=COMMUNITY_LEVEL\n",
|
|
")\n",
|
|
"\n",
|
|
"print(f\"Total report count: {len(report_df)}\")\n",
|
|
"print(\n",
|
|
" f\"Report count after filtering by community level {COMMUNITY_LEVEL}: {len(reports)}\"\n",
|
|
")\n",
|
|
"\n",
|
|
"report_df.head()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"#### Build global context with dynamic community selection\n",
|
|
"\n",
|
|
"The goal of dynamic community selection reduce the number of community reports that need to be processed in the map-reduce operation. To that end, we take advantage of the hierachical structure of the indexed dataset. We first ask the LLM to rate how relevant each level 0 community is with respect to the user query, we then traverse down the child node(s) if the current community report is deemed relevant.\n",
|
|
"\n",
|
|
"You can still set a `COMMUNITY_LEVEL` to filter out lower level community reports and apply dynamic community selection on the filtered reports.\n",
|
|
"\n",
|
|
"Note that the dataset is quite small, with only consist of 20 communities from 2 levels (level 0 and 1). Dynamic community selection is more effective when there are large amount of content to be filtered out."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"context_builder = GlobalCommunityContext(\n",
|
|
" community_reports=reports,\n",
|
|
" communities=communities,\n",
|
|
" entities=entities, # default to None if you don't want to use community weights for ranking\n",
|
|
" token_encoder=token_encoder,\n",
|
|
" dynamic_community_selection=True,\n",
|
|
" dynamic_community_selection_kwargs={\n",
|
|
" \"model\": model,\n",
|
|
" \"token_encoder\": token_encoder,\n",
|
|
" },\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"#### Perform global search with dynamic community selection"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"context_builder_params = {\n",
|
|
" \"use_community_summary\": False, # False means using full community reports. True means using community short summaries.\n",
|
|
" \"shuffle_data\": True,\n",
|
|
" \"include_community_rank\": True,\n",
|
|
" \"min_community_rank\": 0,\n",
|
|
" \"community_rank_name\": \"rank\",\n",
|
|
" \"include_community_weight\": True,\n",
|
|
" \"community_weight_name\": \"occurrence weight\",\n",
|
|
" \"normalize_community_weight\": True,\n",
|
|
" \"max_tokens\": 12_000, # change this based on the token limit you have on your model (if you are using a model with 8k limit, a good setting could be 5000)\n",
|
|
" \"context_name\": \"Reports\",\n",
|
|
"}\n",
|
|
"\n",
|
|
"map_llm_params = {\n",
|
|
" \"max_tokens\": 1000,\n",
|
|
" \"temperature\": 0.0,\n",
|
|
" \"response_format\": {\"type\": \"json_object\"},\n",
|
|
"}\n",
|
|
"\n",
|
|
"reduce_llm_params = {\n",
|
|
" \"max_tokens\": 2000, # change this based on the token limit you have on your model (if you are using a model with 8k limit, a good setting could be 1000-1500)\n",
|
|
" \"temperature\": 0.0,\n",
|
|
"}"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"search_engine = GlobalSearch(\n",
|
|
" model=model,\n",
|
|
" context_builder=context_builder,\n",
|
|
" token_encoder=token_encoder,\n",
|
|
" max_data_tokens=12_000, # change this based on the token limit you have on your model (if you are using a model with 8k limit, a good setting could be 5000)\n",
|
|
" map_llm_params=map_llm_params,\n",
|
|
" reduce_llm_params=reduce_llm_params,\n",
|
|
" allow_general_knowledge=False, # set this to True will add instruction to encourage the LLM to incorporate general knowledge in the response, which may increase hallucinations, but could be useful in some use cases.\n",
|
|
" json_mode=True, # set this to False if your LLM model does not support JSON mode.\n",
|
|
" context_builder_params=context_builder_params,\n",
|
|
" concurrent_coroutines=32,\n",
|
|
" response_type=\"multiple paragraphs\", # free form text describing the response type and format, can be anything, e.g. prioritized list, single paragraph, multiple paragraphs, multiple-page report\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"result = await search_engine.search(\"What is operation dulce?\")\n",
|
|
"\n",
|
|
"print(result.response)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# inspect the data used to build the context for the LLM responses\n",
|
|
"result.context_data[\"reports\"]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# inspect number of LLM calls and tokens in dynamic community selection\n",
|
|
"llm_calls = result.llm_calls_categories[\"build_context\"]\n",
|
|
"prompt_tokens = result.prompt_tokens_categories[\"build_context\"]\n",
|
|
"output_tokens = result.output_tokens_categories[\"build_context\"]\n",
|
|
"print(\n",
|
|
" f\"Build context ({llm_model})\\nLLM calls: {llm_calls}. Prompt tokens: {prompt_tokens}. Output tokens: {output_tokens}.\"\n",
|
|
")\n",
|
|
"# inspect number of LLM calls and tokens in map-reduce\n",
|
|
"llm_calls = result.llm_calls_categories[\"map\"] + result.llm_calls_categories[\"reduce\"]\n",
|
|
"prompt_tokens = (\n",
|
|
" result.prompt_tokens_categories[\"map\"] + result.prompt_tokens_categories[\"reduce\"]\n",
|
|
")\n",
|
|
"output_tokens = (\n",
|
|
" result.output_tokens_categories[\"map\"] + result.output_tokens_categories[\"reduce\"]\n",
|
|
")\n",
|
|
"print(\n",
|
|
" f\"Map-reduce ({llm_model})\\nLLM calls: {llm_calls}. Prompt tokens: {prompt_tokens}. Output tokens: {output_tokens}.\"\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
|
|
}
|