加载中…
加载中…
The GFM-RAG is the first graph foundation model-powered RAG pipeline that combines the power of graph neural networks to reason over knowledge graphs and retrieve relevant documents for question answering.
We first build a knowledge graph index (KG-index) from the documents to capture the relationships between knowledge. Then, we feed the query and constructed KG-index into the pre-trained graph foundation model (GFM) retriever to obtain relevant documents for LLM generation. The GFM retriever experiences large-scale training and can be directly applied to unseen datasets without fine-tuning.
For more details, please refer to our project page and paper.
Conda provides an easy way to install the CUDA development toolkit which is required by GFM-RAG
Install packages
conda create -n gfmrag python=3.12
conda activate gfmrag
conda install cuda-toolkit -c nvidia/label/cuda-12.4.1 # Replace with your desired CUDA version
pip install gfmrag
Install other dependencies
TORCH=$(python -c "import torch; print(torch.__version__)")
pip install torch_scatter torch_sparse -f https://data.pyg.org/whl/torch-${TORCH}.html
You need to prepare the following files:
dataset_corpus.json: A JSON file containing the entire document corpus.train.json (optional): A JSON file containing the training data.test.json (optional): A JSON file containing the test data.Place your files in the following structure:
data_name/
├── raw/
│ ├── dataset_corpus.json
│ ├── train.json # (optional)
│ └── test.json # (optional)
└── processed/ # Output directory
dataset_corpus.jsonThe dataset_corpus.json is a dictionary where each key is the title or unique id of a document and the value is the text of the document.
{
"Fred Gehrke":
"Clarence Fred Gehrke (April 24, 1918 – February 9, 2002) was an American football player and executive. He played in the National Football League (NFL) for the Cleveland / Los Angeles Rams, San Francisco 49ers and Chicago Cardinals from 1940 through 1950. To boost team morale, Gehrke designed and painted the Los Angeles Rams logo in 1948, which was the first painted on the helmets of an NFL team. He later served as the general manager of the Denver Broncos from 1977 through 1981. He is the great-grandfather of Miami Marlin Christian Yelich"
,
"Manny Machado":
"Manuel Arturo Machado (] ; born July 6, 1992) is an American professional baseball third baseman and shortstop for the Baltimore Orioles of Major League Baseball (MLB). He attended Brito High School in Miami and was drafted by the Orioles with the third overall pick in the 2010 Major League Baseball draft. He bats and throws right-handed."
,
...
}
train.json and test.json (optional)If you want to train and evaluate the model, you need to provide training and test data in the form of a JSON file. Each entry in the JSON file should contain the following fields:
id: A unique identifier for the example.question: The question or query.supporting_facts: A list of supporting facts for the question. Each supporting fact is a list containing the title of the document that can be found in the dataset_corpus.json file.Each entry can also contain additional fields depending on the task. For example:
answer: The answer to the question.The additional fields will be copied during the following steps of the pipeline.
Example:
[
{
"id": "5adf5e285542992d7e9f9323",
"question": "When was the judge born who made notable contributions to the trial of the man who tortured, raped, and murdered eight student nurses from South Chicago Community Hospital on the night of July 13-14, 1966?",
"answer": "June 4, 1931",
"supporting_facts": [
"Louis B. Garippo",
"Richard Speck"
]
},
{
"id": "5a7f7b365542992097ad2f80",
"question": "Did the Beaulieu Mine or the McIntyre Mines yield gold and copper?",
"answer": "The McIntyre also yielded a considerable amount of copper",
"supporting_facts": [
"Beaulieu Mine",
"McIntyre Mines"
]
}
...
]
You need to create a KG-index configuration file.
Details of the configuration parameters are explained in the KG-index Configuration page.
python -m gfmrag.workflow.stage1_index_dataset
This method performs two main tasks:
kg.txt and document2entities.json) from the dataset_corpus.json fileFiles created:
kg.txt: Contains knowledge graph triplesdocument2entities.json: Maps documents to their entitiestrain.json: Processed training data (if raw exists)test.json: Processed test data (if raw exists)Directory structure:
root/
└── data_name/
├── raw/
│ ├── dataset_corpus.json
│ ├── train.json (optional)
│ └── test.json (optional)
└── processed/
└── stage1/
├── kg.txt
├── document2entities.json
├── train.json
└── test.json
You need to create a configuration file for inference.
Details of the configuration parameters are explained in the GFM-RAG Configuration page.
You can initialize the GFMRetriever with the following code. It will load the pre-trained GFM-RAG model and the KG-index for retrieval.
import logging
import os
import hydra
from hydra.core.hydra_config import HydraConfig
from omegaconf import DictConfig, OmegaConf
from gfmrag import GFMRetriever
logger = logging.getLogger(__name__)
@hydra.main(
config_path="config", config_name="stage3_qa_ircot_inference", version_base=None
)
def main(cfg: DictConfig) -> None:
output_dir = HydraConfig.get().runtime.output_dir
logger.info(f"Config:\n {OmegaConf.to_yaml(cfg)}")
logger.info(f"Current working directory: {os.getcwd()}")
logger.info(f"Output directory: {output_dir}")
gfmrag_retriever = GFMRetriever.from_config(cfg)
You can use GFM-RAG retriever to reason over the KG-index and obtain documents for a given query.
docs = retriever.retrieve("Who is the president of France?", top_k=5)
from hydra.utils import instantiate
from gfmrag.llms import BaseLanguageModel
from gfmrag.prompt_builder import QAPromptBuilder
llm = instantiate(cfg.llm)
qa_prompt_builder = QAPromptBuilder(cfg.qa_prompt)
message = qa_prompt_builder.build_input_prompt(current_query, retrieved_docs)
answer = llm.generate_sentence(message) # Answer: "Emmanuel Macron"
During fine-tuning, the GFM model will be trained on the query-documents pairs train.json from the labeled dataset to learn complex relationships for retrieval.
It can be conduced on your own dataset to improve the performance of the model on your specific domain.
A example of the training data:
[
{
"id": "5abc553a554299700f9d7871",
"question": "Kyle Ezell is a professor at what School of Architecture building at Ohio State?",
"answer": "Knowlton Hall",
"supporting_facts": [
"Knowlton Hall",
"Kyle Ezell"
],
"question_entities": [
"kyle ezell",
"architectural association school of architecture",
"ohio state"
],
"supporting_entities": [
"10 million donation",
"2004",
"architecture",
"austin e knowlton",
"austin e knowlton school of architecture",
"bachelor s in architectural engineering",
"city and regional planning",
"columbus ohio united states",
"ives hall",
"july 2002",
"knowlton hall",
"ksa",
]
},
...
]
You need to create a configuration file for fine-tuning.
Details of the configuration parameters are explained in the GFM-RAG Fine-tuning Configuration page.
You can fine-tune the pre-trained GFM-RAG model on your dataset using the following command:
python -m gfmrag.workflow.stage2_qa_finetune
# Multi-GPU training
torchrun --nproc_per_node=4 gfmrag.workflow.stage2_qa_finetune
# Multi-node Multi-GPU training
torchrun --nproc_per_node=4 --nnodes=2 gfmrag.workflow.stage2_qa_finetune
We greatly appreciate the following repositories for their help to this project:
If you find this repository helpful, please consider citing our paper:
@article{luo2025gfmrag,
title={GFM-RAG: Graph Foundation Model for Retrieval Augmented Generation},
author={Luo, Linhao and Zhao, Zicheng and Haffari, Gholamreza and Phung, Dinh and Gong, Chen and Pan, Shirui},
journal={arXiv preprint arXiv:2502.01113},
year={2025}
}