Draft: TXT2KG w/ hotpot_qa.py and tech_qa.py examples

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 `LLAMA2` with `GAT` for answering questions based on knowledge graph information                                    |
-| [`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                                                                                  |
+| Example                                      | Description                                                                                                                                                 |
+| -------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| [`g_retriever.py`](./g_retriever.py)         | Example for Retrieval-Augmented Generation (RAG) w/ GNN+LLM by co-training `LLAMA2` with `GAT` for answering questions based on knowledge graph information |
+| [`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                                                                                         |
 | [`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 as an example since it is multihop in nature. |

examples/llm/g_retriever_utils/rag_generate.py

@@ -1,32 +1,30 @@
 # %%
 import argparse
 from itertools import chain
-from typing import Tuple
 
 import pandas as pd
 import torch
 import tqdm
-from rag_backend_utils import create_remote_backend_from_triplets
-from rag_feature_store import SentenceTransformerFeatureStore
-from rag_graph_store import NeighborSamplingRAGGraphStore
 
 from torch_geometric.data import Data
 from torch_geometric.datasets import WebQSPDataset
-from torch_geometric.datasets.web_qsp_dataset import (
-    preprocess_triplet,
-    retrieval_via_pcst,
-)
+from torch_geometric.datasets.web_qsp_dataset import preprocess_triplet
 from torch_geometric.loader import RAGQueryLoader
 from torch_geometric.nn.nlp import SentenceTransformer
+from torch_geometric.utils.rag.backend_utils import (
+    create_remote_backend_from_triplets,
+    make_pcst_filter,
+)
+from torch_geometric.utils.rag.feature_store import (
+    SentenceTransformerFeatureStore,
+)
+from torch_geometric.utils.rag.graph_store import NeighborSamplingRAGGraphStore
 
 # %%
-parser = argparse.ArgumentParser(
-    description="""Generate new WebQSP subgraphs\n""" +
-    """NOTE: Evaluating with smaller samples may result in""" +
-    """ poorer performance for the trained models compared""" +
-    """ to untrained models.""")
+parser = argparse.ArgumentParser(description="""Generate new WebQSP subgraphs
+NOTE: Evaluating with smaller samples may result in poorer performance for the trained models compared to untrained models."""
+                                 )
 # TODO: Add more arguments for configuring rag params
-parser.add_argument("--use_pcst", action="store_true")
 parser.add_argument("--num_samples", type=int, default=4700)
 parser.add_argument("--out_file", default="subg_results.pt")
 args = parser.parse_args()
@@ -56,37 +54,10 @@
     feature_db=SentenceTransformerFeatureStore).load()
 
 # %%
-
-
-def apply_retrieval_via_pcst(graph: Data, query: str, topk: int = 3,
-                             topk_e: int = 3,
-                             cost_e: float = 0.5) -> Tuple[Data, str]:
-    q_emb = model.encode(query)
-    textual_nodes = ds.textual_nodes.iloc[graph["node_idx"]].reset_index()
-    textual_edges = ds.textual_edges.iloc[graph["edge_idx"]].reset_index()
-    out_graph, desc = retrieval_via_pcst(graph, q_emb, textual_nodes,
-                                         textual_edges, topk, topk_e, cost_e)
-    out_graph["desc"] = desc
-    return out_graph
-
-
-def apply_retrieval_with_text(graph: Data, query: str) -> Tuple[Data, str]:
-    textual_nodes = ds.textual_nodes.iloc[graph["node_idx"]].reset_index()
-    textual_edges = ds.textual_edges.iloc[graph["edge_idx"]].reset_index()
-    desc = (
-        textual_nodes.to_csv(index=False) + "\n" +
-        textual_edges.to_csv(index=False, columns=["src", "edge_attr", "dst"]))
-    graph["desc"] = desc
-    return graph
-
-
-transform = apply_retrieval_via_pcst \
-    if args.use_pcst else apply_retrieval_with_text
-
 query_loader = RAGQueryLoader(data=(fs, gs), seed_nodes_kwargs={"k_nodes": 5},
                               seed_edges_kwargs={"k_edges": 5},
                               sampler_kwargs={"num_neighbors": [50] * 2},
-                              local_filter=transform)
+                              local_filter=make_pcst_filter(triplets, model))
 
 
 # %%

examples/llm/hotpot_qa.py

@@ -0,0 +1,148 @@
+import argparse
+import os
+from itertools import chain
+
+import datasets
+import torch
+from tqdm import tqdm
+
+from torch_geometric import seed_everything
+from torch_geometric.loader import RAGQueryLoader
+from torch_geometric.nn.nlp import TXT2KG, SentenceTransformer
+from torch_geometric.utils.rag.backend_utils import (
+    create_remote_backend_from_triplets,
+    make_pcst_filter,
+    preprocess_triplet,
+)
+from torch_geometric.utils.rag.feature_store import (
+    SentenceTransformerFeatureStore,
+)
+from torch_geometric.utils.rag.graph_store import NeighborSamplingRAGGraphStore
+
+if __name__ == '__main__':
+    seed_everything(50)
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--NV_NIM_MODEL', type=str,
+                        default="nvidia/llama-3.1-nemotron-70b-instruct")
+    parser.add_argument('--NV_NIM_KEY', type=str, default="")
+    parser.add_argument('--local_lm', action="store_true")
+    parser.add_argument('--percent_data', type=float, default=1.0)
+    parser.add_argument(
+        '--chunk_size', type=int, default=512, help=
+        "When splitting context documents, the maximum number of characters per chunk."
+    )
+    parser.add_argument('--checkpointing', action="store_true")
+    parser.add_argument('--verbose', action="store_true")
+    args = parser.parse_args()
+    assert args.percent_data <= 100 and args.percent_data > 0
+    if args.local_lm:
+        kg_maker = TXT2KG(
+            local_LM=True,
+            chunk_size=args.chunk_size,
+        )
+    else:
+        kg_maker = TXT2KG(
+            NVIDIA_NIM_MODEL=args.NV_NIM_MODEL,
+            NVIDIA_API_KEY=args.NV_NIM_KEY,
+            chunk_size=args.chunk_size,
+        )
+    if os.path.exists("hotpot_kg.pt"):
+        print("Re-using existing KG...")
+        relevant_triples = torch.load("hotpot_kg.pt")
+    elif os.path.exists("checkpoint_kg.pt"):
+        print("Re-using existing checkpoint...")
+        relevant_triples = torch.load("checkpoint_kg.pt")
+    else:
+        # Use training set for simplicity since our retrieval method is nonparametric
+        raw_dataset = datasets.load_dataset('hotpotqa/hotpot_qa', 'fullwiki',
+                                            trust_remote_code=True)["train"]
+        # Build KG
+        num_data_pts = len(raw_dataset)
+        data_idxs = torch.randperm(num_data_pts)[0:int(num_data_pts *
+                                                       args.percent_data /
+                                                       100.0)]
+        if args.checkpointing:
+            five_percent = int(len(data_idxs) / 20)
+            count = 0
+        for idx in tqdm(data_idxs, desc="Building KG"):
+            data_point = raw_dataset[int(idx)]
+            q = data_point["question"]
+            a = data_point["answer"]
+            context_doc = ''
+            for i in data_point["context"]["sentences"]:
+                for sentence in i:
+                    context_doc += sentence
+
+            QA_pair = (q, a)
+            kg_maker.add_doc_2_KG(
+                txt=context_doc,
+                QA_pair=QA_pair,
+            )
+            if args.checkpointing:
+                count += 1
+                if count == five_percent:
+                    print("checkpointing...")
+                    count = 0
+                    kg_maker.save_kg("checkpoint_kg.pt")
+        kg_maker.save_kg("hotpot_kg.pt")
+        if args.checkpointing:
+            # delete checkpoint
+            os.remove("checkpoint_kg.pt")
+
+        relevant_triples = kg_maker.relevant_triples
+        if args.local_lm:
+            print("Total number of context characters parsed by LLM",
+                  kg_maker.total_chars_parsed)
+            print(
+                "Average number of context characters parsed by LLM per second=",
+                kg_maker.avg_chars_parsed_per_sec)
+
+    print("Size of KG (number of triples) =",
+          sum([len(rel_trips) for rel_trips in relevant_triples.values()]))
+
+    triples = list(
+        chain.from_iterable(triple_set
+                            for triple_set in relevant_triples.values()))
+    # debug
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = SentenceTransformer(
+        model_name='sentence-transformers/all-roberta-large-v1').to(device)
+    fs, gs = create_remote_backend_from_triplets(
+        triplets=triples, node_embedding_model=model,
+        node_method_to_call="encode", path="backend",
+        pre_transform=preprocess_triplet, node_method_kwargs={
+            "batch_size": min(len(triples), 256)
+        }, graph_db=NeighborSamplingRAGGraphStore,
+        feature_db=SentenceTransformerFeatureStore).load()
+    query_loader = RAGQueryLoader(
+        data=(fs, gs), seed_nodes_kwargs={"k_nodes":
+                                          5}, seed_edges_kwargs={"k_edges": 5},
+        sampler_kwargs={"num_neighbors": [50] * 2},
+        local_filter=make_pcst_filter(triples, model))
+    """
+    approx precision = num_relevant_out_of_retrieved/num_retrieved_triples
+    We will use precision as a proxy for recall. This is because for recall,
+    we must know how many relevant triples exist for each question,
+    but this is not known.
+    """
+    precisions = []
+    for QA_pair in relevant_triplges.keys():
+        golden_triples = relevant_triples[QA_pair]
+        q = QA_pair[0]
+        retrieved_subgraph = query_loader.query(q)
+        retrieved_triples = retrieved_subgraph.triples
+
+        if args.verbose:
+            print("Q=", q)
+            print("A=", QA_pair[1])
+            print("retrieved_triples =", retrieved_triples)
+
+        num_relevant_out_of_retrieved = float(
+            sum([
+                int(bool(retrieved_triple in golden_triples))
+                for retrieved_triple in retrieved_triples
+            ]))
+        precisions.append(num_relevant_out_of_retrieved /
+                          len(retrieved_triples))
+    approx_precision = sum(precisions) / len(precisions)
+    print("approx_precision =", approx_precision)