Prompt Tuning ⚙️

GraphRAG provides the ability to create domain adapted prompts for the generation of the knowledge graph. This step is optional, though it is highly encouraged to run it as it will yield better results when executing an Index Run.

These are generated by loading the inputs, splitting them into chunks (text units) and then running a series of LLM invocations and template substitutions to generate the final prompts. We suggest using the default values provided by the script, but in this page you'll find the detail of each in case you want to further explore and tweak the prompt tuning algorithm.

Figure 1: Auto Tuning Conceptual Diagram.

Prerequisites

Before running auto tuning make sure you have already initialized your workspace with the graphrag.index --init command. This will create the necessary configuration files and the default prompts. Refer to the Init Documentation for more information about the initialization process.

Usage

You can run the main script from the command line with various options:

python -m graphrag.prompt_tune [--root ROOT] [--domain DOMAIN]  [--method METHOD] [--limit LIMIT] [--language LANGUAGE] \
[--max-tokens MAX_TOKENS] [--chunk-size CHUNK_SIZE] [--n-subset-max N_SUBSET_MAX] [--k K] \
[--min-examples-required MIN_EXAMPLES_REQUIRED] [--no-entity-types] [--output OUTPUT]

Command-Line Options

• --config (required): The path to the configuration file. This is required to load the data and model settings.

• --root (optional): The data project root directory, including the config files (YML, JSON, or .env). Defaults to the current directory.

• --domain (optional): The domain related to your input data, such as 'space science', 'microbiology', or 'environmental news'. If left empty, the domain will be inferred from the input data.

• --method (optional): The method to select documents. Options are all, random, auto or top. Default is random.

• --limit (optional): The limit of text units to load when using random or top selection. Default is 15.

• --language (optional): The language to use for input processing. If it is different from the inputs' language, the LLM will translate. Default is "" meaning it will be automatically detected from the inputs.

• --max-tokens (optional): Maximum token count for prompt generation. Default is 2000.

• --chunk-size (optional): The size in tokens to use for generating text units from input documents. Default is 200.

• --n-subset-max (optional): The number of text chunks to embed when using auto selection method. Default is 300.

• --k (optional): The number of documents to select when using auto selection method. Default is 15.

• --min-examples-required (optional): The minimum number of examples required for entity extraction prompts. Default is 2.

• --no-entity-types (optional): Use untyped entity extraction generation. We recommend using this when your data covers a lot of topics or it is highly randomized.

• --output (optional): The folder to save the generated prompts. Default is "prompts".

Example Usage

python -m graphrag.prompt_tune --root /path/to/project --config /path/to/settings.yaml --domain "environmental news" \
--method random --limit 10 --language English --max-tokens 2048 --chunk-size 256 --min-examples-required 3 \
--no-entity-types --output /path/to/output

or, with minimal configuration (suggested):

python -m graphrag.prompt_tune --root /path/to/project --config /path/to/settings.yaml --no-entity-types

Document Selection Methods

The auto tuning feature ingests the input data and then divides it into text units the size of the chunk size parameter. After that, it uses one of the following selection methods to pick a sample to work with for prompt generation:

• random: Select text units randomly. This is the default and recommended option.
• top: Select the head n text units.
• all: Use all text units for the generation. Use only with small datasets; this option is not usually recommended.
• auto: Embed text units in a lower-dimensional space and select the k nearest neighbors to the centroid. This is useful when you have a large dataset and want to select a representative sample.

Modify Env Vars

After running auto tuning, you should modify the following environment variables (or config variables) to pick up the new prompts on your index run. Note: Please make sure to update the correct path to the generated prompts, in this example we are using the default "prompts" path.

• GRAPHRAG_ENTITY_EXTRACTION_PROMPT_FILE = "prompts/entity_extraction.txt"

• GRAPHRAG_COMMUNITY_REPORT_PROMPT_FILE = "prompts/community_report.txt"

• GRAPHRAG_SUMMARIZE_DESCRIPTIONS_PROMPT_FILE = "prompts/summarize_descriptions.txt"

or in your yaml config file:

entity_extraction:
  prompt: "prompts/entity_extraction.txt"

summarize_descriptions:
  prompt: "prompts/summarize_descriptions.txt"

community_reports:
  prompt: "prompts/community_report.txt"