pyg-nightly 2.7.0.dev20241009__py3-none-any.whl → 2.8.0.dev20251207__py3-none-any.whl

torch_geometric/llm/utils/graph_store.py

@@ -0,0 +1,199 @@
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+from torch import Tensor
+
+from torch_geometric.data import FeatureStore
+from torch_geometric.distributed import LocalGraphStore
+from torch_geometric.sampler import (
+    BidirectionalNeighborSampler,
+    NodeSamplerInput,
+    SamplerOutput,
+)
+from torch_geometric.utils import index_sort
+
+# A representation of an edge index, following the possible formats:
+#    * default: Tensor, size = [2, num_edges]
+#    *     Tensor[0, :] == row, Tensor[1, :] == col
+#    * COO: (row, col)
+#    * CSC: (row, colptr)
+#    * CSR: (rowptr, col)
+_EdgeTensorType = Union[Tensor, Tuple[Tensor, Tensor]]
+
+
+class NeighborSamplingRAGGraphStore(LocalGraphStore):
+    """Neighbor sampling based graph-store to store & retrieve graph data."""
+    def __init__(  # type: ignore[no-untyped-def]
+        self,
+        feature_store: Optional[FeatureStore] = None,
+        **kwargs,
+    ):
+        """Initializes the graph store.
+        Optional feature store and neighbor sampling settings.
+
+        Args:
+        feature_store (optional): The feature store to use.
+            None if not yet registered.
+        **kwargs (optional):
+            Additional keyword arguments for neighbor sampling.
+        """
+        self.feature_store = feature_store
+        self.sample_kwargs = kwargs
+        self._sampler_is_initialized = False
+        self._config: Dict[str, Any] = {}
+
+        # to be set by the config
+        self.num_neighbors = None
+        super().__init__()
+
+    @property
+    def config(self) -> Dict[str, Any]:
+        """Get the config for the feature store."""
+        return self._config
+
+    def _set_from_config(self, config: Dict[str, Any], attr_name: str) -> None:
+        """Set an attribute from the config.
+
+        Args:
+            config (Dict[str, Any]): Config dictionary
+            attr_name (str): Name of attribute to set
+
+        Raises:
+            ValueError: If required attribute not found in config
+        """
+        if attr_name not in config:
+            raise ValueError(
+                f"Required config parameter '{attr_name}' not found")
+        setattr(self, attr_name, config[attr_name])
+
+    @config.setter  # type: ignore
+    def config(self, config: Dict[str, Any]) -> None:
+        """Set the config for the feature store.
+
+        Args:
+            config (Dict[str, Any]):
+                Config dictionary containing required parameters
+
+        Raises:
+            ValueError: If required parameters missing from config
+        """
+        self._set_from_config(config, "num_neighbors")
+        if hasattr(self, 'sampler'):
+            self.sampler.num_neighbors = (  # type: ignore[has-type]
+                self.num_neighbors)
+
+        self._config = config
+
+    def _init_sampler(self) -> None:
+        """Initializes neighbor sampler with the registered feature store."""
+        if self.feature_store is None:
+            raise AttributeError("Feature store not registered yet.")
+        assert self.num_neighbors is not None, \
+            "Please set num_neighbors through config"
+        self.sampler = BidirectionalNeighborSampler(
+            data=(self.feature_store, self), num_neighbors=self.num_neighbors,
+            **self.sample_kwargs)
+        self._sampler_is_initialized = True
+
+    def register_feature_store(self, feature_store: FeatureStore) -> None:
+        """Registers a feature store with the graph store.
+
+        :param feature_store: The feature store to register.
+        """
+        self.feature_store = feature_store
+        self._sampler_is_initialized = False
+
+    def put_edge_id(  # type: ignore[no-untyped-def]
+            self, edge_id: Tensor, *args, **kwargs) -> bool:
+        """Stores an edge ID in the graph store.
+
+        :param edge_id: The edge ID to store.
+        :return: Whether the operation was successful.
+        """
+        ret = super().put_edge_id(edge_id.contiguous(), *args, **kwargs)
+        self._sampler_is_initialized = False
+        return ret
+
+    @property
+    def edge_index(self) -> _EdgeTensorType:
+        """Gets the edge index of the graph.
+
+        :return: The edge index as a tensor.
+        """
+        return self.get_edge_index(*self.edge_idx_args, **self.edge_idx_kwargs)
+
+    def put_edge_index(  # type: ignore[no-untyped-def]
+            self, edge_index: _EdgeTensorType, *args, **kwargs) -> bool:
+        """Stores an edge index in the graph store.
+
+        :param edge_index: The edge index to store.
+        :return: Whether the operation was successful.
+        """
+        ret = super().put_edge_index(edge_index, *args, **kwargs)
+        # HACK
+        self.edge_idx_args = args
+        self.edge_idx_kwargs = kwargs
+        self._sampler_is_initialized = False
+        return ret
+
+    # HACKY
+    @edge_index.setter  # type: ignore
+    def edge_index(self, edge_index: _EdgeTensorType) -> None:
+        """Sets the edge index of the graph.
+
+        :param edge_index: The edge index to set.
+        """
+        # correct since we make node list from triples
+        if isinstance(edge_index, Tensor):
+            num_nodes = int(edge_index.max()) + 1
+        else:
+            assert isinstance(edge_index, tuple) \
+                and isinstance(edge_index[0], Tensor) \
+                and isinstance(edge_index[1], Tensor), \
+                "edge_index must be a Tensor of [2, num_edges] \
+                or a tuple of Tensors, (row, col)."
+
+            num_nodes = int(edge_index[0].max()) + 1
+        attr = dict(
+            edge_type=None,
+            layout='coo',
+            size=(num_nodes, num_nodes),
+            is_sorted=False,
+        )
+        # edge index needs to be sorted here and the perm saved for later
+        col_sorted, self.perm = index_sort(edge_index[1], num_nodes,
+                                           stable=True)
+        row_sorted = edge_index[0][self.perm]
+        edge_index_sorted = torch.stack([row_sorted, col_sorted], dim=0)
+        self.put_edge_index(edge_index_sorted, **attr)
+
+    def sample_subgraph(
+        self,
+        seed_nodes: Tensor,
+    ) -> SamplerOutput:
+        """Sample the graph starting from the given nodes using the
+        in-built NeighborSampler.
+
+        Args:
+            seed_nodes (InputNodes): Seed nodes to start sampling from.
+            num_neighbors (Optional[NumNeighborsType], optional): Parameters
+                to determine how many hops and number of neighbors per hop.
+                Defaults to None.
+
+        Returns:
+            Union[SamplerOutput, HeteroSamplerOutput]: NeighborSamplerOutput
+                for the input.
+        """
+        # TODO add support for Hetero
+        if not self._sampler_is_initialized:
+            self._init_sampler()
+
+        seed_nodes = seed_nodes.unique().contiguous()
+        node_sample_input = NodeSamplerInput(input_id=None, node=seed_nodes)
+        out = self.sampler.sample_from_nodes(  # type: ignore[has-type]
+            node_sample_input)
+
+        # edge ids need to be remapped to the original indices
+        out.edge = self.perm[out.edge]
+
+        return out

