From 647d8ea8159c24d65ed6268f6a9085107674ecf5 Mon Sep 17 00:00:00 2001
From: Muhammed Fatih BALIN <m.f.balin@gmail.com>
Date: Thu, 29 Aug 2024 19:52:56 -0400
Subject: [PATCH] [GraphBolt][PyG] Add igb datasets to the examples. (#7758)

---
 examples/graphbolt/pyg/labor/load_dataset.py         |  9 +--------
 examples/graphbolt/pyg/labor/node_classification.py  |  3 +++
 .../graphbolt/pyg/node_classification_advanced.py    | 12 ++++++++++--
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/examples/graphbolt/pyg/labor/load_dataset.py b/examples/graphbolt/pyg/labor/load_dataset.py
index 0f9f24da4fc4..c354b91f6692 100644
--- a/examples/graphbolt/pyg/labor/load_dataset.py
+++ b/examples/graphbolt/pyg/labor/load_dataset.py
@@ -39,12 +39,7 @@ def load_dataset(dataset_name, disk_based_feature_keys=None):
         "flickr",
     ]:
         dataset, multilabel = load_dgl(dataset_name)
-    elif dataset_name in [
-        "ogbn-products",
-        "ogbn-arxiv",
-        "ogbn-papers100M",
-        "ogbn-mag240M",
-    ]:
+    else:
         if "mag240M" in dataset_name:
             dataset_name = "ogb-lsc-mag240m"
         dataset = gb.BuiltinDataset(dataset_name)
@@ -56,7 +51,5 @@ def load_dataset(dataset_name, disk_based_feature_keys=None):
             if feature_key in disk_based_feature_keys:
                 feature["in_memory"] = False
         dataset = dataset.load()
-    else:
-        raise ValueError("unknown dataset")
 
     return dataset, multilabel
diff --git a/examples/graphbolt/pyg/labor/node_classification.py b/examples/graphbolt/pyg/labor/node_classification.py
index f6636d2721b7..7129d517c8c6 100644
--- a/examples/graphbolt/pyg/labor/node_classification.py
+++ b/examples/graphbolt/pyg/labor/node_classification.py
@@ -363,6 +363,9 @@ def parse_args():
             "ogbn-arxiv",
             "ogbn-products",
             "ogbn-papers100M",
+            "igb-hom-tiny",
+            "igb-hom-small",
+            "igb-hom-medium",
             "reddit",
             "yelp",
             "flickr",
diff --git a/examples/graphbolt/pyg/node_classification_advanced.py b/examples/graphbolt/pyg/node_classification_advanced.py
index 02df19f6fd38..576335b5bdf1 100644
--- a/examples/graphbolt/pyg/node_classification_advanced.py
+++ b/examples/graphbolt/pyg/node_classification_advanced.py
@@ -335,9 +335,17 @@ def parse_args():
         "--dataset",
         type=str,
         default="ogbn-products",
-        choices=["ogbn-arxiv", "ogbn-products", "ogbn-papers100M"],
+        choices=[
+            "ogbn-arxiv",
+            "ogbn-products",
+            "ogbn-papers100M",
+            "igb-hom-tiny",
+            "igb-hom-small",
+            "igb-hom-medium",
+        ],
         help="The dataset we can use for node classification example. Currently"
-        " ogbn-products, ogbn-arxiv, ogbn-papers100M datasets are supported.",
+        " ogbn-products, ogbn-arxiv, ogbn-papers100M and"
+        " igb-hom-[tiny|small|medium] datasets are supported.",
     )
     parser.add_argument(
         "--fanout",