GFM-RAG: Graph Foundation Model for Retrieval Augmented Generation

GFM-RAG provides two primary workflow families depending on your research or application needs:

GFM-RAG Workflow (gfmrag/workflow/config/gfm_rag/)

Use this for the original GFM-RAG implementation. Presets include:

• Indexing datasets
• KGC (Knowledge Graph Construction) training
• Retrieval and QA inference
• SFT (Supervised Fine-Tuning) training
• Path visualization

G-reasoner Workflow (gfmrag/workflow/config/gfm_reasoner/)

Use this for the G-reasoner implementation. Presets include:

• Indexing datasets
• KGC training
• Retrieval and QA inference
• SFT training
• SFT training with answer supervision
• Path visualization

GFM-RAG logs different metrics depending on the workflow being used.

Pre-training (kgc_training.py)

• Training: train/loss (per batch), train/epoch_loss (avg per epoch), train/epoch.
• Evaluation: eval/{dataset_name}/mr, eval/{dataset_name}/mrr, eval/{dataset_name}/hits@{k} (k=1,3,10), eval/mrr, eval/best_mrr.
• Artifacts: Best model checkpoints and epoch-wise checkpoints.

Fine-tuning (sft_training.py)

• Training: train/{loss_name} (per batch), train/epoch_{loss_name} (avg per epoch), train/epoch.
• Evaluation: eval/{watched_metric}, eval/best_{watched_metric}, test/{dataset_name}/{metric}.
• Artifacts: Best model checkpoints and epoch-wise checkpoints.

G-reasoner presets use a Hydra-based configuration pattern. When inspecting or modifying the YAML files in gfm_reasoner/, note the following structure:

• hydra.run.dir: Defines the root directory for all task outputs.
• defaults: A list of shared component groups that are pulled into the configuration. Common shared groups include ner_model, openie_model, el_model, text_emb_model, and wandb.
• Task-specific overrides: The configuration uses specific sections to override the shared defaults for the current task. These sections include dataset, datasets, graph_retriever, model, trainer, llm, and test.

GFM-RAG uses Hydra-based YAML configuration presets located in gfmrag/workflow/config/gfm_rag/. These presets follow a consistent hierarchical pattern:

1. hydra.run.dir: Defines the root directory where outputs and logs for the specific run will be stored.
2. defaults: A list of shared component groups that provide baseline configurations. Common shared groups include:
   • ner_model (Named Entity Recognition)
   • openie_model (Open Information Extraction)
   • el_model (Entity Linking)
   • text_emb_model (Text Embeddings)
   • doc_ranker (Document Ranking)
   • wandb (Weights & Biases integration)
3. Task-specific overrides: Sections like dataset, datasets, graph_retriever, model, trainer, llm, and test are used to override the baseline settings in defaults for a specific workflow task.

GFM-RAG is a Retrieval Augmented Generation (RAG) pipeline that uses a Graph Foundation Model (GFM) to reason over a graph-index.

The Workflow:

1. Graph Indexing: A graph is built from documents to capture relationships between knowledge entities.
2. Retrieval: A query and the graph-index are fed into a pre-trained GFM retriever.
3. Generation: The retriever identifies relevant documents, which are then used by an LLM for question answering.

Key advantages include efficiency in multi-hop reasoning via single-step retrieval, generalizability to unseen datasets, and interpretability through captured reasoning paths.

The wandb integration is multi-GPU aware. To prevent duplicate logs and ensure data integrity:

• Only the rank 0 process logs to wandb.
• Metrics are synchronized across all processes.
• Model checkpoints are saved only on rank 0.

# Multi-GPU training with wandb
torchrun --nproc_per_node=4 -m gfmrag.workflow.kgc_training --config-path config/gfm_rag

GFM-RAG provides a unified graph interface that allows you to either build a graph from scratch or use an existing one.

• Path A (Start from raw documents): You provide raw JSON files. GFMRetriever.from_index(...) automatically runs the index-building pipeline (NER, OpenIE, Entity Linking, etc.) to create the graph files in processed/stage1/.
• Path B (Bring your own graph): If you already have a graph that follows the GFM-RAG three-file format (nodes.csv, relations.csv, and edges.csv), you can place them in processed/stage1/ and skip the index-building step entirely. This is useful for using existing Knowledge Graphs or graphs from other pipelines.

The project maintains two lines of retriever development:

1. GFM-RAG: The original graph foundation model retriever. It serves as the baseline workflow and provides the published checkpoint family.
2. G-reasoner: The latest version of the retriever, featuring a new architecture and improved performance.

Both lines utilize the same indexing, dataset, QA, and documentation structures provided in this repository.

The wandb integration is enabled by default. You can run pre-training or fine-tuning workflows normally to start logging metrics automatically.

# Pre-training with wandb logging
python -m gfmrag.workflow.kgc_training --config-path config/gfm_rag

# Fine-tuning with wandb logging
python -m gfmrag.workflow.sft_training --config-path config/gfm_rag

To reproduce the results reported in the GFM-RAG paper, follow a three-stage workflow:

1. Graph Indexing: Build a knowledge graph from raw documents.
2. GFM Training: Perform unsupervised KG pre-training (optional) followed by supervised fine-tuning.
3. QA Evaluation: Run single-step (Stage 3a) or multi-step IRCoT (Stage 3b) reasoning to generate answers.

Note on Versioning: The original paper results were produced with v1.0.0. For exact reproduction, use the gfm-rag branch. The current v2.0.0 branch is maintained for backward compatibility.

# The workflow follows these stages:
# 1. Build graph index
# 2. GFM training (pre-training + fine-tuning)
# 3. QA evaluation