torch_geometric/llm/utils/vectorrag.py

@@ -0,0 +1,125 @@
+# mypy: ignore-errors
+import os
+from abc import abstractmethod
+from typing import Any, Callable, Dict, List, Optional, Protocol, Union
+
+import torch
+from torch import Tensor
+
+from torch_geometric.data import Data
+from torch_geometric.llm.models import SentenceTransformer
+from torch_geometric.llm.utils.backend_utils import batch_knn
+
+
+class VectorRetriever(Protocol):
+    """Protocol for VectorRAG."""
+    @abstractmethod
+    def query(self, query: Any, **kwargs: Optional[Dict[str, Any]]) -> Data:
+        """Retrieve a context for a given query."""
+        ...
+
+
+class DocumentRetriever(VectorRetriever):
+    """Retrieve documents from a vector database."""
+    def __init__(self, raw_docs: List[str],
+                 embedded_docs: Optional[Tensor] = None, k_for_docs: int = 2,
+                 model: Optional[Union[SentenceTransformer, torch.nn.Module,
+                                       Callable]] = None,
+                 model_kwargs: Optional[Dict[str, Any]] = None):
+        """Retrieve documents from a vector database.
+
+        Args:
+            raw_docs: List[str]: List of raw documents.
+            embedded_docs: Optional[Tensor]: Embedded documents.
+            k_for_docs: int: Number of documents to retrieve.
+            model: Optional[Union[SentenceTransformer, torch.nn.Module]]:
+                Model to use for encoding.
+            model_kwargs: Optional[Dict[str, Any]]:
+                Keyword arguments to pass to the model.
+        """
+        self.raw_docs = raw_docs
+        self.embedded_docs = embedded_docs
+        self.k_for_docs = k_for_docs
+        self.model = model
+
+        if self.model is not None:
+            self.encoder = self.model
+            self.model_kwargs = model_kwargs
+
+        if self.embedded_docs is None:
+            assert self.model is not None, \
+                "Model must be provided if embedded_docs is not provided"
+            self.model_kwargs = model_kwargs or {}
+            self.embedded_docs = self.encoder(self.raw_docs,
+                                              **self.model_kwargs)
+            # we don't want to print the verbose output in `query`
+            self.model_kwargs.pop("verbose", None)
+
+    def query(self, query: Union[str, Tensor]) -> List[str]:
+        """Retrieve documents from the vector database.
+
+        Args:
+            query: Union[str, Tensor]: Query to retrieve documents for.
+
+        Returns:
+            List[str]: Documents retrieved from the vector database.
+        """
+        if isinstance(query, str):
+            with torch.no_grad():
+                query_enc = self.encoder(query, **self.model_kwargs)
+        else:
+            query_enc = query
+
+        selected_doc_idxs, _ = next(
+            batch_knn(query_enc, self.embedded_docs, self.k_for_docs))
+        return [self.raw_docs[i] for i in selected_doc_idxs]
+
+    def save(self, path: str) -> None:
+        """Save the DocumentRetriever instance to disk.
+
+        Args:
+            path: str: Path where to save the retriever.
+        """
+        os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+
+        # Prepare data to save
+        save_dict = {
+            'raw_docs': self.raw_docs,
+            'embedded_docs': self.embedded_docs,
+            'k_for_docs': self.k_for_docs,
+        }
+
+        # We do not serialize the model
+        torch.save(save_dict, path)
+
+    @classmethod
+    def load(cls, path: str, model: Union[SentenceTransformer, torch.nn.Module,
+                                          Callable],
+             model_kwargs: Optional[Dict[str, Any]] = None) -> VectorRetriever:
+        """Load a DocumentRetriever instance from disk.
+
+        Args:
+            path: str: Path to the saved retriever.
+            model: Union[SentenceTransformer, torch.nn.Module, Callable]:
+                Model to use for encoding.
+                If None, the saved model will be used if available.
+            model_kwargs: Optional[Dict[str, Any]]
+                Key word args to be passed to model
+
+        Returns:
+            DocumentRetriever: The loaded retriever.
+        """
+        if not os.path.exists(path):
+            raise FileNotFoundError(
+                f"No saved document retriever found at {path}")
+
+        save_dict = torch.load(path, weights_only=False)
+        if save_dict['embedded_docs'] is not None \
+                and isinstance(save_dict['embedded_docs'], Tensor)\
+                and model_kwargs is not None:
+            model_kwargs.pop("verbose", None)
+        # Create a new DocumentRetriever with the loaded data
+        return cls(raw_docs=save_dict['raw_docs'],
+                   embedded_docs=save_dict['embedded_docs'],
+                   k_for_docs=save_dict['k_for_docs'], model=model,
+                   model_kwargs=model_kwargs)

