Add llm generated explanations to TAGDataset

CHANGELOG.md

@@ -7,6 +7,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added llm generated explanations to `TAGDataset` ([#9918](https://github.com/pyg-team/pytorch_geometric/pull/9918))
+- Update Dockerfile to use latest from NVIDIA ([#9794](https://github.com/pyg-team/pytorch_geometric/pull/9794))
 - Consolidate Cugraph examples into ogbn_train_cugraph.py and ogbn_train_cugraph_multigpu.py for ogbn-arxiv, ogbn-products and ogbn-papers100M ([#9953](https://github.com/pyg-team/pytorch_geometric/pull/9953))
 - Added `InstructMol` dataset ([#9975](https://github.com/pyg-team/pytorch_geometric/pull/9975))
 - Added support for weighted `LinkPredRecall` metric ([#9947](https://github.com/pyg-team/pytorch_geometric/pull/9947))

examples/llm/glem.py

@@ -40,6 +40,7 @@ def get_n_params(model):
 def main(args):
     gpu = args.gpu
     dataset_name = args.dataset
+    text_type = args.text_type if args.dataset == 'arxiv' else 'raw_text'
     root = osp.join('data', 'ogb')
     hf_model = args.hf_model
     pl_ratio = args.pl_ratio
@@ -83,7 +84,7 @@ def main(args):
 
     tag_dataset = TAGDataset(root, dataset, hf_model,
                              token_on_disk=token_on_disk)
-    text_dataset = tag_dataset.to_text_dataset()
+    text_dataset = tag_dataset.to_text_dataset(text_type)
     print(tag_dataset.num_classes, tag_dataset.raw_file_names)
 
     num_classes = tag_dataset.num_classes
@@ -393,8 +394,12 @@ def load_model(em_phase):
                         help='number of runs')
     parser.add_argument('--num_em_iters', type=int, default=1,
                         help='number of iterations')
-    parser.add_argument("--dataset", type=str, default='products',
+    parser.add_argument("--dataset", type=str, default='arxiv',
                         help='arxiv or products')
+    parser.add_argument(
+        "--text_type", type=str, default='llm_explanation',
+        help="type of text, support raw_text, llm_explanation,"
+        "all for arxiv and raw_text for products")
     parser.add_argument("--pl_ratio", type=float, default=0.5,
                         help="pseudo labels ratio")
     parser.add_argument('--hf_model', type=str, default='prajjwal1/bert-tiny',

test/datasets/test_tag_dataset.py

@@ -0,0 +1,22 @@
+from torch_geometric.datasets import TAGDataset
+from torch_geometric.testing import onlyFullTest, withPackage
+
+
+@onlyFullTest
+@withPackage('ogb')
+def test_tag_dataset() -> None:
+    from ogb.nodeproppred import PygNodePropPredDataset
+
+    root = './data/ogb'
+    hf_model = 'prajjwal1/bert-tiny'
+    token_on_disk = True
+
+    dataset = PygNodePropPredDataset('ogbn-arxiv', root=root)
+    tag_dataset = TAGDataset(root, dataset, hf_model,
+                             token_on_disk=token_on_disk)
+
+    assert 169343 == tag_dataset[0].num_nodes \
+        == len(tag_dataset.text) \
+        == len(tag_dataset.llm_explanation) \
+        == len(tag_dataset.llm_prediction)
+    assert 1166243 == tag_dataset[0].num_edges