pyg-nightly 2.6.0.dev20240319__py3-none-any.whl → 2.7.0.dev20250114__py3-none-any.whl

torch_geometric/datasets/tag_dataset.py

@@ -0,0 +1,350 @@
+import os
+import os.path as osp
+from collections.abc import Sequence
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from torch import Tensor
+from tqdm import tqdm
+
+from torch_geometric.data import InMemoryDataset, download_google_url
+from torch_geometric.data.data import BaseData
+
+try:
+    from pandas import DataFrame, read_csv
+    WITH_PANDAS = True
+except ImportError:
+    WITH_PANDAS = False
+
+IndexType = Union[slice, Tensor, np.ndarray, Sequence]
+
+
+class TAGDataset(InMemoryDataset):
+    r"""The Text Attributed Graph datasets from the
+    `"Learning on Large-scale Text-attributed Graphs via Variational Inference
+    " <https://arxiv.org/abs/2210.14709>`_ paper.
+    This dataset is aiming on transform `ogbn products`, `ogbn arxiv`
+    into Text Attributed Graph that each node in graph is associate with a
+    raw text, that dataset can be adapt to DataLoader (for LM training) and
+    NeighborLoader(for GNN training). In addition, this class can be use as a
+    wrapper class by convert a InMemoryDataset with Tokenizer and text into
+    Text Attributed Graph.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        dataset (InMemoryDataset): The name of the dataset
+            (:obj:`"ogbn-products"`, :obj:`"ogbn-arxiv"`).
+        tokenizer_name (str): The tokenizer name for language model,
+            Be sure to use same tokenizer name as your `model id` of model repo
+            on huggingface.co.
+        text (List[str]): list of raw text associate with node, the order of
+            list should be align with node list
+        split_idx (Optional[Dict[str, torch.Tensor]]): Optional dictionary,
+            for saving split index, it is required that if your dataset doesn't
+            have get_split_idx function
+        tokenize_batch_size (int): batch size of tokenizing text, the
+            tokenizing process will run on cpu, default: 256
+        token_on_disk (bool): save token as .pt file on disk or not,
+            default: False
+        text_on_disk (bool): save given text(list of str) as dataframe on disk
+            or not, default: False
+        force_reload (bool): default: False
+    .. note::
+        See `example/llm_plus_gnn/glem.py` for example usage
+    """
+    raw_text_id = {
+        'ogbn-arxiv': '1g3OOVhRyiyKv13LY6gbp8GLITocOUr_3',
+        'ogbn-products': '1I-S176-W4Bm1iPDjQv3hYwQBtxE0v8mt'
+    }
+
+    def __init__(self, root: str, dataset: InMemoryDataset,
+                 tokenizer_name: str, text: Optional[List[str]] = None,
+                 split_idx: Optional[Dict[str, Tensor]] = None,
+                 tokenize_batch_size: int = 256, token_on_disk: bool = False,
+                 text_on_disk: bool = False,
+                 force_reload: bool = False) -> None:
+        # list the vars you want to pass in before run download & process
+        self.name = dataset.name
+        self.text = text
+        self.tokenizer_name = tokenizer_name
+        from transformers import AutoTokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        self.dir_name = '_'.join(dataset.name.split('-'))
+        self.root = osp.join(root, self.dir_name)
+        missing_str_list = []
+        if not WITH_PANDAS:
+            missing_str_list.append('pandas')
+        if len(missing_str_list) > 0:
+            missing_str = ' '.join(missing_str_list)
+            error_out = f"`pip install {missing_str}` to use this dataset."
+            raise ImportError(error_out)
+        if hasattr(dataset, 'get_idx_split'):
+            self.split_idx = dataset.get_idx_split()
+        elif split_idx is not None:
+            self.split_idx = split_idx
+        else:
+            raise ValueError("TAGDataset need split idx for generating "
+                             "is_gold mask, please pass splited index "
+                             "in format of dictionaty with 'train', 'valid' "
+                             "'test' index tensor to 'split_idx'")
+        if text is not None and text_on_disk:
+            self.save_node_text(text)
+        self.text_on_disk = text_on_disk
+        # init will call download and process
+        super().__init__(self.root, transform=None, pre_transform=None,
+                         pre_filter=None, force_reload=force_reload)
+        # after processing and download
+        # Dataset has to have BaseData as _data
+        assert dataset._data is not None
+        self._data = dataset._data  # reassign reference
+        assert self._data is not None
+        assert dataset._data.y is not None
+        assert isinstance(self._data, BaseData)
+        assert self._data.num_nodes is not None
+        assert isinstance(dataset._data.num_nodes, int)
+        assert isinstance(self._data.num_nodes, int)
+        self._n_id = torch.arange(self._data.num_nodes)
+        is_good_tensor = self.load_gold_mask()
+        self._is_gold = is_good_tensor.squeeze()
+        self._data['is_gold'] = is_good_tensor
+        if self.text is not None and len(self.text) != self._data.num_nodes:
+            raise ValueError("The number of text sequence in 'text' should be "
+                             "equal to number of nodes!")
+        self.token_on_disk = token_on_disk
+        self.tokenize_batch_size = tokenize_batch_size
+        self._token = self.tokenize_graph(self.tokenize_batch_size)
+        self.__num_classes__ = dataset.num_classes
+
+    @property
+    def num_classes(self) -> int:
+        return self.__num_classes__
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        file_names = []
+        for root, _, files in os.walk(osp.join(self.root, 'raw')):
+            for file in files:
+                file_names.append(file)
+        return file_names
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        return [
+            'geometric_data_processed.pt', 'pre_filter.pt',
+            'pre_transformed.pt'
+        ]
+
+    @property
+    def token(self) -> Dict[str, Tensor]:
+        if self._token is None:  # lazy load
+            self._token = self.tokenize_graph()
+        return self._token
+
+    # load is_gold after init
+    @property
+    def is_gold(self) -> Tensor:
+        if self._is_gold is None:
+            print('lazy load is_gold!!')
+            self._is_gold = self.load_gold_mask()
+        return self._is_gold
+
+    def get_n_id(self, node_idx: IndexType) -> Tensor:
+        if self._n_id is None:
+            assert self._data is not None
+            assert self._data.num_nodes is not None
+            assert isinstance(self._data.num_nodes, int)
+            self._n_id = torch.arange(self._data.num_nodes)
+        return self._n_id[node_idx]
+
+    def load_gold_mask(self) -> Tensor:
+        r"""Use original train split as gold split, generating is_gold mask
+        for picking ground truth labels and pseudo labels.
+        """
+        train_split_idx = self.get_idx_split()['train']
+        assert self._data is not None
+        assert self._data.num_nodes is not None
+        assert isinstance(self._data.num_nodes, int)
+        is_good_tensor = torch.zeros(self._data.num_nodes,
+                                     dtype=torch.bool).view(-1, 1)
+        is_good_tensor[train_split_idx] = True
+        return is_good_tensor
+
+    def get_gold(self, node_idx: IndexType) -> Tensor:
+        r"""Get gold mask for given node_idx.
+
+        Args:
+            node_idx (torch.tensor): a tensor contain node idx
+        """
+        if self._is_gold is None:
+            self._is_gold = self.is_gold
+        return self._is_gold[node_idx]
+
+    def get_idx_split(self) -> Dict[str, Tensor]:
+        return self.split_idx
+
+    def download(self) -> None:
+        print('downloading raw text')
+        raw_text_path = download_google_url(id=self.raw_text_id[self.name],
+                                            folder=f'{self.root}/raw',
+                                            filename='node-text.csv.gz',
+                                            log=True)
+        text_df = read_csv(raw_text_path)
+        self.text = list(text_df['text'])
+
+    def process(self) -> None:
+        if osp.exists(osp.join(self.root, 'raw', 'node-text.csv.gz')):
+            text_df = read_csv(osp.join(self.root, 'raw', 'node-text.csv.gz'))
+            self.text = list(text_df['text'])
+        elif self.name in self.raw_text_id:
+            self.download()
+        else:
+            print('The dataset is not ogbn-products nor ogbn-arxiv,'
+                  'please pass in your raw text string list to `text`')
+        if self.text is None:
+            raise ValueError("The TAGDataset only have ogbn-products and "
+                             "ogbn-arxiv raw text in default "
+                             "The raw text of each node is not specified"
+                             "Please pass in 'text' when convert your dataset "
+                             "to Text Attribute Graph Dataset")
+
+    def save_node_text(self, text: List[str]) -> None:
+        node_text_path = osp.join(self.root, 'raw', 'node-text.csv.gz')
+        if osp.exists(node_text_path):
+            print(f'The raw text is existed at {node_text_path}')
+        else:
+            print(f'Saving raw text file at {node_text_path}')
+            os.makedirs(f'{self.root}/raw', exist_ok=True)
+            text_df = DataFrame(text, columns=['text'])
+            text_df.to_csv(osp.join(node_text_path), compression='gzip',
+                           index=False)
+
+    def tokenize_graph(self, batch_size: int = 256) -> Dict[str, Tensor]:
+        r"""Tokenizing the text associate with each node, running in cpu.
+
+        Args:
+            batch_size (Optional[int]): batch size of list of text for
+                generating emebdding
+        Returns:
+            Dict[str, torch.Tensor]: tokenized graph
+        """
+        data_len = 0
+        if self.text is not None:
+            data_len = len(self.text)
+        else:
+            raise ValueError("The TAGDataset need text for tokenization")
+        token_keys = ['input_ids', 'token_type_ids', 'attention_mask']
+        path = os.path.join(self.processed_dir, 'token', self.tokenizer_name)
+        # Check if the .pt files already exist
+        token_files_exist = any(
+            os.path.exists(os.path.join(path, f'{k}.pt')) for k in token_keys)
+
+        if token_files_exist and self.token_on_disk:
+            print('Found tokenized file, loading may take several minutes...')
+            all_encoded_token = {
+                k: torch.load(os.path.join(path, f'{k}.pt'), weights_only=True)
+                for k in token_keys
+                if os.path.exists(os.path.join(path, f'{k}.pt'))
+            }
+            return all_encoded_token
+
+        all_encoded_token = {k: [] for k in token_keys}
+        pbar = tqdm(total=data_len)
+
+        pbar.set_description('Tokenizing Text Attributed Graph')
+        for i in range(0, data_len, batch_size):
+            end_index = min(data_len, i + batch_size)
+            token = self.tokenizer(self.text[i:min(i + batch_size, data_len)],
+                                   padding='max_length', truncation=True,
+                                   max_length=512, return_tensors="pt")
+            for k in token.keys():
+                all_encoded_token[k].append(token[k])
+            pbar.update(end_index - i)
+        pbar.close()
+
+        all_encoded_token = {
+            k: torch.cat(v)
+            for k, v in all_encoded_token.items() if len(v) > 0
+        }
+        if self.token_on_disk:
+            os.makedirs(path, exist_ok=True)
+            print('Saving tokens on Disk')
+            for k, tensor in all_encoded_token.items():
+                torch.save(tensor, os.path.join(path, f'{k}.pt'))
+                print('Token saved:', os.path.join(path, f'{k}.pt'))
+        os.environ["TOKENIZERS_PARALLELISM"] = 'true'  # supressing warning
+        return all_encoded_token
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}()'
+
+    class TextDataset(torch.utils.data.Dataset):
+        r"""This nested dataset provides textual data for each node in
+        the graph. Factory method to create TextDataset from TAGDataset.
+
+        Args:
+            tag_dataset (TAGDataset): the parent dataset
+        """
+        def __init__(self, tag_dataset: 'TAGDataset') -> None:
+            self.tag_dataset = tag_dataset
+            self.token = tag_dataset.token
+            assert tag_dataset._data is not None
+            self._data = tag_dataset._data
+
+            assert tag_dataset._data.y is not None
+            self.labels = tag_dataset._data.y
+
+        def get_token(self, node_idx: IndexType) -> Dict[str, Tensor]:
+            r"""This function will be called in __getitem__().
+
+            Args:
+                node_idx (IndexType): selected node idx in each batch
+            Returns:
+                items (Dict[str, Tensor]): input for LM
+            """
+            items = {k: v[node_idx] for k, v in self.token.items()}
+            return items
+
+        # for LM training
+        def __getitem__(
+                self, node_id: IndexType
+        ) -> Dict[str, Union[Tensor, Dict[str, Tensor]]]:
+            r"""This function will override the function in
+            torch.utils.data.Dataset, and will be called when you
+            iterate batch in the dataloader, make sure all following
+            key value pairs are present in the return dict.
+
+            Args:
+                node_id (List[int]): list of node idx for selecting tokens,
+                    labels etc. when iterating data loader for LM
+            Returns:
+                items (dict): input k,v pairs for Language model training and
+                    inference
+            """
+            item: Dict[str, Union[Tensor, Dict[str, Tensor]]] = {}
+            item['input'] = self.get_token(node_id)
+            item['labels'] = self.labels[node_id]
+            item['is_gold'] = self.tag_dataset.get_gold(node_id)
+            item['n_id'] = self.tag_dataset.get_n_id(node_id)
+            return item
+
+        def __len__(self) -> int:
+            assert self._data.num_nodes is not None
+            return self._data.num_nodes
+
+        def get(self, idx: int) -> BaseData:
+            return self._data
+
+        def __repr__(self) -> str:
+            return f'{self.__class__.__name__}()'
+
+    def to_text_dataset(self) -> TextDataset:
+        r"""Factory Build text dataset from Text Attributed Graph Dataset
+        each data point is node's associated text token.
+        """
+        return TAGDataset.TextDataset(self)