torch_geometric/loader/cluster.py

@@ -235,9 +235,9 @@ class ClusterData(torch.utils.data.Dataset):
 class ClusterLoader(torch.utils.data.DataLoader):
     r"""The data loader scheme from the `"Cluster-GCN: An Efficient Algorithm
     for Training Deep and Large Graph Convolutional Networks"
-    <https://arxiv.org/abs/1905.07953>`_ paper which merges partioned subgraphs
-    and their between-cluster links from a large-scale graph data object to
-    form a mini-batch.
+    <https://arxiv.org/abs/1905.07953>`_ paper which merges partitioned
+    subgraphs and their between-cluster links from a large-scale graph data
+    object to form a mini-batch.
 
     .. note::
 
@@ -252,7 +252,7 @@ class ClusterLoader(torch.utils.data.DataLoader):
 
     Args:
         cluster_data (torch_geometric.loader.ClusterData): The already
-            partioned data object.
+            partitioned data object.
         **kwargs (optional): Additional arguments of
             :class:`torch.utils.data.DataLoader`, such as :obj:`batch_size`,
             :obj:`shuffle`, :obj:`drop_last` or :obj:`num_workers`.

torch_geometric/loader/ibmb_loader.py

@@ -148,7 +148,7 @@ def indices_complete_check(
         if isinstance(aux, Tensor):
             aux = aux.cpu().numpy()
 
-        assert np.all(np.in1d(out,
+        assert np.all(np.isin(out,
                               aux)), "Not all output nodes are in aux nodes!"
         outs.append(out)
 
@@ -236,7 +236,7 @@ def create_batchwise_out_aux_pairs(
             logits[tele_set, i] = 1. / len(tele_set)
 
         new_logits = logits.clone()
-        for i in range(num_iter):
+        for _ in range(num_iter):
             new_logits = adj @ new_logits * (1 - alpha) + alpha * logits
 
         inds = new_logits.argsort(0)
@@ -498,7 +498,7 @@ class IBMBBaseLoader(torch.utils.data.DataLoader):
             assert adj is not None
 
         for out, aux in pbar:
-            mask = torch.from_numpy(np.in1d(aux, out))
+            mask = torch.from_numpy(np.isin(aux, out))
             if isinstance(aux, np.ndarray):
                 aux = torch.from_numpy(aux)
             subg = get_subgraph(aux, graph, return_edge_index_type, adj,
@@ -541,7 +541,7 @@ class IBMBBaseLoader(torch.utils.data.DataLoader):
         out, aux = zip(*data_list)
         out = np.concatenate(out)
         aux = np.unique(np.concatenate(aux))
-        mask = torch.from_numpy(np.in1d(aux, out))
+        mask = torch.from_numpy(np.isin(aux, out))
         aux = torch.from_numpy(aux)
 
         subg = get_subgraph(aux, self.graph, self.return_edge_index_type,

torch_geometric/loader/link_loader.py

@@ -70,7 +70,7 @@ class LinkLoader(
             :obj:`edge_label_index`. If set, temporal sampling will be
             used such that neighbors are guaranteed to fulfill temporal
             constraints, *i.e.*, neighbors have an earlier timestamp than
-            the ouput edge. The :obj:`time_attr` needs to be set for this
+            the output edge. The :obj:`time_attr` needs to be set for this
             to work. (default: :obj:`None`)
         neg_sampling (NegativeSampling, optional): The negative sampling
             configuration.

torch_geometric/loader/link_neighbor_loader.py

@@ -117,7 +117,7 @@ class LinkNeighborLoader(LinkLoader):
             :obj:`edge_label_index`. If set, temporal sampling will be
             used such that neighbors are guaranteed to fulfill temporal
             constraints, *i.e.*, neighbors have an earlier timestamp than
-            the ouput edge. The :obj:`time_attr` needs to be set for this
+            the output edge. The :obj:`time_attr` needs to be set for this
             to work. (default: :obj:`None`)
         replace (bool, optional): If set to :obj:`True`, will sample with
             replacement. (default: :obj:`False`)
@@ -170,6 +170,7 @@ class LinkNeighborLoader(LinkLoader):
             negative sampling mode.
             If set to :obj:`None`, no negative sampling strategy is applied.
             (default: :obj:`None`)
+            For example use obj:`neg_sampling=dict(mode= 'binary', amount=0.5)`
         neg_sampling_ratio (int or float, optional): The ratio of sampled
             negative edges to the number of positive edges.
             Deprecated in favor of the :obj:`neg_sampling` argument.

torch_geometric/loader/mixin.py

@@ -106,9 +106,9 @@ class MultithreadingMixin:
     def _mt_init_fn(self, worker_id: int) -> None:
         try:
             torch.set_num_threads(int(self._worker_threads))
-        except IndexError:
+        except IndexError as e:
             raise ValueError(f"Cannot set {self.worker_threads} threads "
-                             f"in worker {worker_id}")
+                             f"in worker {worker_id}") from e
 
         # Chain worker init functions:
         self._old_worker_init_fn(worker_id)
@@ -213,9 +213,9 @@ class AffinityMixin:
 
             psutil.Process().cpu_affinity(worker_cores)
 
-        except IndexError:
+        except IndexError as e:
             raise ValueError(f"Cannot use CPU affinity for worker ID "
-                             f"{worker_id} on CPU {self.loader_cores}")
+                             f"{worker_id} on CPU {self.loader_cores}") from e
 
         # Chain worker init functions:
         self._old_worker_init_fn(worker_id)
@@ -248,7 +248,8 @@ class AffinityMixin:
             warnings.warn(
                 "Due to conflicting parallelization methods it is not advised "
                 "to use affinitization with 'HeteroData' datasets. "
-                "Use `enable_multithreading` for better performance.")
+                "Use `enable_multithreading` for better performance.",
+                stacklevel=2)
 
         self.loader_cores = loader_cores[:] if loader_cores else None
         if self.loader_cores is None:

torch_geometric/loader/neighbor_loader.py

@@ -14,7 +14,7 @@ class NeighborLoader(NodeLoader):
     This loader allows for mini-batch training of GNNs on large-scale graphs
     where full-batch training is not feasible.
 
-    More specifically, :obj:`num_neighbors` denotes how much neighbors are
+    More specifically, :obj:`num_neighbors` denotes how many neighbors are
     sampled for each node in each iteration.
     :class:`~torch_geometric.loader.NeighborLoader` takes in this list of
     :obj:`num_neighbors` and iteratively samples :obj:`num_neighbors[i]` for

torch_geometric/loader/neighbor_sampler.py

@@ -72,9 +72,9 @@ class NeighborSampler(torch.utils.data.DataLoader):
         `examples/reddit.py
         <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
         reddit.py>`_ or
-        `examples/ogbn_products_sage.py
+        `examples/ogbn_train.py
         <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
-        ogbn_products_sage.py>`_.
+        ogbn_train.py>`_.
 
     Args:
         edge_index (Tensor or SparseTensor): A :obj:`torch.LongTensor` or a

torch_geometric/loader/prefetch.py

@@ -27,8 +27,9 @@ class DeviceHelper:
 
         if ((self.device.type == 'cuda' and not with_cuda)
                 or (self.device.type == 'xpu' and not with_xpu)):
-            warnings.warn(f"Requested device '{self.device.type}' is not "
-                          f"available, falling back to CPU")
+            warnings.warn(
+                f"Requested device '{self.device.type}' is not "
+                f"available, falling back to CPU", stacklevel=2)
             self.device = torch.device('cpu')
 
         self.stream = None

torch_geometric/loader/temporal_dataloader.py

@@ -6,7 +6,7 @@ from torch_geometric.data import TemporalData
 
 
 class TemporalDataLoader(torch.utils.data.DataLoader):
-    r"""A data loader which merges succesive events of a
+    r"""A data loader which merges successive events of a
     :class:`torch_geometric.data.TemporalData` to a mini-batch.
 
     Args:
@@ -15,7 +15,7 @@ class TemporalDataLoader(torch.utils.data.DataLoader):
         batch_size (int, optional): How many samples per batch to load.
             (default: :obj:`1`)
         neg_sampling_ratio (float, optional): The ratio of sampled negative
-            destination nodes to the number of postive destination nodes.
+            destination nodes to the number of positive destination nodes.
             (default: :obj:`0.0`)
         **kwargs (optional): Additional arguments of
             :class:`torch.utils.data.DataLoader`.

torch_geometric/loader/utils.py

@@ -178,7 +178,7 @@ def filter_hetero_data(
     out = copy.copy(data)
 
     for node_type in out.node_types:
-        # Handle the case of disconneted graph sampling:
+        # Handle the case of disconnected graph sampling:
         if node_type not in node_dict:
             node_dict[node_type] = torch.empty(0, dtype=torch.long)
 
@@ -186,7 +186,7 @@ def filter_hetero_data(
                            node_dict[node_type])
 
     for edge_type in out.edge_types:
-        # Handle the case of disconneted graph sampling:
+        # Handle the case of disconnected graph sampling:
         if edge_type not in row_dict:
             row_dict[edge_type] = torch.empty(0, dtype=torch.long)
         if edge_type not in col_dict:
@@ -256,14 +256,6 @@ def filter_custom_hetero_store(
     # Construct a new `HeteroData` object:
     data = custom_cls() if custom_cls is not None else HeteroData()
 
-    # Filter edge storage:
-    # TODO support edge attributes
-    for attr in graph_store.get_all_edge_attrs():
-        key = attr.edge_type
-        if key in row_dict and key in col_dict:
-            edge_index = torch.stack([row_dict[key], col_dict[key]], dim=0)
-            data[attr.edge_type].edge_index = edge_index
-
     # Filter node storage:
     required_attrs = []
     for attr in feature_store.get_all_tensor_attrs():
@@ -280,6 +272,14 @@ def filter_custom_hetero_store(
     for i, attr in enumerate(required_attrs):
         data[attr.group_name][attr.attr_name] = tensors[i]
 
+    # Filter edge storage:
+    # TODO support edge attributes
+    for attr in graph_store.get_all_edge_attrs():
+        key = attr.edge_type
+        if key in row_dict and key in col_dict:
+            edge_index = torch.stack([row_dict[key], col_dict[key]], dim=0)
+            data[attr.edge_type].edge_index = edge_index
+
     return data
 
 

torch_geometric/metrics/__init__.py

@@ -1,21 +1,35 @@
 # flake8: noqa
 
 from .link_pred import (
+    LinkPredMetric,
+    LinkPredMetricCollection,
     LinkPredPrecision,
     LinkPredRecall,
     LinkPredF1,
     LinkPredMAP,
     LinkPredNDCG,
     LinkPredMRR,
+    LinkPredHitRatio,
+    LinkPredCoverage,
+    LinkPredDiversity,
+    LinkPredPersonalization,
+    LinkPredAveragePopularity,
 )
 
 link_pred_metrics = [
+    'LinkPredMetric',
+    'LinkPredMetricCollection',
     'LinkPredPrecision',
     'LinkPredRecall',
     'LinkPredF1',
     'LinkPredMAP',
     'LinkPredNDCG',
     'LinkPredMRR',
+    'LinkPredHitRatio',
+    'LinkPredCoverage',
+    'LinkPredDiversity',
+    'LinkPredPersonalization',
+    'LinkPredAveragePopularity',
 ]
 
 __all__ = link_pred_metrics