Add no PCST option in WebQSPDataset (#9722)

This PR adds an option to disable PCST in WebQSPDataset, it allows the user to use some custom retrieval method (such as GNN-RAG). --------- Co-authored-by: Matthias Fey <[email protected]>
pyg-team · Oct 28, 2024 · b823c7e · b823c7e
1 parent 78e3f39
commit b823c7e
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 15 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added the `use_pcst` option to `WebQSPDataset` ([#9722](https://github.com/pyg-team/pytorch_geometric/pull/9722))
 - Allowed users to pass `edge_weight` to `GraphUNet` models ([#9737](https://github.com/pyg-team/pytorch_geometric/pull/9737))
 - Consolidated `examples/ogbn_{papers_100m,products_gat,products_sage}.py` into `examples/ogbn_train.py` ([#9467](https://github.com/pyg-team/pytorch_geometric/pull/9467))
 

diff --git a/torch_geometric/datasets/web_qsp_dataset.py b/torch_geometric/datasets/web_qsp_dataset.py
@@ -21,12 +21,6 @@ def retrieval_via_pcst(
     cost_e: float = 0.5,
 ) -> Tuple[Data, str]:
     c = 0.01
-    if len(textual_nodes) == 0 or len(textual_edges) == 0:
-        desc = textual_nodes.to_csv(index=False) + "\n" + textual_edges.to_csv(
-            index=False,
-            columns=["src", "edge_attr", "dst"],
-        )
-        return data, desc
 
     from pcst_fast import pcst_fast
 
@@ -135,13 +129,17 @@ class WebQSPDataset(InMemoryDataset):
             If :obj:`"test"`, loads the test dataset. (default: :obj:`"train"`)
         force_reload (bool, optional): Whether to re-process the dataset.
             (default: :obj:`False`)
+        use_pcst (bool, optional): Whether to preprocess the dataset's graph
+            with PCST or return the full graphs. (default: :obj:`True`)
     """
     def __init__(
         self,
         root: str,
         split: str = "train",
         force_reload: bool = False,
+        use_pcst: bool = True,
     ) -> None:
+        self.use_pcst = use_pcst
         super().__init__(root, force_reload=force_reload)
 
         if split not in {'train', 'val', 'test'}:
@@ -224,15 +222,22 @@ def process(self) -> None:
                     edge_index=edge_index,
                     edge_attr=edge_attr,
                 )
-                data, desc = retrieval_via_pcst(
-                    data,
-                    question_embs[i],
-                    nodes,
-                    edges,
-                    topk=3,
-                    topk_e=5,
-                    cost_e=0.5,
-                )
+                if self.use_pcst and len(nodes) > 0 and len(edges) > 0:
+                    data, desc = retrieval_via_pcst(
+                        data,
+                        question_embs[i],
+                        nodes,
+                        edges,
+                        topk=3,
+                        topk_e=5,
+                        cost_e=0.5,
+                    )
+                else:
+                    desc = nodes.to_csv(index=False) + "\n" + edges.to_csv(
+                        index=False,
+                        columns=["src", "edge_attr", "dst"],
+                    )
+
                 data.question = question
                 data.label = label
                 data.desc = desc