E²GraphRAG

To set up your environment for E²GraphRAG, install the required Python packages using the provided requirements.txt file.

pip install -r requirements.txt

To use the InfiniteBench dataset, you need two specific files: longbook_choice_eng.jsonl and longbook_qa_eng.jsonl.

You can obtain them in two ways:

1. Manual Download: Download the files from the InfiniteBench Hugging Face repository and place them in the ./data/InfiniteBench directory.
2. Automated Script: Use the provided download.sh script located in the repository.

The E²GraphRAG pipeline automates document tree construction, graph structure extraction, and answer generation. The entry point is main.py.

To run the pipeline, you must provide a YAML configuration file that defines the key parameters for the process.

Example configuration path: ./configs/example_config.yaml

python main.py --config <path_to_config_file>

The NovelQA dataset is partially available via Hugging Face. To obtain the available portion, you must install the huggingface_hub library, log in via the CLI, and clone the repository. Note that some books in NovelQA are not publicly available; for the full dataset, you must contact the original authors of the NovelQA paper.

pip install huggingface_hub
huggingface-cli login
git clone https://huggingface.co/datasets/NovelQA/NovelQA

The complete pipeline (tree construction, graph extraction, and answer generation) is managed by main.py.

Before running, you must prepare a YAML configuration file to define your parameters. An example configuration can be found at ./configs/example_config.yaml.

python main.py --config <path_to_config_file>

E²GraphRAG supports the following datasets:

• NovelQA: Partially open-source; full access requires contacting the original authors.
• InfiniteBench: Fully open-source and publicly available.

Usage Note: After downloading your data, you must specify the data path when initializing the Dataloader class. Detailed instructions for data acquisition can be found in ./data/README.md.

The E²GraphRAG pipeline is executed by running main.py with a configuration file. The pipeline orchestrates text splitting, tree building, graph extraction, and Question Answering (QA) using a retrieved graph structure.

By default, the script looks for a config.yaml file. You can specify a different configuration file using the --config flag.

Workflow Overview:

1. Load Dataset: Loads data based on dataset_name and dataset_path in the config.
2. Text Splitting: Splits text using sequential or nn (newline) methods.
3. Parallel Processing: Executes build_tree_task and extract_graph_task in parallel using multiprocessing.
4. Retrieval & QA: Uses a Retriever to find relevant graph components and an LLM to generate answers or select multiple-choice options.

python main.py --config your_config.yaml

The pipeline is designed to be resumable. If a crash occurs during dataset processing, you can skip the failed items by updating the resumeIndex in your configuration file.

Error Handling:

• The script uses multiprocessing.set_start_method("spawn") to ensure safe parallel execution on CUDA devices.
• If an error occurs, the script attempts to terminate active child processes and clean up CUDA memory using clean_cuda_memory and torch.cuda.empty_cache().

The pipeline behavior is controlled by a YAML configuration file. Key sections include:

• llm: Defines llm_path (model location) and llm_device (e.g., cuda:0).
• dataset: Specifies dataset_name (e.g., NovelQA, InfiniteChoice, InfiniteQALoader) and dataset_path.
• cluster: Controls text chunking with length, overlap, and merge_num.
• extractor: Defines the method and language for graph extraction.
• retriever: Contains kwargs passed to the Retriever class.
• paths: Defines cache_path for intermediate files and answer_path for results.
• resume: Uses resumeIndex to skip already processed items in a dataset.
• split_method: Options are sequential or nn.