torch_geometric/datasets/upfd.py

@@ -3,7 +3,6 @@ import os.path as osp
 from typing import Callable, List, Optional
 
 import numpy as np
-import scipy.sparse as sp
 import torch
 
 from torch_geometric.data import (
@@ -130,6 +129,8 @@ class UPFD(InMemoryDataset):
         os.remove(path)
 
     def process(self) -> None:
+        import scipy.sparse as sp
+
         x = sp.load_npz(
             osp.join(self.raw_dir, f'new_{self.feature}_feature.npz'))
         x = torch.from_numpy(x.todense()).to(torch.float)

torch_geometric/datasets/web_qsp_dataset.py

@@ -0,0 +1,246 @@
+# Code adapted from the G-Retriever paper: https://arxiv.org/abs/2402.07630
+from typing import Any, Dict, List, Tuple, no_type_check
+
+import numpy as np
+import torch
+from torch import Tensor
+from tqdm import tqdm
+
+from torch_geometric.data import Data, InMemoryDataset
+from torch_geometric.nn.nlp import SentenceTransformer
+
+
+@no_type_check
+def retrieval_via_pcst(
+    data: Data,
+    q_emb: Tensor,
+    textual_nodes: Any,
+    textual_edges: Any,
+    topk: int = 3,
+    topk_e: int = 3,
+    cost_e: float = 0.5,
+) -> Tuple[Data, str]:
+    c = 0.01
+
+    from pcst_fast import pcst_fast
+
+    root = -1
+    num_clusters = 1
+    pruning = 'gw'
+    verbosity_level = 0
+    if topk > 0:
+        n_prizes = torch.nn.CosineSimilarity(dim=-1)(q_emb, data.x)
+        topk = min(topk, data.num_nodes)
+        _, topk_n_indices = torch.topk(n_prizes, topk, largest=True)
+
+        n_prizes = torch.zeros_like(n_prizes)
+        n_prizes[topk_n_indices] = torch.arange(topk, 0, -1).float()
+    else:
+        n_prizes = torch.zeros(data.num_nodes)
+
+    if topk_e > 0:
+        e_prizes = torch.nn.CosineSimilarity(dim=-1)(q_emb, data.edge_attr)
+        topk_e = min(topk_e, e_prizes.unique().size(0))
+
+        topk_e_values, _ = torch.topk(e_prizes.unique(), topk_e, largest=True)
+        e_prizes[e_prizes < topk_e_values[-1]] = 0.0
+        last_topk_e_value = topk_e
+        for k in range(topk_e):
+            indices = e_prizes == topk_e_values[k]
+            value = min((topk_e - k) / sum(indices), last_topk_e_value - c)
+            e_prizes[indices] = value
+            last_topk_e_value = value * (1 - c)
+        # reduce the cost of the edges such that at least one edge is selected
+        cost_e = min(cost_e, e_prizes.max().item() * (1 - c / 2))
+    else:
+        e_prizes = torch.zeros(data.num_edges)
+
+    costs = []
+    edges = []
+    virtual_n_prizes = []
+    virtual_edges = []
+    virtual_costs = []
+    mapping_n = {}
+    mapping_e = {}
+    for i, (src, dst) in enumerate(data.edge_index.t().numpy()):
+        prize_e = e_prizes[i]
+        if prize_e <= cost_e:
+            mapping_e[len(edges)] = i
+            edges.append((src, dst))
+            costs.append(cost_e - prize_e)
+        else:
+            virtual_node_id = data.num_nodes + len(virtual_n_prizes)
+            mapping_n[virtual_node_id] = i
+            virtual_edges.append((src, virtual_node_id))
+            virtual_edges.append((virtual_node_id, dst))
+            virtual_costs.append(0)
+            virtual_costs.append(0)
+            virtual_n_prizes.append(prize_e - cost_e)
+
+    prizes = np.concatenate([n_prizes, np.array(virtual_n_prizes)])
+    num_edges = len(edges)
+    if len(virtual_costs) > 0:
+        costs = np.array(costs + virtual_costs)
+        edges = np.array(edges + virtual_edges)
+
+    vertices, edges = pcst_fast(edges, prizes, costs, root, num_clusters,
+                                pruning, verbosity_level)
+
+    selected_nodes = vertices[vertices < data.num_nodes]
+    selected_edges = [mapping_e[e] for e in edges if e < num_edges]
+    virtual_vertices = vertices[vertices >= data.num_nodes]
+    if len(virtual_vertices) > 0:
+        virtual_vertices = vertices[vertices >= data.num_nodes]
+        virtual_edges = [mapping_n[i] for i in virtual_vertices]
+        selected_edges = np.array(selected_edges + virtual_edges)
+
+    edge_index = data.edge_index[:, selected_edges]
+    selected_nodes = np.unique(
+        np.concatenate(
+            [selected_nodes, edge_index[0].numpy(), edge_index[1].numpy()]))
+
+    n = textual_nodes.iloc[selected_nodes]
+    e = textual_edges.iloc[selected_edges]
+    desc = n.to_csv(index=False) + '\n' + e.to_csv(
+        index=False, columns=['src', 'edge_attr', 'dst'])
+
+    mapping = {n: i for i, n in enumerate(selected_nodes.tolist())}
+    src = [mapping[i] for i in edge_index[0].tolist()]
+    dst = [mapping[i] for i in edge_index[1].tolist()]
+
+    data = Data(
+        x=data.x[selected_nodes],
+        edge_index=torch.tensor([src, dst]),
+        edge_attr=data.edge_attr[selected_edges],
+    )
+
+    return data, desc
+
+
+class WebQSPDataset(InMemoryDataset):
+    r"""The WebQuestionsSP dataset of the `"The Value of Semantic Parse
+    Labeling for Knowledge Base Question Answering"
+    <https://aclanthology.org/P16-2033/>`_ paper.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        split (str, optional): If :obj:`"train"`, loads the training dataset.
+            If :obj:`"val"`, loads the validation dataset.
+            If :obj:`"test"`, loads the test dataset. (default: :obj:`"train"`)
+        force_reload (bool, optional): Whether to re-process the dataset.
+            (default: :obj:`False`)
+        use_pcst (bool, optional): Whether to preprocess the dataset's graph
+            with PCST or return the full graphs. (default: :obj:`True`)
+    """
+    def __init__(
+        self,
+        root: str,
+        split: str = "train",
+        force_reload: bool = False,
+        use_pcst: bool = True,
+    ) -> None:
+        self.use_pcst = use_pcst
+        super().__init__(root, force_reload=force_reload)
+
+        if split not in {'train', 'val', 'test'}:
+            raise ValueError(f"Invalid 'split' argument (got {split})")
+
+        path = self.processed_paths[['train', 'val', 'test'].index(split)]
+        self.load(path)
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        return ['train_data.pt', 'val_data.pt', 'test_data.pt']
+
+    def process(self) -> None:
+        import datasets
+        import pandas as pd
+
+        datasets = datasets.load_dataset('rmanluo/RoG-webqsp')
+
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        model_name = 'sentence-transformers/all-roberta-large-v1'
+        model = SentenceTransformer(model_name).to(device)
+        model.eval()
+
+        for dataset, path in zip(
+            [datasets['train'], datasets['validation'], datasets['test']],
+                self.processed_paths,
+        ):
+            questions = [example["question"] for example in dataset]
+            question_embs = model.encode(
+                questions,
+                batch_size=256,
+                output_device='cpu',
+            )
+
+            data_list = []
+            for i, example in enumerate(tqdm(dataset)):
+                raw_nodes: Dict[str, int] = {}
+                raw_edges = []
+                for tri in example["graph"]:
+                    h, r, t = tri
+                    h = h.lower()
+                    t = t.lower()
+                    if h not in raw_nodes:
+                        raw_nodes[h] = len(raw_nodes)
+                    if t not in raw_nodes:
+                        raw_nodes[t] = len(raw_nodes)
+                    raw_edges.append({
+                        "src": raw_nodes[h],
+                        "edge_attr": r,
+                        "dst": raw_nodes[t]
+                    })
+                nodes = pd.DataFrame([{
+                    "node_id": v,
+                    "node_attr": k,
+                } for k, v in raw_nodes.items()],
+                                     columns=["node_id", "node_attr"])
+                edges = pd.DataFrame(raw_edges,
+                                     columns=["src", "edge_attr", "dst"])
+
+                nodes.node_attr = nodes.node_attr.fillna("")
+                x = model.encode(
+                    nodes.node_attr.tolist(),
+                    batch_size=256,
+                    output_device='cpu',
+                )
+                edge_attr = model.encode(
+                    edges.edge_attr.tolist(),
+                    batch_size=256,
+                    output_device='cpu',
+                )
+                edge_index = torch.tensor([
+                    edges.src.tolist(),
+                    edges.dst.tolist(),
+                ], dtype=torch.long)
+
+                question = f"Question: {example['question']}\nAnswer: "
+                label = ('|').join(example['answer']).lower()
+                data = Data(
+                    x=x,
+                    edge_index=edge_index,
+                    edge_attr=edge_attr,
+                )
+                if self.use_pcst and len(nodes) > 0 and len(edges) > 0:
+                    data, desc = retrieval_via_pcst(
+                        data,
+                        question_embs[i],
+                        nodes,
+                        edges,
+                        topk=3,
+                        topk_e=5,
+                        cost_e=0.5,
+                    )
+                else:
+                    desc = nodes.to_csv(index=False) + "\n" + edges.to_csv(
+                        index=False,
+                        columns=["src", "edge_attr", "dst"],
+                    )
+
+                data.question = question
+                data.label = label
+                data.desc = desc
+                data_list.append(data)
+
+            self.save(data_list, path)

torch_geometric/datasets/webkb.py

@@ -102,7 +102,7 @@ class WebKB(InMemoryDataset):
             download_url(f'{self.url}/splits/{f}', self.raw_dir)
 
     def process(self) -> None:
-        with open(self.raw_paths[0], 'r') as f:
+        with open(self.raw_paths[0]) as f:
             lines = f.read().split('\n')[1:-1]
             xs = [[float(value) for value in line.split('\t')[1].split(',')]
                   for line in lines]
@@ -111,7 +111,7 @@ class WebKB(InMemoryDataset):
             ys = [int(line.split('\t')[2]) for line in lines]
             y = torch.tensor(ys, dtype=torch.long)
 
-        with open(self.raw_paths[1], 'r') as f:
+        with open(self.raw_paths[1]) as f:
             lines = f.read().split('\n')[1:-1]
             edge_indices = [[int(value) for value in line.split('\t')]
                             for line in lines]

torch_geometric/datasets/wikics.py

@@ -65,7 +65,7 @@ class WikiCS(InMemoryDataset):
             download_url(f'{self.url}/{name}', self.raw_dir)
 
     def process(self) -> None:
-        with open(self.raw_paths[0], 'r') as f:
+        with open(self.raw_paths[0]) as f:
             data = json.load(f)
 
         x = torch.tensor(data['features'], dtype=torch.float)

torch_geometric/datasets/wikidata.py

@@ -10,6 +10,7 @@ from torch_geometric.data import (
     download_url,
     extract_tar,
 )
+from torch_geometric.io import fs
 
 
 class Wikidata5M(InMemoryDataset):
@@ -99,7 +100,7 @@ class Wikidata5M(InMemoryDataset):
                 values = line.strip().split('\t')
                 entity_to_id[values[0]] = i
 
-        x = torch.load(self.raw_paths[1])
+        x = fs.torch_load(self.raw_paths[1])
 
         edge_indices = []
         edge_types = []
@@ -107,7 +108,7 @@ class Wikidata5M(InMemoryDataset):
 
         rel_to_id: Dict[str, int] = {}
         for split, path in enumerate(self.raw_paths[2:]):
-            with open(path, 'r') as f:
+            with open(path) as f:
                 for line in f:
                     head, rel, tail = line[:-1].split('\t')
                     src = entity_to_id[head]

torch_geometric/datasets/wikipedia_network.py

@@ -105,7 +105,7 @@ class WikipediaNetwork(InMemoryDataset):
 
     def process(self) -> None:
         if self.geom_gcn_preprocess:
-            with open(self.raw_paths[0], 'r') as f:
+            with open(self.raw_paths[0]) as f:
                 lines = f.read().split('\n')[1:-1]
             xs = [[float(value) for value in line.split('\t')[1].split(',')]
                   for line in lines]
@@ -113,7 +113,7 @@ class WikipediaNetwork(InMemoryDataset):
             ys = [int(line.split('\t')[2]) for line in lines]
             y = torch.tensor(ys, dtype=torch.long)
 
-            with open(self.raw_paths[1], 'r') as f:
+            with open(self.raw_paths[1]) as f:
                 lines = f.read().split('\n')[1:-1]
                 edge_indices = [[int(value) for value in line.split('\t')]
                                 for line in lines]