Prompt Tuning ⚙️

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

The templates 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 template generation algorithm.

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] [--max-tokens MAX_TOKENS] [--chunk-size CHUNK_SIZE] [--no-entity-types] [--output OUTPUT]

Command-Line Options

• --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, or top. Default is random.

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

• --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.

• --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 --domain "environmental news" --method random --limit 10 --max-tokens 2048 --chunk-size 256 --no-entity-types --output /path/to/output

or, with minimal configuration (suggested):

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

Document Selection Methods

The auto template 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 template 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.

Modify Env Vars

After running auto-templating, 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"