pyg-nightly 2.6.0.dev20240704__py3-none-any.whl → 2.8.0.dev20251207__py3-none-any.whl

torch_geometric/datasets/teeth3ds.py

@@ -0,0 +1,269 @@
+import json
+import os
+import os.path as osp
+from glob import glob
+from typing import Callable, Dict, List, Optional
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from torch_geometric.data import (
+    Data,
+    InMemoryDataset,
+    download_url,
+    extract_zip,
+)
+
+
+class Teeth3DS(InMemoryDataset):
+    r"""The Teeth3DS+ dataset from the `"An Extended Benchmark for Intra-oral
+    3D Scans Analysis" <https://crns-smartvision.github.io/teeth3ds/>`_ paper.
+
+    This dataset is the first comprehensive public benchmark designed to
+    advance the field of intra-oral 3D scan analysis developed as part of the
+    3DTeethSeg 2022 and 3DTeethLand 2024 MICCAI challenges, aiming to drive
+    research in teeth identification, segmentation, labeling, 3D modeling,
+    and dental landmark identification.
+    The dataset includes at least 1,800 intra-oral scans (containing 23,999
+    annotated teeth) collected from 900 patients, covering both upper and lower
+    jaws separately.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        split (str): The split name (one of :obj:`"Teeth3DS"`,
+            :obj:`"3DTeethSeg22_challenge"` or :obj:`"3DTeethLand_challenge"`).
+        train (bool, optional): If :obj:`True`, loads the training dataset,
+            otherwise the test dataset. (default: :obj:`True`)
+        num_samples (int, optional): Number of points to sample from each mesh.
+            (default: :obj:`30000`)
+        transform (callable, optional): A function/transform that takes in an
+            :obj:`torch_geometric.data.Data` object and returns a transformed
+            version. The data object will be transformed before every access.
+            (default: :obj:`None`)
+        pre_transform (callable, optional): A function/transform that takes in
+            an :obj:`torch_geometric.data.Data` object and returns a
+            transformed version. The data object will be transformed before
+            being saved to disk. (default: :obj:`None`)
+        force_reload (bool, optional): Whether to re-process the dataset.
+            (default: :obj:`False`)
+    """
+    urls = {
+        'data_part_1.zip':
+        'https://osf.io/download/qhprs/',
+        'data_part_2.zip':
+        'https://osf.io/download/4pwnr/',
+        'data_part_3.zip':
+        'https://osf.io/download/frwdp/',
+        'data_part_4.zip':
+        'https://osf.io/download/2arn4/',
+        'data_part_5.zip':
+        'https://osf.io/download/xrz5f/',
+        'data_part_6.zip':
+        'https://osf.io/download/23hgq/',
+        'data_part_7.zip':
+        'https://osf.io/download/u83ad/',
+        'train_test_split':
+        'https://files.de-1.osf.io/v1/'
+        'resources/xctdy/providers/osfstorage/?zip='
+    }
+
+    sample_url = {
+        'teeth3ds_sample': 'https://osf.io/download/vr38s/',
+    }
+
+    landmarks_urls = {
+        '3DTeethLand_landmarks_train.zip': 'https://osf.io/download/k5hbj/',
+        '3DTeethLand_landmarks_test.zip': 'https://osf.io/download/sqw5e/',
+    }
+
+    def __init__(
+        self,
+        root: str,
+        split:
+        str = 'Teeth3DS',  # [3DTeethSeg22_challenge, 3DTeethLand_challenge]
+        train: bool = True,
+        num_samples: int = 30000,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        force_reload: bool = False,
+    ) -> None:
+
+        self.mode = 'training' if train else 'testing'
+        self.split = split
+        self.num_samples = num_samples
+
+        super().__init__(root, transform, pre_transform,
+                         force_reload=force_reload)
+
+    @property
+    def processed_dir(self) -> str:
+        return os.path.join(self.root, f'processed_{self.split}_{self.mode}')
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return ['license.txt']
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        # Directory containing train/test split files:
+        split_subdir = 'teeth3ds_sample' if self.split == 'sample' else ''
+        split_dir = osp.join(
+            self.raw_dir,
+            split_subdir,
+            f'{self.split}_train_test_split',
+        )
+
+        split_files = glob(osp.join(split_dir, f'{self.mode}*.txt'))
+
+        # Collect all file names from the split files:
+        combined_list = []
+        for file_path in split_files:
+            with open(file_path) as file:
+                combined_list.extend(file.read().splitlines())
+
+        # Generate the list of processed file paths:
+        return [f'{file_name}.pt' for file_name in combined_list]
+
+    def download(self) -> None:
+        if self.split == 'sample':
+            for key, url in self.sample_url.items():
+                path = download_url(url, self.root, filename=key)
+                extract_zip(path, self.raw_dir)
+                os.unlink(path)
+        else:
+            for key, url in self.urls.items():
+                path = download_url(url, self.root, filename=key)
+                extract_zip(path, self.raw_dir)
+                os.unlink(path)
+            for key, url in self.landmarks_urls.items():
+                path = download_url(url, self.root, filename=key)
+                extract_zip(path, self.raw_dir)  # Extract each downloaded part
+                os.unlink(path)
+
+    def process_file(self, file_path: str) -> Optional[Data]:
+        """Processes the input file path to load mesh data, annotations,
+        and prepare the input features for a graph-based deep learning model.
+        """
+        import trimesh
+        from fpsample import bucket_fps_kdline_sampling
+
+        mesh = trimesh.load_mesh(file_path)
+
+        if isinstance(mesh, list):
+            # Handle the case where a list of Geometry objects is returned
+            mesh = mesh[0]
+
+        vertices = mesh.vertices
+        vertex_normals = mesh.vertex_normals
+
+        # Perform sampling on mesh vertices:
+        if len(vertices) < self.num_samples:
+            sampled_indices = np.random.choice(
+                len(vertices),
+                self.num_samples,
+                replace=True,
+            )
+        else:
+            sampled_indices = bucket_fps_kdline_sampling(
+                vertices,
+                self.num_samples,
+                h=5,
+                start_idx=0,
+            )
+
+        if len(sampled_indices) != self.num_samples:
+            raise RuntimeError(f"Sampled points mismatch, expected "
+                               f"{self.num_samples} points, but got "
+                               f"{len(sampled_indices)} for '{file_path}'")
+
+        # Extract features and annotations for the sampled points:
+        pos = torch.tensor(vertices[sampled_indices], dtype=torch.float)
+        x = torch.tensor(vertex_normals[sampled_indices], dtype=torch.float)
+
+        # Load segmentation annotations:
+        seg_annotation_path = file_path.replace('.obj', '.json')
+        if osp.exists(seg_annotation_path):
+            with open(seg_annotation_path) as f:
+                seg_annotations = json.load(f)
+            y = torch.tensor(
+                np.asarray(seg_annotations['labels'])[sampled_indices],
+                dtype=torch.float)
+            instances = torch.tensor(
+                np.asarray(seg_annotations['instances'])[sampled_indices],
+                dtype=torch.float)
+        else:
+            y = torch.empty(0, 3)
+            instances = torch.empty(0, 3)
+
+        # Load landmarks annotations:
+        landmarks_annotation_path = file_path.replace('.obj', '__kpt.json')
+
+        # Parse keypoint annotations into structured tensors:
+        keypoints_dict: Dict[str, List] = {
+            key: []
+            for key in [
+                'Mesial', 'Distal', 'Cusp', 'InnerPoint', 'OuterPoint',
+                'FacialPoint'
+            ]
+        }
+        keypoint_tensors: Dict[str, torch.Tensor] = {
+            key: torch.empty(0, 3)
+            for key in [
+                'Mesial', 'Distal', 'Cusp', 'InnerPoint', 'OuterPoint',
+                'FacialPoint'
+            ]
+        }
+        if osp.exists(landmarks_annotation_path):
+            with open(landmarks_annotation_path) as f:
+                landmarks_annotations = json.load(f)
+
+            for keypoint in landmarks_annotations['objects']:
+                keypoints_dict[keypoint['class']].extend(keypoint['coord'])
+
+            keypoint_tensors = {
+                k: torch.tensor(np.asarray(v),
+                                dtype=torch.float).reshape(-1, 3)
+                for k, v in keypoints_dict.items()
+            }
+
+        data = Data(
+            pos=pos,
+            x=x,
+            y=y,
+            instances=instances,
+            jaw=file_path.split('.obj')[0].split('_')[1],
+            mesial=keypoint_tensors['Mesial'],
+            distal=keypoint_tensors['Distal'],
+            cusp=keypoint_tensors['Cusp'],
+            inner_point=keypoint_tensors['InnerPoint'],
+            outer_point=keypoint_tensors['OuterPoint'],
+            facial_point=keypoint_tensors['FacialPoint'],
+        )
+
+        if self.pre_transform is not None:
+            data = self.pre_transform(data)
+
+        return data
+
+    def process(self) -> None:
+        for file in tqdm(self.processed_file_names):
+            name = file.split('.')[0]
+            path = osp.join(self.raw_dir, '**', '*', name + '.obj')
+            paths = glob(path)
+            if len(paths) == 1:
+                data = self.process_file(paths[0])
+                torch.save(data, osp.join(self.processed_dir, file))
+
+    def len(self) -> int:
+        return len(self.processed_file_names)
+
+    def get(self, idx: int) -> Data:
+        return torch.load(
+            osp.join(self.processed_dir, self.processed_file_names[idx]),
+            weights_only=False,
+        )
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({len(self)}, '
+                f'mode={self.mode}, split={self.split})')

