E²GraphRAG

https://github.com/yibozhao624/e-2graphrag
Admin
E²GraphRAG is a lightweight and modular framework designed to enhance efficiency and effectiveness...
Tokens:5,022
Snippets:34
Trust Score:8.3
License:MIT
Update:6 months ago
Show doc for...
Tokens:
Raw
### Install Hugging Face Hub and Clone NovelQA Dataset

Source: https://github.com/yibozhao624/e-2graphrag/blob/master/data/README.md

This snippet shows how to install the huggingface_hub library, log in to your Hugging Face account, and clone the NovelQA dataset repository. Ensure you have git and Python installed.

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

--------------------------------

### Loading Datasets (Python)

Source: https://context7.com/yibozhao624/e-2graphrag/llms.txt

Provides examples for loading various long-context Question Answering datasets supported by E²GraphRAG, including NovelQA, InfiniteChoice, and InfiniteQA, using both utility functions and specific loaders.

```python
from dataloader import NovelQALoader, InfiniteChoiceLoader, InfiniteQALoader
from utils import load_dataset

# Load using utility function
dataset = load_dataset(
    dataset_name="InfiniteChoice",  # or "NovelQA", "InfiniteQALoader"
    dataset_path="./data/InfiniteBench/longbook_choice_eng.jsonl"
)

# NovelQA dataset (requires books and QA data)
novel_loader = NovelQALoader("./data/NovelQA")
for i in range(len(novel_loader)):
    data = novel_loader[i]
    book_text = data["book"]  # Full novel text
    qa_pairs = data["qa"]     # List of QA dictionaries

    for qa in qa_pairs:
        question = qa["question"]  # Formatted with options A-D
        answer = qa["answer"]      # Correct option
        evidence = qa.get("evidence")  # Supporting text spans

# InfiniteBench Choice dataset (multiple choice)
choice_loader = InfiniteChoiceLoader("./data/InfiniteBench/longbook_choice_eng.jsonl")
sample = choice_loader[0]
# {
#     "book": "Long context text...",
#     "qa": [{"question": "Question?\nA. Option1\nB. Option2...", "answer": "A"}]
# }

# InfiniteBench QA dataset (open-ended answers)
qa_loader = InfiniteQALoader("./data/InfiniteBench/longbook_qa_eng.jsonl")
sample = qa_loader[0]
# {
#     "book": "Long context text...",
#     "qa": [{"question": "Who is the main character?", "answer": "John Smith"}]
# }
```

--------------------------------

### Install E²GraphRAG Dependencies

Source: https://github.com/yibozhao624/e-2graphrag/blob/master/README.md

Installs the necessary Python packages required to run the E²GraphRAG framework. This command should be executed in your terminal within the project environment.

```bash
pip install -r requirements.txt
```

--------------------------------

### Download InfiniteBench Datasets

Source: https://github.com/yibozhao624/e-2graphrag/blob/master/data/README.md

Instructions for downloading the InfiniteBench dataset files 'longbook_choice_eng.jsonl' and 'longbook_qa_eng.jsonl'. These files should be placed in the './data/InfiniteBench' directory. A download script is also available.

```bash
# Instructions to download files from https://huggingface.co/datasets/xinrongzhang2022/InfiniteBench/tree/main
# and place them in ./data/InfiniteBench
# A download.sh script is also provided in the original repo.
```

--------------------------------

### Running the E²GraphRAG Full Pipeline (Python)

Source: https://context7.com/yibozhao624/e-2graphrag/llms.txt

This snippet demonstrates how to run the complete E²GraphRAG pipeline, including tree building, graph extraction, and question answering, using a provided configuration file. It orchestrates these processes in parallel for efficiency.

