add yelp3k dataset

Author: yifanfeng97

dhg/data/__init__.py

@@ -1,7 +1,7 @@
 from .base import BaseData
 from .cooking_200 import Cooking200
 from .movielens import MovieLens1M
-from .yelp import Yelp2018, YelpRestaurant
+from .yelp import Yelp2018, YelpRestaurant, Yelp3k
 from .gowalla import Gowalla
 from .amazon import AmazonBook
 from .walmart import WalmartTrips
@@ -52,5 +52,6 @@
     "News20",
     "IMDB4k",
     "Recipe100k",
-    "Recipe200k"
+    "Recipe200k",
+    "Yelp3k"
 ]

dhg/data/yelp.py

@@ -100,3 +100,47 @@ def __init__(self, data_root: Optional[str] = None) -> None:
             },
         }
 
+
+
+class Yelp3k(BaseData):
+    r"""The Yelp3k dataset is a subset of Yelp-Restaurant dataset for vertex classification task.  
+    It is a restaurant-review network. All businesses in the “restaurant” catalog are selected as our nodes, 
+    and formed hyperedges by selecting restaurants visited by the same user. 
+    We use the state of the business as the corresponding node label.
+
+    The content of the Yelp-Restaurant dataset includes the following:
+
+    - ``num_classes``: The number of classes: :math:`6`.
+    - ``num_vertices``: The number of vertices: :math:`3,855`.     
+    - ``num_edges``: The number of edges: :math:`24,137`.
+    - ``dim_features``: The dimension of features: :math:`1,862`.
+    - ``features``: The vertex feature matrix. ``torch.Tensor`` with size :math:`(3,855 \times 1,862)`.
+    - ``edge_list``: The edge list. ``List`` with length :math:`24,137`.
+    - ``labels``: The label list. ``torch.LongTensor`` with size :math:`(3,855, )`.
+
+    Args:
+        ``data_root`` (``str``, optional): The ``data_root`` has stored the data. If set to ``None``, this function will auto-download from server and save into the default direction ``~/.dhg/datasets/``. Defaults to ``None``.
+    """
+
+    def __init__(self, data_root: Optional[str] = None) -> None:
+        super().__init__("yelp_3k", data_root)
+        self._content = {
+            "num_classes": 6,
+            "num_vertices": 3855,
+            "num_edges": 24137,
+            "dim_features": 1862,
+            "features": {
+                "upon": [{"filename": "features.pkl", "md5": "b107876d2e9ba5b0f7f67175e36bfdf4"}],
+                "loader": load_from_pickle,
+                "preprocess": [to_tensor,],
+            },
+            "edge_list": {
+                "upon": [{"filename": "edge_list.pkl", "md5": "322cfd002c0f557cf7bc4b9b2654b987"}],
+                "loader": load_from_pickle,
+            },
+            "labels": {
+                "upon": [{"filename": "labels.pkl", "md5": "16ead8bcbb8a3621202c88f64c8cbc07"}],
+                "loader": load_from_pickle,
+                "preprocess": [to_long_tensor],
+            },
+        }

docs/source/api/data.rst

@@ -66,6 +66,7 @@ Hypergraph Datasets
     dhg.data.IMDB4k
     dhg.data.Recipe100k
     dhg.data.Recipe200k
+    dhg.data.Yelp3k
 
 
 **Welcome to contribute datasets!**