We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
SNAPDataset ram usage
### 🐛 Describe the bug

Ran the following code on python 3.10/3.11 and the process got killed by the os(tried on windows/wsl/mac)
for using too much RAM (tried to run both on a laptop with 16gb of memory and a desktop pc with 64gb of memory).

```python
from torch_geometric.datasets import SNAPDataset
dataset = SNAPDataset("./datasets/snap/twitter", "ego-twitter")
```


### Versions

PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.0 (x86_64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.3.9.4)
CMake version: Could not collect
Libc version: N/A

Python version: 3.11.1 (main, Jan 26 2023, 14:19:45) [Clang 13.0.0 (clang-1300.0.29.30)] (64-bit runtime)
Python platform: macOS-14.0-x86_64-i386-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Intel(R) Core(TM) i5-1038NG7 CPU @ 2.00GHz

Versions of relevant libraries:
[pip3] mypy==0.991
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.24.1
[pip3] torch==2.0.1
[pip3] torch_geometric==2.4.0
[pip3] torchaudio==2.0.2
[pip3] torchvision==0.15.2
[conda] Could not collect

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torch_geometric/datasets/snap_dataset.py`
Content:
```
1 import os
2 import os.path as osp
3 from typing import Any, Callable, Dict, List, Optional, Union
4 
5 import fsspec
6 import numpy as np
7 import torch
8 
9 from torch_geometric.data import Data, InMemoryDataset
10 from torch_geometric.io import fs
11 from torch_geometric.utils import coalesce
12 
13 
14 class EgoData(Data):
15     def __inc__(self, key: str, value: Any, *args: Any, **kwargs: Any) -> Any:
16         if key == 'circle':
17             return self.num_nodes
18         elif key == 'circle_batch':
19             return int(value.max()) + 1 if value.numel() > 0 else 0
20         return super().__inc__(key, value, *args, **kwargs)
21 
22 
23 def read_ego(files: List[str], name: str) -> List[EgoData]:
24     import pandas as pd
25 
26     all_featnames = []
27     files = [
28         x for x in files if x.split('.')[-1] in
29         ['circles', 'edges', 'egofeat', 'feat', 'featnames']
30     ]
31     for i in range(4, len(files), 5):
32         featnames_file = files[i]
33         with fsspec.open(featnames_file, 'r') as f:
34             featnames = f.read().split('\n')[:-1]
35             featnames = [' '.join(x.split(' ')[1:]) for x in featnames]
36             all_featnames += featnames
37     all_featnames = sorted(list(set(all_featnames)))
38     all_featnames_dict = {key: i for i, key in enumerate(all_featnames)}
39 
40     data_list = []
41     for i in range(0, len(files), 5):
42         circles_file = files[i]
43         edges_file = files[i + 1]
44         egofeat_file = files[i + 2]
45         feat_file = files[i + 3]
46         featnames_file = files[i + 4]
47 
48         x = None
49         if name != 'gplus':  # Don't read node features on g-plus:
50             x_ego = pd.read_csv(egofeat_file, sep=' ', header=None,
51                                 dtype=np.float32)
52             x_ego = torch.from_numpy(x_ego.values)
53 
54             x = pd.read_csv(feat_file, sep=' ', header=None, dtype=np.float32)
55             x = torch.from_numpy(x.values)[:, 1:]
56 
57             x_all = torch.cat([x, x_ego], dim=0)
58 
59             # Reorder `x` according to `featnames` ordering.
60             x_all = torch.zeros(x.size(0), len(all_featnames))
61             with fsspec.open(featnames_file, 'r') as f:
62                 featnames = f.read().split('\n')[:-1]
63                 featnames = [' '.join(x.split(' ')[1:]) for x in featnames]
64             indices = [all_featnames_dict[featname] for featname in featnames]
65             x_all[:, torch.tensor(indices)] = x
66             x = x_all
67 
68         idx = pd.read_csv(feat_file, sep=' ', header=None, dtype=str,
69                           usecols=[0]).squeeze()
70 
71         idx_assoc: Dict[str, int] = {}
72         for i, j in enumerate(idx):
73             idx_assoc[j] = i
74 
75         circles: List[int] = []
76         circles_batch: List[int] = []
77         with fsspec.open(circles_file, 'r') as f:
78             for i, line in enumerate(f.read().split('\n')[:-1]):
79                 circle_indices = [idx_assoc[c] for c in line.split()[1:]]
80                 circles += circle_indices
81                 circles_batch += [i] * len(circle_indices)
82         circle = torch.tensor(circles)
83         circle_batch = torch.tensor(circles_batch)
84 
85         try:
86             row = pd.read_csv(edges_file, sep=' ', header=None, dtype=str,
87                               usecols=[0]).squeeze()
88             col = pd.read_csv(edges_file, sep=' ', header=None, dtype=str,
89                               usecols=[1]).squeeze()
90         except Exception:
91             continue
92 
93         row = torch.tensor([idx_assoc[i] for i in row])
94         col = torch.tensor([idx_assoc[i] for i in col])
95 
96         N = max(int(row.max()), int(col.max())) + 2
97         N = x.size(0) if x is not None else N
98 
99         row_ego = torch.full((N - 1, ), N - 1, dtype=torch.long)
100         col_ego = torch.arange(N - 1)
101 
102         # Ego node should be connected to every other node.
103         row = torch.cat([row, row_ego, col_ego], dim=0)
104         col = torch.cat([col, col_ego, row_ego], dim=0)
105         edge_index = torch.stack([row, col], dim=0)
106         edge_index = coalesce(edge_index, num_nodes=N)
107 
108         data = EgoData(x=x, edge_index=edge_index, circle=circle,
109                        circle_batch=circle_batch)
110 
111         data_list.append(data)
112 
113     return data_list
114 
115 
116 def read_soc(files: List[str], name: str) -> List[Data]:
117     import pandas as pd
118 
119     skiprows = 4
120     if name == 'pokec':
121         skiprows = 0
122 
123     edge_index = pd.read_csv(files[0], sep='\t', header=None,
124                              skiprows=skiprows, dtype=np.int64)
125     edge_index = torch.from_numpy(edge_index.values).t()
126     num_nodes = edge_index.max().item() + 1
127     edge_index = coalesce(edge_index, num_nodes=num_nodes)
128 
129     return [Data(edge_index=edge_index, num_nodes=num_nodes)]
130 
131 
132 def read_wiki(files: List[str], name: str) -> List[Data]:
133     import pandas as pd
134 
135     edge_index = pd.read_csv(files[0], sep='\t', header=None, skiprows=4,
136                              dtype=np.int64)
137     edge_index = torch.from_numpy(edge_index.values).t()
138 
139     idx = torch.unique(edge_index.flatten())
140     idx_assoc = torch.full((edge_index.max() + 1, ), -1, dtype=torch.long)
141     idx_assoc[idx] = torch.arange(idx.size(0))
142 
143     edge_index = idx_assoc[edge_index]
144     num_nodes = edge_index.max().item() + 1
145     edge_index = coalesce(edge_index, num_nodes=num_nodes)
146 
147     return [Data(edge_index=edge_index, num_nodes=num_nodes)]
148 
149 
150 class SNAPDataset(InMemoryDataset):
151     r"""A variety of graph datasets collected from `SNAP at Stanford University
152     <https://snap.stanford.edu/data>`_.
153 
154     Args:
155         root (str): Root directory where the dataset should be saved.
156         name (str): The name of the dataset.
157         transform (callable, optional): A function/transform that takes in an
158             :obj:`torch_geometric.data.Data` object and returns a transformed
159             version. The data object will be transformed before every access.
160             (default: :obj:`None`)
161         pre_transform (callable, optional): A function/transform that takes in
162             an :obj:`torch_geometric.data.Data` object and returns a
163             transformed version. The data object will be transformed before
164             being saved to disk. (default: :obj:`None`)
165         pre_filter (callable, optional): A function that takes in an
166             :obj:`torch_geometric.data.Data` object and returns a boolean
167             value, indicating whether the data object should be included in the
168             final dataset. (default: :obj:`None`)
169         force_reload (bool, optional): Whether to re-process the dataset.
170             (default: :obj:`False`)
171     """
172 
173     url = 'https://snap.stanford.edu/data'
174 
175     available_datasets = {
176         'ego-facebook': ['facebook.tar.gz'],
177         'ego-gplus': ['gplus.tar.gz'],
178         'ego-twitter': ['twitter.tar.gz'],
179         'soc-ca-astroph': ['ca-AstroPh.txt.gz'],
180         'soc-ca-grqc': ['ca-GrQc.txt.gz'],
181         'soc-epinions1': ['soc-Epinions1.txt.gz'],
182         'soc-livejournal1': ['soc-LiveJournal1.txt.gz'],
183         'soc-pokec': ['soc-pokec-relationships.txt.gz'],
184         'soc-slashdot0811': ['soc-Slashdot0811.txt.gz'],
185         'soc-slashdot0922': ['soc-Slashdot0902.txt.gz'],
186         'wiki-vote': ['wiki-Vote.txt.gz'],
187     }
188 
189     def __init__(
190         self,
191         root: str,
192         name: str,
193         transform: Optional[Callable] = None,
194         pre_transform: Optional[Callable] = None,
195         pre_filter: Optional[Callable] = None,
196         force_reload: bool = False,
197     ) -> None:
198         self.name = name.lower()
199         assert self.name in self.available_datasets.keys()
200         super().__init__(root, transform, pre_transform, pre_filter,
201                          force_reload=force_reload)
202         self.load(self.processed_paths[0])
203 
204     @property
205     def raw_dir(self) -> str:
206         return osp.join(self.root, self.name, 'raw')
207 
208     @property
209     def processed_dir(self) -> str:
210         return osp.join(self.root, self.name, 'processed')
211 
212     @property
213     def processed_file_names(self) -> str:
214         return 'data.pt'
215 
216     def _download(self) -> None:
217         if osp.isdir(self.raw_dir) and len(os.listdir(self.raw_dir)) > 0:
218             return
219 
220         fs.makedirs(self.raw_dir, exist_ok=True)
221         self.download()
222 
223     def download(self) -> None:
224         for name in self.available_datasets[self.name]:
225             fs.cp(f'{self.url}/{name}', self.raw_dir, extract=True)
226 
227     def process(self) -> None:
228         raw_dir = self.raw_dir
229         filenames = fs.ls(self.raw_dir)
230         if len(filenames) == 1 and fs.isdir(filenames[0]):
231             raw_dir = filenames[0]
232 
233         raw_files = fs.ls(raw_dir)
234 
235         data_list: Union[List[Data], List[EgoData]]
236         if self.name[:4] == 'ego-':
237             data_list = read_ego(raw_files, self.name[4:])
238         elif self.name[:4] == 'soc-':
239             data_list = read_soc(raw_files, self.name[:4])
240         elif self.name[:5] == 'wiki-':
241             data_list = read_wiki(raw_files, self.name[5:])
242         else:
243             raise NotImplementedError
244 
245         if len(data_list) > 1 and self.pre_filter is not None:
246             data_list = [data for data in data_list if self.pre_filter(data)]
247 
248         if self.pre_transform is not None:
249             data_list = [self.pre_transform(data) for data in data_list]
250 
251         self.save(data_list, self.processed_paths[0])
252 
253     def __repr__(self) -> str:
254         return f'SNAP-{self.name}({len(self)})'
255
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/torch_geometric/datasets/snap_dataset.py b/torch_geometric/datasets/snap_dataset.py
--- a/torch_geometric/datasets/snap_dataset.py
+++ b/torch_geometric/datasets/snap_dataset.py
@@ -22,6 +22,9 @@
 
 def read_ego(files: List[str], name: str) -> List[EgoData]:
     import pandas as pd