```python
# Run the complete pipeline via command line
# python main.py --config ./configs/example_config.yaml

# Example YAML configuration (configs/example_config.yaml)
"""
dataset:
  dataset_name: InfiniteChoice
  dataset_path: ./data/InfiniteBench/longbook_choice_eng.jsonl

llm:
  llm_name: Qwen2.5-7B-Instruct
  llm_path: /path/to/your/model
  llm_device: cuda:0

paths:
  log_path: ./log
  answer_path: ./answer/test
  cache_path: ./cache/test

extractor:
  language: "en"
  method: "Spacy"  # Options: Spacy, NLTK, BERT_NER_POS

cluster:
  length: 1200
  overlap: 100
  merge_num: 5
  force_Reextract: False

resume:
  resumeIndex: 0

retriever:
  kwargs:
    device: cuda
    shortest_path_k: 4
    debug: True
    merge_num: 5
    overlap: 100
    tokenizer: /path/to/tokenizer
    max_chunk_setting: 25
"""

```

--------------------------------

### Building Hierarchical Summary Trees (Python)

Source: https://context7.com/yibozhao624/e-2graphrag/llms.txt

This code demonstrates how to build a hierarchical summarization tree from text chunks using an LLM. It involves loading a tokenizer and LLM pipeline, splitting text, and then calling the `build_tree` function. Caching mechanisms are supported to load previously generated trees.

```python
from build_tree import build_tree, load_cache_summary
from transformers import pipeline, AutoTokenizer
import torch

# Load tokenizer and LLM pipeline
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
llm = pipeline(
    "text-generation",
    model="Qwen/Qwen2.5-7B-Instruct",
    tokenizer=tokenizer,
    device="cuda:0",
    torch_dtype=torch.float16,
    max_new_tokens=1200
)

# Split text into chunks
text_chunks = [
    "Chapter 1: The protagonist John meets Mary in the old library...",
    "Chapter 2: They discover a hidden passage behind the bookshelf...",
    "Chapter 3: The mysterious letter reveals a family secret...",
]

# Build the hierarchical tree
cache_folder = "./cache/book_1"
tree, time_cost = build_tree(
    text_chunks=text_chunks,
    llm=llm,
    cache_folder=cache_folder,
    tokenizer=tokenizer,
    length=1200,      # chunk length in tokens
    overlap=100,      # overlap between chunks
    merge_num=5,      # chunks to merge per summary
    language="en"
)

# Tree structure example:
# {
#     "leaf_0": {"text": "Chapter 1...", "children": None, "parent": "summary_0_0"},
#     "leaf_1": {"text": "Chapter 2...", "children": None, "parent": "summary_0_0"},
#     "summary_0_0": {"text": "Summary of chapters 1-5...", "children": ["leaf_0", "leaf_1", ...], "parent": "summary_1_0"},
#     "summary_1_0": {"text": "Higher-level summary...", "children": ["summary_0_0", ...], "parent": []}
# }

# Load cached tree
cached_tree = load_cache_summary("./cache/book_1/tree.json")

```

--------------------------------

### Build Graph from Triplets (Python)

Source: https://context7.com/yibozhao624/e-2graphrag/llms.txt

Demonstrates how to construct a graph data structure from a list of co-occurrence triplets. This is a foundational step for graph-based retrieval.

```python
triplets = [
    ("John", "Mary", 3),   # John and Mary co-occur 3 times
    ("John", "Google", 2),
    ("Mary", "Microsoft", 1)
]
G = build_graph(triplets)
```

--------------------------------

### Evaluating Results (Python)

Source: https://context7.com/yibozhao624/e-2graphrag/llms.txt

Demonstrates how to calculate standard evaluation metrics for generated answers, specifically Exact Match (EM) and Rouge-L (RL) scores, using utility functions provided by the E²GraphRAG project.

```python
from utils import EM_score, RL_score
from evaluate import calculate_metrics, calculate_time_cost

# Calculate individual scores
pred = "John met Mary in the library"
gold = "john met mary in the library"

em = EM_score(pred, gold)  # Returns 1.0 (exact match after normalization)
rl = RL_score(pred, gold)  # Returns Rouge-L F1 score (0.0-1.0)
```

--------------------------------

### Querying with Retriever (Python)

Source: https://context7.com/yibozhao624/e-2graphrag/llms.txt

Initializes and uses the Retriever class for hybrid search, combining graph-based local search with dense vector retrieval. It can query for relevant context based on a question and update its internal state with new documents.

