TXT2KG w/ hotpot_qa.pyGraph Retrieval Example and tech_qa.py HybridRAG e2e workflow example

CHANGELOG.md

@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Adds TXT2KG class with example on HotPotQA ([#9846](https://github.com/pyg-team/pytorch_geometric/pull/9846))
 - Improving documentation for [cuGraph](https://github.com/rapidsai/cugraph) ([#10083](https://github.com/pyg-team/pytorch_geometric/pull/10083))
 - Added `HashTensor` ([#10072](https://github.com/pyg-team/pytorch_geometric/pull/10072))
 - Added `SGFormer` model and example ([#9904](https://github.com/pyg-team/pytorch_geometric/pull/9904))

examples/llm/README.md

@@ -1,11 +1,12 @@
 # Examples for Co-training LLMs and GNNs
 
-| Example                                      | Description                                                                                                                                                                                                                                                                                                    |
-| -------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| [`g_retriever.py`](./g_retriever.py)         | Example for Retrieval-Augmented Generation (RAG) w/ GNN+LLM by co-training `LLAMA3` with `GAT` for answering questions based on knowledge graph information. We also have an [example repo](https://github.com/neo4j-product-examples/neo4j-gnn-llm-example) for integration with [Neo4j Graph DBs][neo4j.com] |
-| [`g_retriever_utils/`](./g_retriever_utils/) | Contains multiple scripts for benchmarking GRetriever's architecture and evaluating different retrieval methods.                                                                                                                                                                                               |
-| [`multihop_rag/`](./multihop_rag/)           | Contains starter code and an example run for building a Multi-hop dataset using WikiHop5M and 2WikiMultiHopQA                                                                                                                                                                                                  |
-| [`nvtx_examples/`](./nvtx_examples/)         | Contains examples of how to wrap functions using the NVTX profiler for CUDA runtime analysis.                                                                                                                                                                                                                  |
-| [`molecule_gpt.py`](./molecule_gpt.py)       | Example for MoleculeGPT: Instruction Following Large Language Models for Molecular Property Prediction. Supports MoleculeGPT and InstructMol dataset                                                                                                                                                           |
-| [`glem.py`](./glem.py)                       | Example for [GLEM](https://arxiv.org/abs/2210.14709), a GNN+LLM co-training model via variational Expectation-Maximization (EM) framework on node classification tasks to achieve SOTA results                                                                                                                 |
-| [`git_mol.py`](./git_mol.py)                 | Example for GIT-Mol: A Multi-modal Large Language Model for Molecular Science with Graph, Image, and Text                                                                                                                                                                                                      |
+| Example                                      | Description                                                                                                                                                                                                                                                                                                                                        |
+| -------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| [`g_retriever.py`](./g_retriever.py)         | Example for Retrieval-Augmented Generation (RAG) w/ GNN+LLM by co-training `LLAMA3` with `GAT` for answering questions based on knowledge graph information. We also have an [example repo](https://github.com/neo4j-product-examples/neo4j-gnn-llm-example) for integration with [Neo4j Graph DBs][neo4j.com]                                     |
+| [`g_retriever_utils/`](./g_retriever_utils/) | Contains multiple scripts for benchmarking GRetriever's architecture and evaluating different retrieval methods.                                                                                                                                                                                                                                   |
+| [`nvtx_examples/`](./nvtx_examples/)         | Contains examples of how to wrap functions using the NVTX profiler for CUDA runtime analysis.                                                                                                                                                                                                                                                      |
+| [`molecule_gpt.py`](./molecule_gpt.py)       | Example for MoleculeGPT: Instruction Following Large Language Models for Molecular Property Prediction. Supports MoleculeGPT and InstructMol dataset                                                                                                                                                                                               |
+| [`glem.py`](./glem.py)                       | Example for [GLEM](https://arxiv.org/abs/2210.14709), a GNN+LLM co-training model via variational Expectation-Maximization (EM) framework on node classification tasks to achieve SOTA results                                                                                                                                                     |
+| [`git_mol.py`](./git_mol.py)                 | Example for GIT-Mol: A Multi-modal Large Language Model for Molecular Science with Graph, Image, and Text                                                                                                                                                                                                                                          |
+| [`hotpot_qa.py`](./hotpot_qa.py)             | Example for converting adapting the retrieval step of conventional Retrieval-Augmented Generation (RAG) for use with G-retriever, and how to approximate the precision/recall of a subgraph retrieval method. Uses the HotPotQA dataset from [Hugging Face](https://huggingface.co/datasets/hotpotqa/hotpot_qa). This is it is multihop in nature. |
+| [`tech_qa.py`](./tech_qa.py)                 | Full end 2 end HybridRAG pipeline combining txt2kg and retrieval from `hotpot_qa.py` and training/testing from g_retriever.py with VectorRAG. Uses the techQA dataset from [Hugging Face](https://huggingface.co/datasets/rojagtap/tech-qa)                                                                                                        |

examples/llm/g_retriever.py

@@ -131,7 +131,7 @@ def load_params_dict(model, save_path):
     return model
 
 
-def get_loss(model, batch, model_save_name: str) -> Tensor:
+def get_loss(model, batch, model_save_name="gnn+llm") -> Tensor:
     """Compute the loss for a given model and batch of data.
 
     Args:
@@ -148,17 +148,17 @@ def get_loss(model, batch, model_save_name: str) -> Tensor:
         return model(batch.question, batch.label, batch.desc)
     else:  # (GNN+LLM)
         return model(
-            batch.question,
-            batch.x,  # node features
-            batch.edge_index,  # edge indices
-            batch.batch,  # batch indices
-            batch.label,  # answers (labels)
+            batch.question,  # ["list", "of", "questions", "here"]
+            batch.x,  # [num_nodes, num_features]
+            batch.edge_index,  # [2, num_edges]
+            batch.batch,  # which node belongs to which batch index
+            batch.label,  # list answers (labels)
             batch.edge_attr,  # edge attributes
-            batch.desc  # description
+            batch.desc  # list of text graph descriptions
         )
 
 
-def inference_step(model, batch, model_save_name):
+def inference_step(model, batch, model_save_name="gnn+llm"):
     """Performs inference on a given batch of data using the provided model.
 
     Args:
@@ -174,14 +174,31 @@ def inference_step(model, batch, model_save_name):
         # Perform inference on the question and textual graph description
         return model.inference(batch.question, batch.desc)
     else:  # (GNN+LLM)
-        return model.inference(
-            batch.question,
-            batch.x,  # node features
-            batch.edge_index,  # edge indices
-            batch.batch,  # batch indices
-            batch.edge_attr,  # edge attributes
-            batch.desc  # description
-        )
+        return model.inference(batch.question, batch.x, batch.edge_index,
+                               batch.batch, batch.edge_attr, batch.desc)
+
+
+def adjust_learning_rate(param_group, LR, epoch, num_epochs):
+    """Decay learning rate with half-cycle cosine after warmup.
+
+    Args:
+        param_group (dict): Parameter group.
+        LR (float): Learning rate.
+        num_epochs (int): Current epoch.
+
+    Returns:
+        float: Adjusted learning rate.
+    """
+    min_lr = 5e-6
+    warmup_epochs = 1
+    if epoch < warmup_epochs:
+        lr = LR
+    else:
+        lr = min_lr + (LR - min_lr) * 0.5 * (
+            1.0 + math.cos(math.pi * (epoch - warmup_epochs) /
+                           (num_epochs - warmup_epochs)))
+    param_group['lr'] = lr
+    return lr
 
 
 def train(
@@ -216,28 +233,6 @@ def train(
     Returns:
         None
     """
-    def adjust_learning_rate(param_group, LR, epoch):
-        """Decay learning rate with half-cycle cosine after warmup.
-
-        Args:
-            param_group (dict): Parameter group.
-            LR (float): Learning rate.
-            epoch (int): Current epoch.
-
-        Returns:
-            float: Adjusted learning rate.
-        """
-        min_lr = 5e-6
-        warmup_epochs = 1
-        if epoch < warmup_epochs:
-            lr = LR
-        else:
-            lr = min_lr + (LR - min_lr) * 0.5 * (
-                1.0 + math.cos(math.pi * (epoch - warmup_epochs) /
-                               (num_epochs - warmup_epochs)))
-        param_group['lr'] = lr
-        return lr
-
     # Start training time
     start_time = time.time()
 
@@ -322,7 +317,8 @@ def adjust_learning_rate(param_group, LR, epoch):
 
             if (step + 1) % 2 == 0:
                 adjust_learning_rate(optimizer.param_groups[0], lr,
-                                     step / len(train_loader) + epoch)
+                                     step / len(train_loader) + epoch,
+                                     num_epochs)
 
             optimizer.step()
             epoch_loss = epoch_loss + float(loss)

examples/llm/g_retriever_utils/README.md

@@ -1,11 +1,8 @@
 # Examples for LLM and GNN co-training
 
-| Example                                                          | Description                                                                                                                                                               |
-| ---------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| [`rag_feature_store.py`](./rag_feature_store.py)                 | A Proof of Concept Implementation of a RAG enabled FeatureStore that can serve as a starting point for implementing a custom RAG Remote Backend                           |
-| [`rag_graph_store.py`](./rag_graph_store.py)                     | A Proof of Concept Implementation of a RAG enabled GraphStore that can serve as a starting point for implementing a custom RAG Remote Backend                             |
-| [`rag_backend_utils.py`](./rag_backend_utils.py)                 | Utility functions used for loading a series of Knowledge Graph Triplets into the Remote Backend defined by a FeatureStore and GraphStore                                  |
-| [`rag_generate.py`](./rag_generate.py)                           | Script for generating a unique set of subgraphs from the WebQSP dataset using a custom defined retrieval algorithm (defaults to the FeatureStore and GraphStore provided) |
-| [`benchmark_model_archs_rag.py`](./benchmark_model_archs_rag.py) | Script for running a GNN/LLM benchmark on GRetriever while grid searching relevent architecture parameters and datasets.                                                  |
+| Example                                                          | Description                                                                                                              |
+| ---------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------ |
+| [`benchmark_model_archs_rag.py`](./benchmark_model_archs_rag.py) | Script for running a GNN/LLM benchmark on GRetriever while grid searching relevent architecture parameters and datasets. |
+| [`minimal_demo.py`](./minimal_demo.py)                           | Minimal demo for WebQSP dataset comparing GNN+LLM vs LLM                                                                 |
 
 NOTE: Evaluating performance on GRetriever with smaller sample sizes may result in subpar performance. It is not unusual for the fine-tuned model/LLM to perform worse than an untrained LLM on very small sample sizes.