+    import tqdm
+
+    files = sorted(files)
 
     all_featnames = []
     files = [
@@ -38,7 +41,7 @@
     all_featnames_dict = {key: i for i, key in enumerate(all_featnames)}
 
     data_list = []
-    for i in range(0, len(files), 5):
+    for i in tqdm.tqdm(range(0, len(files), 5)):
         circles_file = files[i]
         edges_file = files[i + 1]
         egofeat_file = files[i + 2]
@@ -65,6 +68,9 @@
             x_all[:, torch.tensor(indices)] = x
             x = x_all
 
+            if x.size(1) > 100_000:
+                x = x.to_sparse_csr()
+
         idx = pd.read_csv(feat_file, sep=' ', header=None, dtype=str,
                           usecols=[0]).squeeze()

{"golden_diff": "diff --git a/torch_geometric/datasets/snap_dataset.py b/torch_geometric/datasets/snap_dataset.py\n--- a/torch_geometric/datasets/snap_dataset.py\n+++ b/torch_geometric/datasets/snap_dataset.py\n@@ -22,6 +22,9 @@\n \n def read_ego(files: List[str], name: str) -> List[EgoData]:\n     import pandas as pd\n+    import tqdm\n+\n+    files = sorted(files)\n \n     all_featnames = []\n     files = [\n@@ -38,7 +41,7 @@\n     all_featnames_dict = {key: i for i, key in enumerate(all_featnames)}\n \n     data_list = []\n-    for i in range(0, len(files), 5):\n+    for i in tqdm.tqdm(range(0, len(files), 5)):\n         circles_file = files[i]\n         edges_file = files[i + 1]\n         egofeat_file = files[i + 2]\n@@ -65,6 +68,9 @@\n             x_all[:, torch.tensor(indices)] = x\n             x = x_all\n \n+            if x.size(1) > 100_000:\n+                x = x.to_sparse_csr()\n+\n         idx = pd.read_csv(feat_file, sep=' ', header=None, dtype=str,\n                           usecols=[0]).squeeze()\n", "issue": "SNAPDataset ram usage\n### \ud83d\udc1b Describe the bug\n\nRan the following code on python 3.10/3.11 and the process got killed by the os(tried on windows/wsl/mac)\r\nfor using too much RAM (tried to run both on a laptop with 16gb of memory and a desktop pc with 64gb of memory).\r\n\r\n```python\r\nfrom torch_geometric.datasets import SNAPDataset\r\ndataset = SNAPDataset(\"./datasets/snap/twitter\", \"ego-twitter\")\r\n```\r\n\n\n### Versions\n\nPyTorch version: 2.0.1\r\nIs debug build: False\r\nCUDA used to build PyTorch: None\r\nROCM used to build PyTorch: N/A\r\n\r\nOS: macOS 14.0 (x86_64)\r\nGCC version: Could not collect\r\nClang version: 15.0.0 (clang-1500.3.9.4)\r\nCMake version: Could not collect\r\nLibc version: N/A\r\n\r\nPython version: 3.11.1 (main, Jan 26 2023, 14:19:45) [Clang 13.0.0 (clang-1300.0.29.30)] (64-bit runtime)\r\nPython platform: macOS-14.0-x86_64-i386-64bit\r\nIs CUDA available: False\r\nCUDA runtime version: No CUDA\r\nCUDA_MODULE_LOADING set to: N/A\r\nGPU models and configuration: No CUDA\r\nNvidia driver version: No CUDA\r\ncuDNN version: No CUDA\r\nHIP runtime version: N/A\r\nMIOpen runtime version: N/A\r\nIs XNNPACK available: True\r\n\r\nCPU:\r\nIntel(R) Core(TM) i5-1038NG7 CPU @ 2.00GHz\r\n\r\nVersions of relevant libraries:\r\n[pip3] mypy==0.991\r\n[pip3] mypy-extensions==0.4.3\r\n[pip3] numpy==1.24.1\r\n[pip3] torch==2.0.1\r\n[pip3] torch_geometric==2.4.0\r\n[pip3] torchaudio==2.0.2\r\n[pip3] torchvision==0.15.2\r\n[conda] Could not collect\n", "before_files": [{"content": "import os\nimport os.path as osp\nfrom typing import Any, Callable, Dict, List, Optional, Union\n\nimport fsspec\nimport numpy as np\nimport torch\n\nfrom torch_geometric.data import Data, InMemoryDataset\nfrom torch_geometric.io import fs\nfrom torch_geometric.utils import coalesce\n\n\nclass EgoData(Data):\n    def __inc__(self, key: str, value: Any, *args: Any, **kwargs: Any) -> Any:\n        if key == 'circle':\n            return self.num_nodes\n        elif key == 'circle_batch':\n            return int(value.max()) + 1 if value.numel() > 0 else 0\n        return super().__inc__(key, value, *args, **kwargs)\n\n\ndef read_ego(files: List[str], name: str) -> List[EgoData]:\n    import pandas as pd\n\n    all_featnames = []\n    files = [\n        x for x in files if x.split('.')[-1] in\n        ['circles', 'edges', 'egofeat', 'feat', 'featnames']\n    ]\n    for i in range(4, len(files), 5):\n        featnames_file = files[i]\n        with fsspec.open(featnames_file, 'r') as f:\n            featnames = f.read().split('\\n')[:-1]\n            featnames = [' '.join(x.split(' ')[1:]) for x in featnames]\n            all_featnames += featnames\n    all_featnames = sorted(list(set(all_featnames)))\n    all_featnames_dict = {key: i for i, key in enumerate(all_featnames)}\n\n    data_list = []\n    for i in range(0, len(files), 5):\n        circles_file = files[i]\n        edges_file = files[i + 1]\n        egofeat_file = files[i + 2]\n        feat_file = files[i + 3]\n        featnames_file = files[i + 4]\n\n        x = None\n        if name != 'gplus':  # Don't read node features on g-plus:\n            x_ego = pd.read_csv(egofeat_file, sep=' ', header=None,\n                                dtype=np.float32)\n            x_ego = torch.from_numpy(x_ego.values)\n\n            x = pd.read_csv(feat_file, sep=' ', header=None, dtype=np.float32)\n            x = torch.from_numpy(x.values)[:, 1:]\n\n            x_all = torch.cat([x, x_ego], dim=0)\n\n            # Reorder `x` according to `featnames` ordering.\n            x_all = torch.zeros(x.size(0), len(all_featnames))\n            with fsspec.open(featnames_file, 'r') as f:\n                featnames = f.read().split('\\n')[:-1]\n                featnames = [' '.join(x.split(' ')[1:]) for x in featnames]\n            indices = [all_featnames_dict[featname] for featname in featnames]\n            x_all[:, torch.tensor(indices)] = x\n            x = x_all\n\n        idx = pd.read_csv(feat_file, sep=' ', header=None, dtype=str,\n                          usecols=[0]).squeeze()\n\n        idx_assoc: Dict[str, int] = {}\n        for i, j in enumerate(idx):\n            idx_assoc[j] = i\n\n        circles: List[int] = []\n        circles_batch: List[int] = []\n        with fsspec.open(circles_file, 'r') as f:\n            for i, line in enumerate(f.read().split('\\n')[:-1]):\n                circle_indices = [idx_assoc[c] for c in line.split()[1:]]\n                circles += circle_indices\n                circles_batch += [i] * len(circle_indices)\n        circle = torch.tensor(circles)\n        circle_batch = torch.tensor(circles_batch)\n\n        try:\n            row = pd.read_csv(edges_file, sep=' ', header=None, dtype=str,\n                              usecols=[0]).squeeze()\n            col = pd.read_csv(edges_file, sep=' ', header=None, dtype=str,\n                              usecols=[1]).squeeze()\n        except Exception:\n            continue\n\n        row = torch.tensor([idx_assoc[i] for i in row])\n        col = torch.tensor([idx_assoc[i] for i in col])\n\n        N = max(int(row.max()), int(col.max())) + 2\n        N = x.size(0) if x is not None else N\n\n        row_ego = torch.full((N - 1, ), N - 1, dtype=torch.long)\n        col_ego = torch.arange(N - 1)\n\n        # Ego node should be connected to every other node.\n        row = torch.cat([row, row_ego, col_ego], dim=0)\n        col = torch.cat([col, col_ego, row_ego], dim=0)\n        edge_index = torch.stack([row, col], dim=0)\n        edge_index = coalesce(edge_index, num_nodes=N)\n\n        data = EgoData(x=x, edge_index=edge_index, circle=circle,\n                       circle_batch=circle_batch)\n\n        data_list.append(data)\n\n    return data_list\n\n\ndef read_soc(files: List[str], name: str) -> List[Data]:\n    import pandas as pd\n\n    skiprows = 4\n    if name == 'pokec':\n        skiprows = 0\n\n    edge_index = pd.read_csv(files[0], sep='\\t', header=None,\n                             skiprows=skiprows, dtype=np.int64)\n    edge_index = torch.from_numpy(edge_index.values).t()\n    num_nodes = edge_index.max().item() + 1\n    edge_index = coalesce(edge_index, num_nodes=num_nodes)\n\n    return [Data(edge_index=edge_index, num_nodes=num_nodes)]\n\n\ndef read_wiki(files: List[str], name: str) -> List[Data]:\n    import pandas as pd\n\n    edge_index = pd.read_csv(files[0], sep='\\t', header=None, skiprows=4,\n                             dtype=np.int64)\n    edge_index = torch.from_numpy(edge_index.values).t()\n\n    idx = torch.unique(edge_index.flatten())\n    idx_assoc = torch.full((edge_index.max() + 1, ), -1, dtype=torch.long)\n    idx_assoc[idx] = torch.arange(idx.size(0))\n\n    edge_index = idx_assoc[edge_index]\n    num_nodes = edge_index.max().item() + 1\n    edge_index = coalesce(edge_index, num_nodes=num_nodes)\n\n    return [Data(edge_index=edge_index, num_nodes=num_nodes)]\n\n\nclass SNAPDataset(InMemoryDataset):\n    r\"\"\"A variety of graph datasets collected from `SNAP at Stanford University\n    <https://snap.stanford.edu/data>`_.\n\n    Args:\n        root (str): Root directory where the dataset should be saved.\n        name (str): The name of the dataset.\n        transform (callable, optional): A function/transform that takes in an\n            :obj:`torch_geometric.data.Data` object and returns a transformed\n            version. The data object will be transformed before every access.\n            (default: :obj:`None`)\n        pre_transform (callable, optional): A function/transform that takes in\n            an :obj:`torch_geometric.data.Data` object and returns a\n            transformed version. The data object will be transformed before\n            being saved to disk. (default: :obj:`None`)\n        pre_filter (callable, optional): A function that takes in an\n            :obj:`torch_geometric.data.Data` object and returns a boolean\n            value, indicating whether the data object should be included in the\n            final dataset. (default: :obj:`None`)\n        force_reload (bool, optional): Whether to re-process the dataset.\n            (default: :obj:`False`)\n    \"\"\"\n\n    url = 'https://snap.stanford.edu/data'\n\n    available_datasets = {\n        'ego-facebook': ['facebook.tar.gz'],\n        'ego-gplus': ['gplus.tar.gz'],\n        'ego-twitter': ['twitter.tar.gz'],\n        'soc-ca-astroph': ['ca-AstroPh.txt.gz'],\n        'soc-ca-grqc': ['ca-GrQc.txt.gz'],\n        'soc-epinions1': ['soc-Epinions1.txt.gz'],\n        'soc-livejournal1': ['soc-LiveJournal1.txt.gz'],\n        'soc-pokec': ['soc-pokec-relationships.txt.gz'],\n        'soc-slashdot0811': ['soc-Slashdot0811.txt.gz'],\n        'soc-slashdot0922': ['soc-Slashdot0902.txt.gz'],\n        'wiki-vote': ['wiki-Vote.txt.gz'],\n    }\n\n    def __init__(\n        self,\n        root: str,\n        name: str,\n        transform: Optional[Callable] = None,\n        pre_transform: Optional[Callable] = None,\n        pre_filter: Optional[Callable] = None,\n        force_reload: bool = False,\n    ) -> None:\n        self.name = name.lower()\n        assert self.name in self.available_datasets.keys()\n        super().__init__(root, transform, pre_transform, pre_filter,\n                         force_reload=force_reload)\n        self.load(self.processed_paths[0])\n\n    @property\n    def raw_dir(self) -> str:\n        return osp.join(self.root, self.name, 'raw')\n\n    @property\n    def processed_dir(self) -> str:\n        return osp.join(self.root, self.name, 'processed')\n\n    @property\n    def processed_file_names(self) -> str:\n        return 'data.pt'\n\n    def _download(self) -> None:\n        if osp.isdir(self.raw_dir) and len(os.listdir(self.raw_dir)) > 0:\n            return\n\n        fs.makedirs(self.raw_dir, exist_ok=True)\n        self.download()\n\n    def download(self) -> None:\n        for name in self.available_datasets[self.name]:\n            fs.cp(f'{self.url}/{name}', self.raw_dir, extract=True)\n\n    def process(self) -> None:\n        raw_dir = self.raw_dir\n        filenames = fs.ls(self.raw_dir)\n        if len(filenames) == 1 and fs.isdir(filenames[0]):\n            raw_dir = filenames[0]\n\n        raw_files = fs.ls(raw_dir)\n\n        data_list: Union[List[Data], List[EgoData]]\n        if self.name[:4] == 'ego-':\n            data_list = read_ego(raw_files, self.name[4:])\n        elif self.name[:4] == 'soc-':\n            data_list = read_soc(raw_files, self.name[:4])\n        elif self.name[:5] == 'wiki-':\n            data_list = read_wiki(raw_files, self.name[5:])\n        else:\n            raise NotImplementedError\n\n        if len(data_list) > 1 and self.pre_filter is not None:\n            data_list = [data for data in data_list if self.pre_filter(data)]\n\n        if self.pre_transform is not None:\n            data_list = [self.pre_transform(data) for data in data_list]\n\n        self.save(data_list, self.processed_paths[0])\n\n    def __repr__(self) -> str:\n        return f'SNAP-{self.name}({len(self)})'\n", "path": "torch_geometric/datasets/snap_dataset.py"}], "after_files": [{"content": "import os\nimport os.path as osp\nfrom typing import Any, Callable, Dict, List, Optional, Union\n\nimport fsspec\nimport numpy as np\nimport torch\n\nfrom torch_geometric.data import Data, InMemoryDataset\nfrom torch_geometric.io import fs\nfrom torch_geometric.utils import coalesce\n\n\nclass EgoData(Data):\n    def __inc__(self, key: str, value: Any, *args: Any, **kwargs: Any) -> Any:\n        if key == 'circle':\n            return self.num_nodes\n        elif key == 'circle_batch':\n            return int(value.max()) + 1 if value.numel() > 0 else 0\n        return super().__inc__(key, value, *args, **kwargs)\n\n\ndef read_ego(files: List[str], name: str) -> List[EgoData]:\n    import pandas as pd\n    import tqdm\n\n    files = sorted(files)\n\n    all_featnames = []\n    files = [\n        x for x in files if x.split('.')[-1] in\n        ['circles', 'edges', 'egofeat', 'feat', 'featnames']\n    ]\n    for i in range(4, len(files), 5):\n        featnames_file = files[i]\n        with fsspec.open(featnames_file, 'r') as f:\n            featnames = f.read().split('\\n')[:-1]\n            featnames = [' '.join(x.split(' ')[1:]) for x in featnames]\n            all_featnames += featnames\n    all_featnames = sorted(list(set(all_featnames)))\n    all_featnames_dict = {key: i for i, key in enumerate(all_featnames)}\n\n    data_list = []\n    for i in tqdm.tqdm(range(0, len(files), 5)):\n        circles_file = files[i]\n        edges_file = files[i + 1]\n        egofeat_file = files[i + 2]\n        feat_file = files[i + 3]\n        featnames_file = files[i + 4]\n\n        x = None\n        if name != 'gplus':  # Don't read node features on g-plus:\n            x_ego = pd.read_csv(egofeat_file, sep=' ', header=None,\n                                dtype=np.float32)\n            x_ego = torch.from_numpy(x_ego.values)\n\n            x = pd.read_csv(feat_file, sep=' ', header=None, dtype=np.float32)\n            x = torch.from_numpy(x.values)[:, 1:]\n\n            x_all = torch.cat([x, x_ego], dim=0)\n\n            # Reorder `x` according to `featnames` ordering.\n            x_all = torch.zeros(x.size(0), len(all_featnames))\n            with fsspec.open(featnames_file, 'r') as f:\n                featnames = f.read().split('\\n')[:-1]\n                featnames = [' '.join(x.split(' ')[1:]) for x in featnames]\n            indices = [all_featnames_dict[featname] for featname in featnames]\n            x_all[:, torch.tensor(indices)] = x\n            x = x_all\n\n            if x.size(1) > 100_000:\n                x = x.to_sparse_csr()\n\n        idx = pd.read_csv(feat_file, sep=' ', header=None, dtype=str,\n                          usecols=[0]).squeeze()\n\n        idx_assoc: Dict[str, int] = {}\n        for i, j in enumerate(idx):\n            idx_assoc[j] = i\n\n        circles: List[int] = []\n        circles_batch: List[int] = []\n        with fsspec.open(circles_file, 'r') as f:\n            for i, line in enumerate(f.read().split('\\n')[:-1]):\n                circle_indices = [idx_assoc[c] for c in line.split()[1:]]\n                circles += circle_indices\n                circles_batch += [i] * len(circle_indices)\n        circle = torch.tensor(circles)\n        circle_batch = torch.tensor(circles_batch)\n\n        try:\n            row = pd.read_csv(edges_file, sep=' ', header=None, dtype=str,\n                              usecols=[0]).squeeze()\n            col = pd.read_csv(edges_file, sep=' ', header=None, dtype=str,\n                              usecols=[1]).squeeze()\n        except Exception:\n            continue\n\n        row = torch.tensor([idx_assoc[i] for i in row])\n        col = torch.tensor([idx_assoc[i] for i in col])\n\n        N = max(int(row.max()), int(col.max())) + 2\n        N = x.size(0) if x is not None else N\n\n        row_ego = torch.full((N - 1, ), N - 1, dtype=torch.long)\n        col_ego = torch.arange(N - 1)\n\n        # Ego node should be connected to every other node.\n        row = torch.cat([row, row_ego, col_ego], dim=0)\n        col = torch.cat([col, col_ego, row_ego], dim=0)\n        edge_index = torch.stack([row, col], dim=0)\n        edge_index = coalesce(edge_index, num_nodes=N)\n\n        data = EgoData(x=x, edge_index=edge_index, circle=circle,\n                       circle_batch=circle_batch)\n\n        data_list.append(data)\n\n    return data_list\n\n\ndef read_soc(files: List[str], name: str) -> List[Data]:\n    import pandas as pd\n\n    skiprows = 4\n    if name == 'pokec':\n        skiprows = 0\n\n    edge_index = pd.read_csv(files[0], sep='\\t', header=None,\n                             skiprows=skiprows, dtype=np.int64)\n    edge_index = torch.from_numpy(edge_index.values).t()\n    num_nodes = edge_index.max().item() + 1\n    edge_index = coalesce(edge_index, num_nodes=num_nodes)\n\n    return [Data(edge_index=edge_index, num_nodes=num_nodes)]\n\n\ndef read_wiki(files: List[str], name: str) -> List[Data]:\n    import pandas as pd\n\n    edge_index = pd.read_csv(files[0], sep='\\t', header=None, skiprows=4,\n                             dtype=np.int64)\n    edge_index = torch.from_numpy(edge_index.values).t()\n\n    idx = torch.unique(edge_index.flatten())\n    idx_assoc = torch.full((edge_index.max() + 1, ), -1, dtype=torch.long)\n    idx_assoc[idx] = torch.arange(idx.size(0))\n\n    edge_index = idx_assoc[edge_index]\n    num_nodes = edge_index.max().item() + 1\n    edge_index = coalesce(edge_index, num_nodes=num_nodes)\n\n    return [Data(edge_index=edge_index, num_nodes=num_nodes)]\n\n\nclass SNAPDataset(InMemoryDataset):\n    r\"\"\"A variety of graph datasets collected from `SNAP at Stanford University\n    <https://snap.stanford.edu/data>`_.\n\n    Args:\n        root (str): Root directory where the dataset should be saved.\n        name (str): The name of the dataset.\n        transform (callable, optional): A function/transform that takes in an\n            :obj:`torch_geometric.data.Data` object and returns a transformed\n            version. The data object will be transformed before every access.\n            (default: :obj:`None`)\n        pre_transform (callable, optional): A function/transform that takes in\n            an :obj:`torch_geometric.data.Data` object and returns a\n            transformed version. The data object will be transformed before\n            being saved to disk. (default: :obj:`None`)\n        pre_filter (callable, optional): A function that takes in an\n            :obj:`torch_geometric.data.Data` object and returns a boolean\n            value, indicating whether the data object should be included in the\n            final dataset. (default: :obj:`None`)\n        force_reload (bool, optional): Whether to re-process the dataset.\n            (default: :obj:`False`)\n    \"\"\"\n\n    url = 'https://snap.stanford.edu/data'\n\n    available_datasets = {\n        'ego-facebook': ['facebook.tar.gz'],\n        'ego-gplus': ['gplus.tar.gz'],\n        'ego-twitter': ['twitter.tar.gz'],\n        'soc-ca-astroph': ['ca-AstroPh.txt.gz'],\n        'soc-ca-grqc': ['ca-GrQc.txt.gz'],\n        'soc-epinions1': ['soc-Epinions1.txt.gz'],\n        'soc-livejournal1': ['soc-LiveJournal1.txt.gz'],\n        'soc-pokec': ['soc-pokec-relationships.txt.gz'],\n        'soc-slashdot0811': ['soc-Slashdot0811.txt.gz'],\n        'soc-slashdot0922': ['soc-Slashdot0902.txt.gz'],\n        'wiki-vote': ['wiki-Vote.txt.gz'],\n    }\n\n    def __init__(\n        self,\n        root: str,\n        name: str,\n        transform: Optional[Callable] = None,\n        pre_transform: Optional[Callable] = None,\n        pre_filter: Optional[Callable] = None,\n        force_reload: bool = False,\n    ) -> None:\n        self.name = name.lower()\n        assert self.name in self.available_datasets.keys()\n        super().__init__(root, transform, pre_transform, pre_filter,\n                         force_reload=force_reload)\n        self.load(self.processed_paths[0])\n\n    @property\n    def raw_dir(self) -> str:\n        return osp.join(self.root, self.name, 'raw')\n\n    @property\n    def processed_dir(self) -> str:\n        return osp.join(self.root, self.name, 'processed')\n\n    @property\n    def processed_file_names(self) -> str:\n        return 'data.pt'\n\n    def _download(self) -> None:\n        if osp.isdir(self.raw_dir) and len(os.listdir(self.raw_dir)) > 0:\n            return\n\n        fs.makedirs(self.raw_dir, exist_ok=True)\n        self.download()\n\n    def download(self) -> None:\n        for name in self.available_datasets[self.name]:\n            fs.cp(f'{self.url}/{name}', self.raw_dir, extract=True)\n\n    def process(self) -> None:\n        raw_dir = self.raw_dir\n        filenames = fs.ls(self.raw_dir)\n        if len(filenames) == 1 and fs.isdir(filenames[0]):\n            raw_dir = filenames[0]\n\n        raw_files = fs.ls(raw_dir)\n\n        data_list: Union[List[Data], List[EgoData]]\n        if self.name[:4] == 'ego-':\n            data_list = read_ego(raw_files, self.name[4:])\n        elif self.name[:4] == 'soc-':\n            data_list = read_soc(raw_files, self.name[:4])\n        elif self.name[:5] == 'wiki-':\n            data_list = read_wiki(raw_files, self.name[5:])\n        else:\n            raise NotImplementedError\n\n        if len(data_list) > 1 and self.pre_filter is not None:\n            data_list = [data for data in data_list if self.pre_filter(data)]\n\n        if self.pre_transform is not None:\n            data_list = [self.pre_transform(data) for data in data_list]\n\n        self.save(data_list, self.processed_paths[0])\n\n    def __repr__(self) -> str:\n        return f'SNAP-{self.name}({len(self)})'\n", "path": "torch_geometric/datasets/snap_dataset.py"}]}