Adding 'wiki' dataset to LINKXDataset

CHANGELOG.md

@@ -5,7 +5,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [2.0.5] - 2022-MM-DD
 ### Added
-- Added the `Genius` datasets to `nn.datasets.LINKXDataset` ([#4570](https://github.com/pyg-team/pytorch_geometric/pull/4570))
+- Added the `Genius` and `Wiki` datasets to `nn.datasets.LINKXDataset` ([#4570](https://github.com/pyg-team/pytorch_geometric/pull/4570), [#4600](https://github.com/pyg-team/pytorch_geometric/pull/4600))
 - Added `nn.glob.GlobalPooling` module with support for multiple aggregations ([#4582](https://github.com/pyg-team/pytorch_geometric/pull/4582))
 - Added support for graph-level outputs in `to_hetero` ([#4582](https://github.com/pyg-team/pytorch_geometric/pull/4582))
 - Added `CHANGELOG.md` ([#4581](https://github.com/pyg-team/pytorch_geometric/pull/4581))

torch_geometric/data/download.py

@@ -2,11 +2,13 @@
 import ssl
 import sys
 import urllib
+from typing import Optional
 
 from .makedirs import makedirs
 
 
-def download_url(url: str, folder: str, log: bool = True):
+def download_url(url: str, folder: str, log: bool = True,
+                 filename: Optional[str] = None):
     r"""Downloads the content of an URL to a specific folder.
 
     Args:
@@ -16,8 +18,10 @@ def download_url(url: str, folder: str, log: bool = True):
             console. (default: :obj:`True`)
     """
 
-    filename = url.rpartition('/')[2]
-    filename = filename if filename[0] == '?' else filename.split('?')[0]
+    if filename is None:
+        filename = url.rpartition('/')[2]
+        filename = filename if filename[0] == '?' else filename.split('?')[0]
+
     path = osp.join(folder, filename)
 
     if osp.exists(path):  # pragma: no cover
@@ -34,6 +38,11 @@ def download_url(url: str, folder: str, log: bool = True):
     data = urllib.request.urlopen(url, context=context)
 
     with open(path, 'wb') as f:
-        f.write(data.read())
+        # workaround for https://bugs.python.org/issue42853
+        while True:
+            chunk = data.read(10 * 1024 * 1024)
+            if not chunk:
+                break
+            f.write(chunk)
 
     return path

torch_geometric/datasets/linkx_dataset.py

@@ -32,23 +32,45 @@ class LINKXDataset(InMemoryDataset):
             being saved to disk. (default: :obj:`None`)
     """
 
-    url = 'https://github.com/CUAI/Non-Homophily-Large-Scale/raw/master/data'
+    github_url = 'https://github.com/CUAI/Non-Homophily-Large-Scale/' \
+                 'raw/master/data'
+    gdrive_url = 'https://drive.google.com/uc?confirm=t&'
 
     facebook_datasets = [
         'penn94', 'reed98', 'amherst41', 'cornell5', 'johnshopkins55'
     ]
 
     datasets = {
-        'penn94': f'{url}/facebook100/Penn94.mat',
-        'reed98': f'{url}/facebook100/Reed98.mat',
-        'amherst41': f'{url}/facebook100/Amherst41.mat',
-        'cornell5': f'{url}/facebook100/Cornell5.mat',
-        'johnshopkins55': f'{url}/facebook100/Johns%20Hopkins55.mat',
-        'genius': f'{url}/genius.mat'
+        'penn94': {
+            'data.mat': f'{github_url}/facebook100/Penn94.mat'
+        },
+        'reed98': {
+            'data.mat': f'{github_url}/facebook100/Reed98.mat'
+        },
+        'amherst41': {
+            'data.mat': f'{github_url}/facebook100/Amherst41.mat',
+        },
+        'cornell5': {
+            'data.mat': f'{github_url}/facebook100/Cornell5.mat'
+        },
+        'johnshopkins55': {
+            'data.mat': f'{github_url}/facebook100/Johns%20Hopkins55.mat'
+        },
+        'genius': {
+            'data.mat': f'{github_url}/genius.mat'
+        },
+        'wiki': {
+            'wiki_views2M.pt':
+            f'{gdrive_url}id=1p5DlVHrnFgYm3VsNIzahSsvCD424AyvP',
+            'wiki_edges2M.pt':
+            f'{gdrive_url}id=14X7FlkjrlUgmnsYtPwdh-gGuFla4yb5u',
+            'wiki_features2M.pt':
+            f'{gdrive_url}id=1ySNspxbK-snNoAZM7oxiWGvOnTRdSyEK'
+        }
     }
 
     splits = {
-        'penn94': f'{url}/splits/fb100-Penn94-splits.npy',
+        'penn94': f'{github_url}/splits/fb100-Penn94-splits.npy',
     }
 
     def __init__(self, root: str, name: str,
@@ -69,7 +91,7 @@ def processed_dir(self) -> str:
 
     @property
     def raw_file_names(self) -> List[str]:
-        names = [self.datasets[self.name].split('/')[-1]]
+        names = list(self.datasets[self.name].keys())
         if self.name in self.splits:
             names += [self.splits[self.name].split('/')[-1]]
         return names
@@ -79,10 +101,21 @@ def processed_file_names(self) -> str:
         return 'data.pt'
 
     def download(self):
-        download_url(self.datasets[self.name], self.raw_dir)
+        for filename, path in self.datasets[self.name].items():
+            download_url(path, self.raw_dir, filename=filename)
         if self.name in self.splits:
             download_url(self.splits[self.name], self.raw_dir)
 
+    def _process_wiki(self):
+
+        paths = {x.split("/")[-1]: x for x in self.raw_paths}
+        print(paths)
+        x = torch.load(paths['wiki_features2M.pt'])
+        edge_index = torch.load(paths['wiki_edges2M.pt']).T
+        y = torch.load(paths['wiki_views2M.pt'])
+
+        return Data(x=x, edge_index=edge_index, y=y)
+
     def _process_facebook(self):
         from scipy.io import loadmat
 
@@ -134,8 +167,11 @@ def process(self):
             data = self._process_facebook()
         elif self.name == 'genius':
             data = self._process_genius()
+        elif self.name == 'wiki':
+            data = self._process_wiki()
         else:
-            raise NotImplementedError
+            raise NotImplementedError(
+                f"chosen dataset '{self.name}' is not implemented")
 
         if self.pre_transform is not None:
             data = self.pre_transform(data)