```python
from query import Retriever
from extract_graph import load_nlp
import networkx as nx

# Initialize components
nlp = load_nlp("en", "Spacy")

# Load cached tree and graph (or use freshly built ones)
# cache_tree: hierarchical summary tree
# G: entity co-occurrence graph
# index: entity-to-chunk mapping
# appearance_count: entity frequency statistics

retriever = Retriever(
    cache_tree=tree,
    G=graph,
    index=index,
    appearance_count=appearance_count,
    nlp=nlp,
    device="cuda:0",
    merge_num=5,
    min_count=2,
    overlap=100,
    tokenizer="/path/to/tokenizer",
    embedder="BAAI/bge-m3"  # Sentence embedding model for dense retrieval
)

# Query for relevant context
question = "What happened when John met Mary in the library?"

result = retriever.query(
    query=question,
    shortest_path_k=4,       # Max path length for graph traversal
    max_chunk_setting=25,    # Max chunks to retrieve
    debug=True
)

# Result structure:
# {
#     "chunks": "John_Mary: John first met Mary in the old library on a rainy afternoon...",
#     "entities": ["John", "Mary", "library"],
#     "keys": ["John_Mary"],
#     "len_chunks": 15,
#     "retrieval_type": "Local, Loop for 2 times"  # or "Global Search", "Occurrence Rerank"
# }

# Update retriever for a new document
retriever.update(new_tree, new_graph, new_index, new_appearance_count)

# Direct dense retrieval (when no entities found in query)
dense_result = retriever.dense_retrieval(query="Describe the atmosphere", k=10)
```

--------------------------------

### Text Processing Utilities: Splitting and Timing

Source: https://context7.com/yibozhao624/e-2graphrag/llms.txt

Provides helper functions for text splitting using a transformer tokenizer and for timing operations using a Timer class. It demonstrates splitting a long text into chunks with specified length and overlap, and timing different stages of a process like tree building and graph extraction.

```python
from utils import sequential_split, Timer, setup_logging
from transformers import AutoTokenizer
import logging

# Setup logging
logger = setup_logging(level=logging.INFO, log_file="./log/run.log")

# Split long text into overlapping chunks
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
long_text = "Very long document text..."

chunks = sequential_split(
    text=long_text,
    tokenizer=tokenizer,
    length=1200,    # tokens per chunk
    overlap=100     # overlapping tokens
)

# Using the Timer class for performance tracking
timer = Timer()

with timer.timer("tree_building"):
    # ... tree building code ...
    pass

with timer.timer("graph_extraction"):
    # ... graph extraction code ...
    pass

print(f"Tree building took: {timer['tree_building']:.2f} seconds")
print(f"Graph extraction took: {timer['graph_extraction']:.2f} seconds")
print(timer.summary())
```

--------------------------------

### Execute E²GraphRAG Pipeline

Source: https://github.com/yibozhao624/e-2graphrag/blob/master/README.md

Runs the main pipeline which includes tree construction, graph extraction, and answer generation. It requires a path to a valid YAML configuration file as an argument.

```bash
python main.py --config <path_to_config_file>
```

--------------------------------

### Calculate Evaluation Metrics and Processing Time

Source: https://context7.com/yibozhao624/e-2graphrag/llms.txt

Calculates evaluation metrics such as average exact match (EM) and relevance (RL) for answers, and also computes processing time statistics for build and extraction phases. It takes folder paths and dataset names as input and prints the results.

```python
avg_em, avg_rl, total_questions = calculate_metrics(
    answer_folder="./answer/test",
    dataset_name="InfiniteChoice"
)
print(f"EM: {avg_em:.4f}, RL: {avg_rl:.4f}, Total: {total_questions}")

avg_build, avg_extract, total_samples = calculate_time_cost(
    cache_folder="./cache/test",
    dataset_name="InfiniteChoice"
)
print(f"Avg Build: {avg_build:.2f}s, Avg Extract: {avg_extract:.2f}s")
```

--------------------------------

### Prompt Templates for E²GraphRAG

Source: https://context7.com/yibozhao624/e-2graphrag/llms.txt

Defines structured prompt templates used by E²GraphRAG for various tasks including summarization of leaf nodes, merging summaries, and question answering (both multiple-choice and open-ended). It shows how to format these prompts with specific content.

