GraphRAG

An end-to-end GraphRAG pipeline built from scratch: extract entities and relations from a text corpus, cluster them into a knowledge graph, generate descriptions, embed everything, and answer questions via local search over the graph.

Runs on an NVIDIA A40 with 16GB of system ram.

Targets a slice of the wiki_dpr corpus (~21M passages) (for now) and runs entirely on local infrastructure — no external LLM API calls.

Pipeline

1. Extraction (EntityRelationExtractor.py) — chunks the corpus, runs GLiNER-Relex for joint NER/RE, deduplicates entities by embedding similarity, and writes entities/relations to Parquet.
2. Description generation (DescriptionGenerator.py) — async LLM-based summarization for both entities and detected communities, with TF-IDF clustering to downsample excerpts for entities that appear too often.
3. Graph construction (GraphGenerator.py) — builds the graph in igraph, runs Leiden community detection, and embeds entities and communities into ChromaDB.
4. Query engine (GraphRAG.py) — GraphRagQueryEngine performs local search: retrieves relevant entities/communities, expands via graph hops, and asks the LLM to answer using only the retrieved context (with citations).

Stack

• NER/RE: GLiNER-Relex
• Embeddings: BGE-M3 (via FlagEmbedding)
• Vector store: ChromaDB
• Graph: igraph (Leiden clustering)
• Data handling: Polars, PyArrow, DuckDB
• LLM inference: Ollama (Llama 3.1 8B, async)

Setup

conda env create -f environment.yml
conda activate graphrag
./start_ollama.sh   # starts Ollama and creates the forced-gpu-model from ./modelfile

Tested on an A40 GPU (48GB VRAM). start_ollama.sh disables llama.cpp's RAM cache and caps parallelism to avoid the OOM issues that come with default settings.

Usage

cd graphrag_src
python main.py

This builds the graph from wiki_dpr end-to-end, then runs a sample local-search query. Intermediate artifacts (chunks, entities, relations, communities, ChromaDB collections) are cached under ./data.

For a quick sanity check after building the graph, run sample_generated_data.py to print random entity/relation/community samples.

Tests

pytest tests/

Covers EntityRelationExtractor and GraphGenerator.

Evaluation

Will eventually integrate BenchmarkQED AutoQ/AutoE for automated question generation and answer evaluation.