torch_geometric/datasets/web_qsp_dataset.py

@@ -0,0 +1,342 @@
+# Code adapted from the G-Retriever paper: https://arxiv.org/abs/2402.07630
+import gc
+import os
+from itertools import chain
+from typing import Any, Dict, Iterator, List, Optional
+
+import torch
+from tqdm import tqdm
+
+from torch_geometric.data import InMemoryDataset
+from torch_geometric.llm.large_graph_indexer import (
+    EDGE_RELATION,
+    LargeGraphIndexer,
+    TripletLike,
+    get_features_for_triplets_groups,
+)
+from torch_geometric.llm.models import SentenceTransformer
+from torch_geometric.llm.utils.backend_utils import (
+    preprocess_triplet,
+    retrieval_via_pcst,
+)
+
+
+class KGQABaseDataset(InMemoryDataset):
+    r"""Base class for the 2 KGQA datasets used in `"Reasoning on Graphs:
+    Faithful and Interpretable Large Language Model Reasoning"
+    <https://arxiv.org/pdf/2310.01061>`_ paper.
+
+    Args:
+        dataset_name (str): HuggingFace `dataset` name.
+        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`)
+        verbose (bool, optional): Whether to print output. Defaults to False.
+        use_pcst (bool, optional): Whether to preprocess the dataset's graph
+            with PCST or return the full graphs. (default: :obj:`True`)
+        load_dataset_kwargs (dict, optional):
+            Keyword arguments for the `datasets.load_dataset` function.
+            (default: :obj:`{}`)
+        retrieval_kwargs (dict, optional):
+            Keyword arguments for the
+            `get_features_for_triplets_groups` function.
+            (default: :obj:`{}`)
+    """
+    def __init__(
+        self,
+        dataset_name: str,
+        root: str,
+        split: str = "train",
+        force_reload: bool = False,
+        verbose: bool = False,
+        use_pcst: bool = True,
+        load_dataset_kwargs: Optional[Dict[str, Any]] = None,
+        retrieval_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        self.split = split
+        self.dataset_name = dataset_name
+        self.use_pcst = use_pcst
+        self.load_dataset_kwargs = load_dataset_kwargs or {}
+        """
+        NOTE: If running into memory issues,
+        try reducing this batch size for the LargeGraphIndexer
+        used to build our KG.
+        Example: self.retrieval_kwargs = {"batch_size": 64}
+        """
+        self.retrieval_kwargs = retrieval_kwargs or {}
+
+        # Caching custom subsets of the dataset results in unsupported behavior
+        if 'split' in self.load_dataset_kwargs:
+            print("WARNING: Caching custom subsets of the dataset \
+                results in unsupported behavior.\
+                Please specify a separate root directory for each split,\
+                or set force_reload=True on subsequent instantiations\
+                of the dataset.")
+
+        self.required_splits = ['train', 'validation', 'test']
+
+        self.verbose = verbose
+        self.force_reload = force_reload
+        super().__init__(root, force_reload=force_reload)
+        """
+        NOTE: Current behavior is to process the entire dataset,
+        and only return the split specified by the user.
+        """
+        if f'{split}_data.pt' not in set(self.processed_file_names):
+            raise ValueError(f"Invalid 'split' argument (got {split})")
+        if split == 'val':
+            split = 'validation'
+
+        self.load(self.processed_paths[self.required_splits.index(split)])
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return ["raw.pt"]
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        return ["train_data.pt", "val_data.pt", "test_data.pt"]
+
+    def download(self) -> None:
+        import datasets
+
+        # HF Load Dataset by dataset name if no path is specified
+        self.load_dataset_kwargs['path'] = self.load_dataset_kwargs.get(
+            'path', self.dataset_name)
+        raw_dataset = datasets.load_dataset(**self.load_dataset_kwargs)
+
+        # Assert that the dataset contains the required splits
+        assert all(split in raw_dataset for split in self.required_splits), \
+            f"Dataset '{self.dataset_name}' is missing required splits: \
+            {self.required_splits}"
+
+        raw_dataset.save_to_disk(self.raw_paths[0])
+
+    def _get_trips(self) -> Iterator[TripletLike]:
+        # Iterate over each element's graph in each split of the dataset
+        # Using chain to lazily iterate without storing all trips in memory
+        split_iterators = []
+
+        for split in self.required_splits:
+            # Create an iterator for each element's graph in the current split
+            split_graphs = (element['graph']
+                            for element in self.raw_dataset[split])
+            split_iterators.append(chain.from_iterable(split_graphs))
+
+        # Chain all split iterators together
+        return chain.from_iterable(split_iterators)
+
+    def _build_graph(self) -> None:
+        print("Encoding graph...")
+        trips = self._get_trips()
+        self.indexer: LargeGraphIndexer = LargeGraphIndexer.from_triplets(
+            trips, pre_transform=preprocess_triplet)
+
+        # Nodes:
+        print("\tEncoding nodes...")
+        nodes = self.indexer.get_unique_node_features()
+        x = self.model.encode(nodes, batch_size=256, output_device='cpu')
+        self.indexer.add_node_feature(new_feature_name="x", new_feature_vals=x)
+
+        # Edges:
+        print("\tEncoding edges...")
+        edges = self.indexer.get_unique_edge_features(
+            feature_name=EDGE_RELATION)
+        edge_attr = self.model.encode(edges, batch_size=256,
+                                      output_device='cpu')
+        self.indexer.add_edge_feature(
+            new_feature_name="edge_attr",
+            new_feature_vals=edge_attr,
+            map_from_feature=EDGE_RELATION,
+        )
+
+        print("\tSaving graph...")
+        self.indexer.save(self.indexer_path)
+
+    def _retrieve_subgraphs(self) -> None:
+        raw_splits = [
+            self.raw_dataset[split] for split in self.required_splits
+        ]
+        zipped = zip(
+            self.required_splits,
+            raw_splits,  # noqa
+            self.processed_paths,
+        )
+        for split_name, dataset, path in zipped:
+            print(f"Processing {split_name} split...")
+
+            print("\tEncoding questions...")
+            split_questions = [str(element['question']) for element in dataset]
+            split_q_embs = self.model.encode(split_questions, batch_size=256,
+                                             output_device='cpu')
+
+            print("\tRetrieving subgraphs...")
+            results_graphs = []
+            retrieval_kwargs = {
+                **self.retrieval_kwargs,
+                **{
+                    'pre_transform': preprocess_triplet,
+                    'verbose': self.verbose,
+                }
+            }
+            graph_gen = get_features_for_triplets_groups(
+                self.indexer, (element['graph'] for element in dataset),
+                **retrieval_kwargs)
+
+            for index in tqdm(range(len(dataset)), disable=not self.verbose):
+                data_i = dataset[index]
+                graph = next(graph_gen)
+                textual_nodes = self.textual_nodes.iloc[
+                    graph["node_idx"]].reset_index()
+                textual_edges = self.textual_edges.iloc[
+                    graph["edge_idx"]].reset_index()
+                if self.use_pcst and len(textual_nodes) > 0 and len(
+                        textual_edges) > 0:
+                    subgraph, desc = retrieval_via_pcst(
+                        graph,
+                        split_q_embs[index],
+                        textual_nodes,
+                        textual_edges,
+                    )
+                else:
+                    desc = textual_nodes.to_csv(
+                        index=False) + "\n" + textual_edges.to_csv(
+                            index=False,
+                            columns=["src", "edge_attr", "dst"],
+                        )
+                    subgraph = graph
+                question = f"Question: {data_i['question']}\nAnswer: "
+                label = ("|").join(data_i["answer"]).lower()
+
+                subgraph["question"] = question
+                subgraph["label"] = label
+                subgraph["desc"] = desc
+                results_graphs.append(subgraph.to("cpu"))
+            print("\tSaving subgraphs...")
+            self.save(results_graphs, path)
+
+    def process(self) -> None:
+        import datasets
+        from pandas import DataFrame
+        self.raw_dataset = datasets.load_from_disk(self.raw_paths[0])
+
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        model_name = 'sentence-transformers/all-roberta-large-v1'
+        self.model: SentenceTransformer = SentenceTransformer(model_name).to(
+            device)
+        self.model.eval()
+        self.indexer_path = os.path.join(self.processed_dir,
+                                         "large_graph_indexer")
+        if self.force_reload or not os.path.exists(self.indexer_path):
+            self._build_graph()
+        else:
+            print("Loading graph...")
+            self.indexer = LargeGraphIndexer.from_disk(self.indexer_path)
+        self.textual_nodes = DataFrame.from_dict(
+            {"node_attr": self.indexer.get_node_features()})
+        self.textual_nodes["node_id"] = self.textual_nodes.index
+        self.textual_nodes = self.textual_nodes[["node_id", "node_attr"]]
+        self.textual_edges = DataFrame(self.indexer.get_edge_features(),
+                                       columns=["src", "edge_attr", "dst"])
+        self.textual_edges["src"] = [
+            self.indexer._nodes[h] for h in self.textual_edges["src"]
+        ]
+        self.textual_edges["dst"] = [
+            self.indexer._nodes[h] for h in self.textual_edges["dst"]
+        ]
+        self._retrieve_subgraphs()
+
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+class WebQSPDataset(KGQABaseDataset):
+    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`)
+        verbose (bool, optional): Whether to print output. Defaults to False.
+        use_pcst (bool, optional): Whether to preprocess the dataset's graph
+            with PCST or return the full graphs. (default: :obj:`True`)
+        load_dataset_kwargs (dict, optional):
+            Keyword arguments for the `datasets.load_dataset` function.
+            (default: :obj:`{}`)
+        retrieval_kwargs (dict, optional):
+            Keyword arguments for the
+            `get_features_for_triplets_groups` function.
+            (default: :obj:`{}`)
+    """
+    def __init__(
+        self,
+        root: str,
+        split: str = "train",
+        force_reload: bool = False,
+        verbose: bool = False,
+        use_pcst: bool = True,
+        load_dataset_kwargs: Optional[Dict[str, Any]] = None,
+        retrieval_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        load_dataset_kwargs = load_dataset_kwargs or {}
+        retrieval_kwargs = retrieval_kwargs or {}
+        # Modify these paramters if running into memory/compute issues
+        default_retrieval_kwargs = {
+            'max_batch_size': 250,  # Lower batch size to reduce memory usage
+            'num_workers':
+            None,  # Use all available workers, or set to number of threads
+        }
+        retrieval_kwargs = {**default_retrieval_kwargs, **retrieval_kwargs}
+        dataset_name = 'rmanluo/RoG-webqsp'
+        super().__init__(dataset_name, root, split, force_reload, verbose,
+                         use_pcst, load_dataset_kwargs=load_dataset_kwargs,
+                         retrieval_kwargs=retrieval_kwargs)
+
+
+class CWQDataset(KGQABaseDataset):
+    r"""The ComplexWebQuestions (CWQ) dataset of the `"The Web as a
+    Knowledge-base forAnswering Complex Questions"
+    <https://arxiv.org/pdf/1803.06643>`_ 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`)
+        verbose (bool, optional): Whether to print output. Defaults to False.
+        use_pcst (bool, optional): Whether to preprocess the dataset's graph
+            with PCST or return the full graphs. (default: :obj:`True`)
+        load_dataset_kwargs (dict, optional):
+            Keyword arguments for the `datasets.load_dataset` function.
+            (default: :obj:`{}`)
+        retrieval_kwargs (dict, optional):
+            Keyword arguments for the
+            `get_features_for_triplets_groups` function.
+            (default: :obj:`{}`)
+    """
+    def __init__(
+        self,
+        root: str,
+        split: str = "train",
+        force_reload: bool = False,
+        verbose: bool = False,
+        use_pcst: bool = True,
+        load_dataset_kwargs: Optional[Dict[str, Any]] = None,
+        retrieval_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        load_dataset_kwargs = load_dataset_kwargs or {}
+        retrieval_kwargs = retrieval_kwargs or {}
+        dataset_name = 'rmanluo/RoG-cwq'
+        super().__init__(dataset_name, root, split, force_reload, verbose,
+                         use_pcst, load_dataset_kwargs=load_dataset_kwargs,
+                         retrieval_kwargs=retrieval_kwargs)

torch_geometric/datasets/wikics.py

@@ -45,7 +45,8 @@ class WikiCS(InMemoryDataset):
             warnings.warn(
                 f"The {self.__class__.__name__} dataset now returns an "
                 f"undirected graph by default. Please explicitly specify "
-                f"'is_undirected=False' to restore the old behavior.")
+                f"'is_undirected=False' to restore the old behavior.",
+                stacklevel=2)
             is_undirected = True
         self.is_undirected = is_undirected
         super().__init__(root, transform, pre_transform,

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 = []

torch_geometric/deprecation.py

@@ -23,7 +23,7 @@ def deprecated(
             out = f"'{name}' is deprecated"
             if details is not None:
                 out += f", {details}"
-            warnings.warn(out)
+            warnings.warn(out, stacklevel=2)
             return func(*args, **kwargs)
 
         return wrapper

torch_geometric/distributed/__init__.py

@@ -1,3 +1,5 @@
+from warnings import warn
+
 from .dist_context import DistContext
 from .local_feature_store import LocalFeatureStore
 from .local_graph_store import LocalGraphStore
@@ -7,6 +9,17 @@ from .dist_loader import DistLoader
 from .dist_neighbor_loader import DistNeighborLoader
 from .dist_link_neighbor_loader import DistLinkNeighborLoader
 
+warn(
+    "`torch_geometric.distributed` has been deprecated since 2.7.0 and will "
+    "no longer be maintained. For distributed training, refer to our "
+    "tutorials on distributed training at "
+    "https://pytorch-geometric.readthedocs.io/en/latest/tutorial/distributed.html "  # noqa: E501
+    "or cuGraph examples at "
+    "https://github.com/rapidsai/cugraph-gnn/tree/main/python/cugraph-pyg/cugraph_pyg/examples",  # noqa: E501
+    stacklevel=2,
+    category=DeprecationWarning,
+)
+
 __all__ = classes = [
     'DistContext',
     'LocalFeatureStore',

torch_geometric/distributed/dist_loader.py

@@ -138,9 +138,9 @@ class DistLoader:
             # close RPC & worker group at exit:
             atexit.register(shutdown_rpc, self.current_ctx_worker.worker_name)
 
-        except RuntimeError:
+        except RuntimeError as e:
             raise RuntimeError(f"`{self}.init_fn()` could not initialize the "
-                               f"worker loop of the neighbor sampler")
+                               f"worker loop of the neighbor sampler") from e
 
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}(pid={self.pid})'