```python
from prompt_dict import Prompts

# Summarization prompt for leaf nodes (original text)
leaf_prompt = Prompts["summarize_details"].format(
    content="Chapter text content here..."
)
# Output: Prompt asking LLM to summarize with main characters, plot, and details

# Summarization prompt for merging summaries
summary_prompt = Prompts["summarize_summary"].format(
    summary="Previous summary content..."
)

# QA prompt for multiple choice questions
qa_prompt = Prompts["QA_prompt_options"].format(
    question="What was John's occupation?\nA. Doctor\nB. Engineer\nC. Teacher\nD. Lawyer",
    evidence="John_occupation: John worked as a software engineer at Google..."
)

# QA prompt for open-ended questions
open_qa_prompt = Prompts["QA_prompt_answer"].format(
    question="Where did John and Mary meet?",
    evidence="John_Mary: They first encountered each other in the old library..."
)

# Chinese language variants also available:
# Prompts["summarize_details_zh"], Prompts["summarize_summary_zh"], Prompts["QA_prompt_answer_zh"]
```

--------------------------------

### BERT-based Entity Extraction (Python)

Source: https://context7.com/yibozhao624/e-2graphrag/llms.txt

Utilizes the BERTExtractor to perform Named Entity Recognition (NER) and Part-of-Speech (POS) tagging using pre-trained transformer models. It processes text to identify nouns, co-occurring entities, and their appearance counts.

```python
from extract_graph import BERTExtractor

# Initialize BERT extractor with custom models
bert_extractor = BERTExtractor(
    language="en",
    ner_model_name="./models/ner",  # Path to NER model
    pos_model_name="./models/pos"   # Path to POS tagging model
)

text = """The European Union announced new regulations affecting tech companies.
Angela Merkel discussed the policy changes with Emmanuel Macron in Brussels."""

result = bert_extractor(text)

# Result includes entities and nouns detected by BERT:
# {
#     "nouns": ["European Union", "regulations", "tech companies", "Angela Merkel", "Emmanuel Macron", "Brussels", "policy", "changes"],
#     "cooccurrence": {("Angela Merkel", "Emmanuel Macron"): 1, ("European Union", "regulations"): 1, ...},
#     "double_nouns": {},
#     "appearance_count": {"European Union": 1, "Angela Merkel": 1, ...}
# }
```

--------------------------------

### Extracting Entity Graphs from Text (Python)

Source: https://context7.com/yibozhao624/e-2graphrag/llms.txt

This code snippet shows how to extract entities and their co-occurrence relationships from text using different NLP methods like Spacy, NLTK, or BERT. It includes functions to load NLP models, extract graphs from single texts or multiple chunks, and optionally use caching.

```python
from extract_graph import load_nlp, extract_graph, build_graph, SpacyExtractor

# Initialize NLP extractor (choose method: Spacy, NLTK, or BERT_NER_POS)
nlp = load_nlp(language="en", method="Spacy")

# Extract entities from a single text chunk
text = """John Doe, a software engineer at Google, visited New York last week.
He met with Jane Smith from Microsoft to discuss a potential partnership.
The meeting took place in the Empire State Building."""

result = nlp.naive_extract_graph(text)

# Result structure:
# {
#     "nouns": ["John", "Doe", "Google", "New York", "Jane", "Smith", "Microsoft", "Empire State Building", "engineer", "partnership", "meeting"],
#     "cooccurrence": {("John", "Doe"): 1, ("Google", "John"): 1, ("Jane", "Microsoft"): 1, ...},
#     "double_nouns": {"John": ["John", "Doe"], "Jane": ["Jane", "Smith"]},
#     "appearance_count": {"John": 2, "Google": 1, "New York": 1, ...}
# }

# Extract graph from multiple text chunks
text_chunks = ["Chapter 1 text...", "Chapter 2 text...", "Chapter 3 text..."]
(graph, index, appearance_count), time_cost = extract_graph(
    text=text_chunks,
    cache_folder="./cache/book_1",
    nlp=nlp,
    use_cache=True,
    reextract=False
)

# graph: NetworkX graph with entities as nodes and co-occurrence weights as edges
# index: {"entity_name": ["leaf_0", "leaf_2", ...]} - maps entities to chunks

```
Tokens:
Raw