pyg-nightly 2.6.0.dev20240909__py3-none-any.whl → 2.6.0.dev20240910__py3-none-any.whl

{pyg_nightly-2.6.0.dev20240909.dist-info → pyg_nightly-2.6.0.dev20240910.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pyg-nightly
-Version: 2.6.0.dev20240909
+Version: 2.6.0.dev20240910
 Summary: Graph Neural Network Library for PyTorch
 Keywords: deep-learning,pytorch,geometric-deep-learning,graph-neural-networks,graph-convolutional-networks
 Author-email: Matthias Fey <matthias@pyg.org>

{pyg_nightly-2.6.0.dev20240909.dist-info → pyg_nightly-2.6.0.dev20240910.dist-info}/RECORD

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=8YENZOJAbIgDzXkgwPPNDmQik7Kny9kIWjNlWiAnJYg,1904
+torch_geometric/__init__.py,sha256=vRKXyHIGqBHJUKx9tLap5c3uB1Mbb6ZOlvVgapW_D6Q,1904
 torch_geometric/_compile.py,sha256=0HAdz6MGmyrgi4g6P-PorTg8dPIKx3Jo4zVJavrlfX0,1139
 torch_geometric/_onnx.py,sha256=V9ffrIKSqhDw6xUZ12lkuSfNs48cQp2EeJ6Z19GfnVw,349
 torch_geometric/backend.py,sha256=lVaf7aLoVaB3M-UcByUJ1G4T4FOK6LXAg0CF4W3E8jo,1575
@@ -416,7 +416,7 @@ torch_geometric/nn/kge/distmult.py,sha256=dGQ0bVzjreZgFN1lXE23_IIidsiOq7ehPrMb-N
 torch_geometric/nn/kge/loader.py,sha256=5Uc1j3OUMQnBYSHDqL7pLCty1siFLzoPkztigYO2zP8,771
 torch_geometric/nn/kge/rotate.py,sha256=XLuO1AbyTt5cJxr97ZzoyAyIEsHKesgW5TvDmnGJAao,3208
 torch_geometric/nn/kge/transe.py,sha256=jlejq5BLMm-sb1wWcLDp7pZqCdelWBgjDIC8ctbjSdU,3088
-torch_geometric/nn/models/__init__.py,sha256=_2KGXTo3eAgwcwAo0xIpw_I41n_cuTimJTZQgG0zKEc,1963
+torch_geometric/nn/models/__init__.py,sha256=RpYFFqaYWq1BVMF3Fs-EQo-QZDdLQjIHPdkl3d2MOW4,2017
 torch_geometric/nn/models/attentive_fp.py,sha256=tkgvw28wg9-JqHIfBllfCwTHrZIUiv85yZJcDqjz3z0,6634
 torch_geometric/nn/models/autoencoder.py,sha256=nGje-zty78Y3hxOJ9o0_6QziJjOvBlknk6z0_fDQwQU,10770
 torch_geometric/nn/models/basic_gnn.py,sha256=PGa0RUMyvrNy_5yRI2jX_zwPsmZXwOQWfsWvxOiHsSk,31225
@@ -426,6 +426,7 @@ torch_geometric/nn/models/deep_graph_infomax.py,sha256=u6j-5-iHBASDCZ776dyfCI1N8
 torch_geometric/nn/models/deepgcn.py,sha256=tIgT03cj8MghYlxEozpoGvGG_CwpJrGDxv1Z0CVIUts,4339
 torch_geometric/nn/models/dimenet.py,sha256=Kc5p-rB5q-0e8lY22l-OdQTscTxJh2lTEpeRFMdL4RY,36186
 torch_geometric/nn/models/dimenet_utils.py,sha256=xP_nbzkSSL25GC3rrZ9KP8x9QZ59S-CZuHzCmQ-K0fI,5062
+torch_geometric/nn/models/g_retriever.py,sha256=uH_aYrFbFNHaAeKQn_LtUgP5ajutLYYD8N9UvSKcpfk,7271
 torch_geometric/nn/models/gnnff.py,sha256=15dkiLgy0LmH1hnUrpeoHioIp4BPTfjpVATpnGRt9E0,7860
 torch_geometric/nn/models/graph_mixer.py,sha256=mthMeCOikR8gseEsu4oJ3Cd9C35zHSv1p32ROwnG-6s,9246
 torch_geometric/nn/models/graph_unet.py,sha256=WFb7d_DBByMGyXh3AdK2CKNmvMmSKsSUt8l8UnSOovs,5395
@@ -448,7 +449,7 @@ torch_geometric/nn/models/signed_gcn.py,sha256=J40CnedFIqtKI1LhW1ITSEFRbA_XiJZL6
 torch_geometric/nn/models/tgn.py,sha256=kEGdfLJybkbMT4UMoAh2nCzfX3_nDjfm1cicuPHEwAM,11878
 torch_geometric/nn/models/visnet.py,sha256=97OFMCsPDEI5BCSi7RhoRcU2CNRp7zck2tEzrltFZj4,43192
 torch_geometric/nn/nlp/__init__.py,sha256=JJESTA7w_K8v60XbCd25IqmrKKHLz5OiNexMHYGV2mE,138
-torch_geometric/nn/nlp/llm.py,sha256=blCLWkm76bKMxGgOj7dxMXcyn9ecngX1LKDkP-MRSW4,10824
+torch_geometric/nn/nlp/llm.py,sha256=KwSXgI55FuHLR_9vhgekDXMaRUodPQceHPD7OCp2KN4,11639
 torch_geometric/nn/nlp/sentence_transformer.py,sha256=DzbQO8wgR34BkKpXfMqQu61hMrK94W2MBa3bZ4fDmVs,3114
 torch_geometric/nn/norm/__init__.py,sha256=u2qIDrkbeuObGVXSAIftAlvSd6ouGTtxznCfD-59UiA,669
 torch_geometric/nn/norm/batch_norm.py,sha256=sJKrinHGwA-noIgteg1RD2W06rd0zskD-rXuY-36glY,8283
@@ -460,13 +461,14 @@ torch_geometric/nn/norm/layer_norm.py,sha256=pWo5q8rLNSaU2fECpP7L8T_airtaukjOztL
 torch_geometric/nn/norm/mean_subtraction_norm.py,sha256=KVHOp413mw7obwAN09Le6XdgobtCXpi4UKpjpG1M550,1322
 torch_geometric/nn/norm/msg_norm.py,sha256=zaQtqhs55LU-e6KPC4ylaSdge4KvEoseqOt7pmAzi2s,1662
 torch_geometric/nn/norm/pair_norm.py,sha256=IfHMiVYw_xsy035NakbPGdQVaVC-Ue3Oxwo651Vc47I,2824
-torch_geometric/nn/pool/__init__.py,sha256=pJsD4qumvCu_oZUtC-orZCHp9nObx-VMWHFlJckFrHc,14129
+torch_geometric/nn/pool/__init__.py,sha256=2Bi-_xlsGIUUKDeOO7BhaTqCc5n6_ixbu_MO9pglMts,14192
 torch_geometric/nn/pool/approx_knn.py,sha256=n7C8Cbar6o5tJcuAbzhM5hqMK26hW8dm5DopuocidO0,3967
 torch_geometric/nn/pool/asap.py,sha256=p8fwpMOeCUyJrdvMmLoTMzr0tI9YCTnefMx8ylIv5xE,6683
 torch_geometric/nn/pool/avg_pool.py,sha256=pwiQh14BCVsT-iULqVAFW-Dxt7DjFOu8CQX_Hu34vZc,3966
+torch_geometric/nn/pool/cluster_pool.py,sha256=et2YaFu1kf-o6Eg9XpqHGp_Cqv68DndWbE88VJHOSPQ,5227
 torch_geometric/nn/pool/consecutive.py,sha256=7dMiMd5IybNeml1RqZq436FI6sod5ZUxTuDWJjr5syo,273
 torch_geometric/nn/pool/decimation.py,sha256=AjbU2h_Gl_EQcfkhF977EnrLJ2kait_e4HyCNKRyxPw,1601
-torch_geometric/nn/pool/edge_pool.py,sha256=ZdDv0t1CYwdgg56V8oFTvYozHHzerMoltEtVsY-9Wv8,8581
+torch_geometric/nn/pool/edge_pool.py,sha256=cXgcN5xF8z5NeycYMX9m1zoAk1jtSdyK42YiNNHTeow,8571
 torch_geometric/nn/pool/glob.py,sha256=RJrq1sgAe8oV15WSGtXgB6yXWj2irSJIWAdQLb0byN4,3492
 torch_geometric/nn/pool/graclus.py,sha256=dL9tasXNM-x2NOMRJn8k6z4CeW46nRzoa49IzG58wow,1349
 torch_geometric/nn/pool/knn.py,sha256=fNZV0q2A4lzhZQyePRLHSrtuWjbxQxvv3V7oeNzBLVk,11343
@@ -615,6 +617,6 @@ torch_geometric/utils/undirected.py,sha256=H_nfpI0_WluOG6VfjPyldvcjL4w5USAKWu2x5
 torch_geometric/visualization/__init__.py,sha256=PyR_4K5SafsJrBr6qWrkjKr6GBL1b7FtZybyXCDEVwY,154
 torch_geometric/visualization/graph.py,sha256=SvbdVx5Zmuy_WSSA4-WWCkqAcCSHVe84mjMfsEWbZCs,4813
 torch_geometric/visualization/influence.py,sha256=CWMvuNA_Nf1sfbJmQgn58yS4OFpeKXeZPe7kEuvkUBw,477
-pyg_nightly-2.6.0.dev20240909.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-pyg_nightly-2.6.0.dev20240909.dist-info/METADATA,sha256=TyeRwIRUgTAzeT0sDoNXKBQ80n5fjfhulW918W_iAWs,63068
-pyg_nightly-2.6.0.dev20240909.dist-info/RECORD,,
+pyg_nightly-2.6.0.dev20240910.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+pyg_nightly-2.6.0.dev20240910.dist-info/METADATA,sha256=4d3E2ca0L5gmd30HK26OTddin11CJkS3iteHVYcnxfI,63068
+pyg_nightly-2.6.0.dev20240910.dist-info/RECORD,,

torch_geometric/__init__.py

@@ -30,7 +30,7 @@ from .lazy_loader import LazyLoader
 contrib = LazyLoader('contrib', globals(), 'torch_geometric.contrib')
 graphgym = LazyLoader('graphgym', globals(), 'torch_geometric.graphgym')
 
-__version__ = '2.6.0.dev20240909'
+__version__ = '2.6.0.dev20240910'
 
 __all__ = [
     'Index',

torch_geometric/nn/models/__init__.py

@@ -28,6 +28,7 @@ from .gnnff import GNNFF
 from .pmlp import PMLP
 from .neural_fingerprint import NeuralFingerprint
 from .visnet import ViSNet
+from .g_retriever import GRetriever
 
 # Deprecated:
 from torch_geometric.explain.algorithm.captum import (to_captum_input,
@@ -75,4 +76,5 @@ __all__ = classes = [
     'PMLP',
     'NeuralFingerprint',
     'ViSNet',
+    'GRetriever',
 ]

torch_geometric/nn/models/g_retriever.py

@@ -0,0 +1,205 @@
+from typing import List, Optional
+
+import torch
+from torch import Tensor
+
+from torch_geometric.nn.models import GAT
+from torch_geometric.nn.nlp.llm import BOS, LLM, MAX_NEW_TOKENS
+from torch_geometric.utils import scatter
+
+
+class GRetriever(torch.nn.Module):
+    r"""The G-Retriever model from the `"G-Retriever: Retrieval-Augmented
+    Generation for Textual Graph Understanding and Question Answering"
+    <https://arxiv.org/abs/2402.07630>`_ paper.
+
+    Args:
+        llm (LLM): The LLM to use.
+        gnn (torch.nn.Module): The GNN to use.
+        use_lora (bool, optional): If set to :obj:`True`, will use LORA from
+            :obj:`peft` for training the LLM, see
+            `here <https://huggingface.co/docs/peft/en/index>`_ for details.
+            (default: :obj:`False`)
+        mlp_out_channels (int, optional): The size of each graph embedding
+            after projection. (default: :obj:`4096`)
+
+    .. warning::
+        This module has been tested with the following HuggingFace models
+
+        * :obj:`llm_to_use="meta-llama/Llama-2-7b-chat-hf"`
+        * :obj:`llm_to_use="google/gemma-7b"`
+
+        and may not work with other models. See other models at `HuggingFace
+        Models <https://huggingface.co/models>`_ and let us know if you
+        encounter any issues.
+
+    .. note::
+        For an example of using :class:`GRetriever`, see
+        `examples/llm/g_retriever.py <https://github.com/pyg-team/
+        pytorch_geometric/blob/master/examples/llm/g_retriever.py>`_.
+    """
+    def __init__(
+        self,
+        llm: LLM,
+        gnn: torch.nn.Module,
+        use_lora: bool = False,
+        gnn_to_use=GAT,
+        mlp_out_channels: int = 4096,
+    ) -> None:
+        super().__init__()
+
+        self.llm = llm
+        self.gnn = gnn.to(self.llm.device)
+
+        self.word_embedding = self.llm.word_embedding
+        self.llm_generator = self.llm.llm
+        if use_lora:
+            from peft import (
+                LoraConfig,
+                get_peft_model,
+                prepare_model_for_kbit_training,
+            )
+            self.llm_generator = prepare_model_for_kbit_training(
+                self.llm_generator)
+            lora_r: int = 8
+            lora_alpha: int = 16
+            lora_dropout: float = 0.05
+            lora_target_modules = ['q_proj', 'v_proj']
+            config = LoraConfig(
+                r=lora_r,
+                lora_alpha=lora_alpha,
+                target_modules=lora_target_modules,
+                lora_dropout=lora_dropout,
+                bias='none',
+                task_type='CAUSAL_LM',
+            )
+            self.llm_generator = get_peft_model(self.llm_generator, config)
+
+        mlp_hidden_channels = self.gnn.out_channels
+        self.projector = torch.nn.Sequential(
+            torch.nn.Linear(mlp_hidden_channels, mlp_hidden_channels),
+            torch.nn.Sigmoid(),
+            torch.nn.Linear(mlp_hidden_channels, mlp_out_channels),
+        ).to(self.llm.device)
+
+    def encode(
+        self,
+        x: Tensor,
+        edge_index: Tensor,
+        batch: Tensor,
+        edge_attr: Optional[Tensor],
+    ) -> Tensor:
+        x = x.to(self.llm.device)
+        edge_index = edge_index.to(self.llm.device)
+        if edge_attr is not None:
+            edge_attr = edge_attr.to(self.llm.device)
+        batch = batch.to(self.llm.device)
+
+        out = self.gnn(x, edge_index, edge_attr=edge_attr)
+        return scatter(out, batch, dim=0, reduce='mean')
+
+    def forward(
+        self,
+        question: List[str],
+        x: Tensor,
+        edge_index: Tensor,
+        batch: Tensor,
+        label: List[str],
+        edge_attr: Optional[Tensor] = None,
+        additional_text_context: Optional[List[str]] = None,
+    ):
+        r"""The forward pass.
+
+        Args:
+            question (List[str]): The questions/prompts.
+            x (torch.Tensor): The input node features.
+            edge_index (torch.Tensor): The edge indices.
+            batch (torch.Tensor): The batch vector
+                :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns
+                each element to a specific example.
+            label (List[str]): The answers/labels.
+            edge_attr (torch.Tensor, optional): The edge features (if supported
+                by the GNN). (default: :obj:`None`)
+            additional_text_context (List[str], optional): Additional context
+                to give to the LLM, such as textified knowledge graphs.
+                (default: :obj:`None`)
+        """
+        x = self.encode(x, edge_index, batch, edge_attr)
+        x = self.projector(x)
+        xs = x.split(x.size(0), dim=0)
+
+        (
+            inputs_embeds,
+            attention_mask,
+            label_input_ids,
+        ) = self.llm._get_embeds(question, additional_text_context, xs, label)
+
+        with self.llm.autocast_context:
+            outputs = self.llm_generator(
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                return_dict=True,
+                labels=label_input_ids,
+            )
+
+        return outputs.loss
+
+    @torch.no_grad()
+    def inference(
+        self,
+        question: List[str],
+        x: Tensor,
+        edge_index: Tensor,
+        batch: Tensor,
+        edge_attr: Optional[Tensor] = None,
+        additional_text_context: Optional[List[str]] = None,
+        max_out_tokens: Optional[int] = MAX_NEW_TOKENS,
+    ):
+        r"""The inference pass.
+
+        Args:
+            question (List[str]): The questions/prompts.
+            x (torch.Tensor): The input node features.
+            edge_index (torch.Tensor): The edge indices.
+            batch (torch.Tensor): The batch vector
+                :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns
+                each element to a specific example.
+            edge_attr (torch.Tensor, optional): The edge features (if supported
+                by the GNN). (default: :obj:`None`)
+            additional_text_context (List[str], optional): Additional context
+                to give to the LLM, such as textified knowledge graphs.
+                (default: :obj:`None`)
+            max_out_tokens (int, optional): How many tokens for the LLM to
+                generate. (default: :obj:`32`)
+        """
+        x = self.encode(x, edge_index, batch, edge_attr)
+        x = self.projector(x)
+        xs = x.split(x.size(0), dim=0)
+
+        inputs_embeds, attention_mask, _ = self.llm._get_embeds(
+            question, additional_text_context, xs)
+
+        bos_token = self.llm.tokenizer(
+            BOS,
+            add_special_tokens=False,
+        ).input_ids[0]
+
+        with self.llm.autocast_context:
+            outputs = self.llm_generator.generate(
+                inputs_embeds=inputs_embeds,
+                max_new_tokens=max_out_tokens,
+                attention_mask=attention_mask,
+                bos_token_id=bos_token,
+                use_cache=True  # Important to set!
+            )
+
+        return self.llm.tokenizer.batch_decode(
+            outputs,
+            skip_special_tokens=True,
+        )
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}(\n'
+                f'  llm={self.llm},\n'
+                f'  gnn={self.gnn},\n'
+                f')')

torch_geometric/nn/nlp/llm.py

@@ -1,10 +1,14 @@
-import warnings
 from contextlib import nullcontext
 from typing import Any, Dict, List, Optional
 
 import torch
 from torch import Tensor
 
+try:
+    from transformers.tokenization_utils_base import BatchEncoding
+except ImportError:
+    BatchEncoding = Dict
+
 BOS = '<s>[INST]'
 EOS_USER = '[/INST]'
 EOS = '[/s]'
@@ -61,23 +65,16 @@ class LLM(torch.nn.Module):
     ) -> None:
         super().__init__()
 
-        from transformers import AutoModelForCausalLM, AutoTokenizer
+        self.model_name = model_name
 
-        if model_name == 'llama2-7b':
-            pretty_model_name = 'LLAMA2'
-            model_name = 'meta-llama/Llama-2-7b-chat-hf'
-        elif model_name == 'gemma':
-            pretty_model_name = 'GEMMA'
-            model_name = 'google/gemma-7b'
-        else:
-            pretty_model_name = model_name
+        from transformers import AutoModelForCausalLM, AutoTokenizer
 
         # A rough heuristic on GPU memory requirements, e.g., we found that
         # LLAMA2 (7B parameters) fits on a 85GB GPU.
         required_memory = 85 * num_params / 7
         kwargs = get_llm_kwargs(required_memory, dtype)
 
-        print(f"Setting up '{pretty_model_name}' with configuration: {kwargs}")
+        print(f"Setting up '{model_name}' with configuration: {kwargs}")
         self.tokenizer = AutoTokenizer.from_pretrained(
             model_name,
             use_fast=False,
@@ -88,17 +85,17 @@ class LLM(torch.nn.Module):
         self.word_embedding = self.llm.model.get_input_embeddings()
 
         if 'max_memory' not in kwargs:  # Pure CPU:
-            self.llm_device = torch.device('cpu')
+            self.device = torch.device('cpu')
             self.autocast_context = nullcontext()
         else:
-            self.llm_device = self.llm.device
+            self.device = self.llm.device
             self.autocast_context = torch.cuda.amp.autocast(dtype=dtype)
 
     def _encode_inputs(
         self,
         question: List[str],
         context: Optional[List[str]] = None,
-    ) -> None:
+    ) -> tuple:
         batch_size = len(question)
         questions = self.tokenizer(question, add_special_tokens=False)
         if context is not None:
@@ -109,14 +106,144 @@ class LLM(torch.nn.Module):
             BOS,
             add_special_tokens=False,
             return_tensors='pt',
-        ).input_ids[0].to(self.llm_device)
+        ).input_ids[0].to(self.device)
         bos_embeds = self.word_embedding(bos_token)
         pad_token = torch.tensor(self.tokenizer.pad_token_id,
-                                 device=self.llm_device)
+                                 device=self.device)
         pad_embeds = self.word_embedding(pad_token).unsqueeze(0)
         return (batch_size, questions, context, eos_user_tokens, bos_embeds,
                 pad_embeds)
 
+    def _label_input_ids(
+        self,
+        i: int,
+        label: BatchEncoding,
+        eos_tokens: BatchEncoding,
+    ) -> List[int]:
+        label_input_ids = label.input_ids[i][:MAX_NEW_TOKENS]
+        label_input_ids = label_input_ids + eos_tokens.input_ids
+        return label_input_ids
+
+    def _input_ids(
+        self,
+        i: int,
+        context: BatchEncoding,
+        question: BatchEncoding,
+        eos_user_tokens: BatchEncoding,
+    ) -> List[int]:
+        input_ids: List[int] = []
+        if context is not None:
+            input_ids += context.input_ids[i][:MAX_TXT_LEN]
+        input_ids += question.input_ids[i]
+        input_ids += eos_user_tokens.input_ids
+        return input_ids
+
+    def _inputs_embeds(
+        self,
+        i: int,
+        input_ids: List[int],
+        bos_embeds: Tensor,
+        embedding: Optional[List[Tensor]] = None,
+    ) -> Tensor:
+        inputs_embeds = self.word_embedding(
+            torch.tensor(input_ids, device=self.device))
+
+        to_cat = [bos_embeds]
+        if embedding is not None and embedding[i] is not None:
+            to_cat.append(embedding[i])
+        to_cat.append(inputs_embeds)
+        return torch.cat(to_cat, dim=0).to(self.device)
+
+    def _append_embeds(
+        self,
+        inputs_embeds: Tensor,
+        batch_inputs_embeds: List[Tensor],
+        batch_attention_mask: List[List[int]],
+        label_input_ids: List[int] = None,
+        batch_label_input_ids: Optional[List[List[int]]] = None,
+    ) -> tuple:
+        batch_inputs_embeds.append(inputs_embeds)
+        batch_attention_mask.append([1] * inputs_embeds.size(0))
+        if label_input_ids is not None:
+            pad = inputs_embeds.size(0) - len(label_input_ids)
+            label_input_ids = [IGNORE_INDEX] * pad + label_input_ids
+            batch_label_input_ids.append(label_input_ids)
+        return batch_inputs_embeds, batch_attention_mask, batch_label_input_ids
+
+    def _pad_embeds(
+        self,
+        pad_embeds: Tensor,
+        batch_inputs_embeds: List[Tensor],
+        batch_attention_mask: List[List[int]],
+        batch_label_input_ids: Optional[List[List[int]]] = None,
+    ) -> tuple:
+        max_length = max([x.size(0) for x in batch_inputs_embeds])
+        batch_size = len(batch_inputs_embeds)
+        for i in range(batch_size):
+            pad = max_length - batch_inputs_embeds[i].size(0)
+            batch_inputs_embeds[i] = torch.cat([
+                pad_embeds.repeat(pad, 1),
+                batch_inputs_embeds[i],
+            ])
+            batch_attention_mask[i] = [0] * pad + batch_attention_mask[i]
+            if batch_label_input_ids is not None:
+                tmp = [IGNORE_INDEX] * pad + batch_label_input_ids[i]
+                batch_label_input_ids[i] = tmp
+        inputs_embeds = torch.stack(batch_inputs_embeds, dim=0)
+        attention_mask = torch.tensor(batch_attention_mask, device=self.device)
+        label_input_ids = None
+        if batch_label_input_ids is not None:
+            label_input_ids = torch.tensor(batch_label_input_ids,
+                                           device=self.device)
+        return inputs_embeds, attention_mask, label_input_ids
+
+    def _get_embeds(
+        self,
+        question: List[str],
+        context: Optional[List[str]] = None,
+        embedding: Optional[List[Tensor]] = None,
+        answer: Optional[List[str]] = None,
+    ) -> tuple:
+        (batch_size, question, context, eos_user_tokens, bos_embeds,
+         pad_embeds) = self._encode_inputs(question, context)
+
+        batch_label_input_ids = None
+        if answer is not None:
+            label = self.tokenizer(answer, add_special_tokens=False)
+            eos_tokens = self.tokenizer(EOS, add_special_tokens=False)
+            batch_label_input_ids = []
+
+        batch_inputs_embeds = []
+        batch_attention_mask = []
+        for i in range(batch_size):
+            input_ids = self._input_ids(i, context, question, eos_user_tokens)
+            if answer is not None:
+                label_input_ids = self._label_input_ids(i, label, eos_tokens)
+                input_ids += label_input_ids
+            else:
+                label_input_ids = None
+
+            inputs_embeds = self._inputs_embeds(i, input_ids, bos_embeds,
+                                                embedding)
+
+            (
+                batch_inputs_embeds,
+                batch_attention_mask,
+                batch_label_input_ids,
+            ) = self._append_embeds(
+                inputs_embeds,
+                batch_inputs_embeds,
+                batch_attention_mask,
+                label_input_ids,
+                batch_label_input_ids,
+            )
+
+        inputs_embeds, attention_mask, label_input_ids = self._pad_embeds(
+            pad_embeds, batch_inputs_embeds, batch_attention_mask,
+            batch_label_input_ids)
+
+        return inputs_embeds, attention_mask, label_input_ids
+
     def forward(
         self,
         question: List[str],
@@ -133,65 +260,11 @@ class LLM(torch.nn.Module):
                 LLM, such as textified knowledge graphs. (default: :obj:`None`)
             embedding (list[torch.Tensor], optional): RAG embedding
                 tensors, *i.e.* the embedded form of :obj:`context`. Either
-                :obj:`context` or :obj:`rag_embeddings` should be used, not
+                :obj:`context` or :obj:`embedding` should be used, not
                 both. (default: :obj:`None`)
         """
-        if context is not None and embedding is not None:
-            warnings.warn("Using both 'context' and 'embedding' is a waste of "
-                          "compute and memory")
-
-        (batch_size, question, context, eos_user_tokens, bos_embeds,
-         pad_embeds) = self._encode_inputs(question, context)
-
-        label = self.tokenizer(answer, add_special_tokens=False)
-        eos_tokens = self.tokenizer(EOS, add_special_tokens=False)
-
-        batch_inputs_embeds = []
-        batch_attention_mask = []
-        batch_label_input_ids = []
-        for i in range(batch_size):
-            label_input_ids = label.input_ids[i][:MAX_NEW_TOKENS]
-            label_input_ids += eos_tokens.input_ids  # Add EOS token.
-
-            input_ids: List[int] = []
-            if context is not None:
-                input_ids += context.input_ids[i][:MAX_TXT_LEN]
-            input_ids += question.input_ids[i]
-            input_ids += eos_user_tokens.input_ids
-            input_ids += label_input_ids
-
-            inputs_embeds = self.word_embedding(
-                torch.tensor(input_ids, device=self.llm_device))
-
-            to_cat = [bos_embeds]
-            if embedding is not None:
-                to_cat.append(embedding[i])
-            to_cat.append(inputs_embeds)
-            inputs_embeds = torch.cat(to_cat, dim=0)
-
-            batch_inputs_embeds.append(inputs_embeds)
-            batch_attention_mask.append([1] * inputs_embeds.size(0))
-            label_input_ids = [IGNORE_INDEX] * (
-                inputs_embeds.size(0) - len(label_input_ids)) + label_input_ids
-            batch_label_input_ids.append(label_input_ids)
-
-        # Pad input embeddings:
-        max_length = max([x.size(0) for x in batch_inputs_embeds])
-        for i in range(batch_size):
-            pad = max_length - batch_inputs_embeds[i].size(0)
-            batch_inputs_embeds[i] = torch.cat([
-                pad_embeds.repeat(pad, 1),
-                batch_inputs_embeds[i],
-            ])
-            batch_attention_mask[i] = [0] * pad + batch_attention_mask[i]
-            batch_label_input_ids[i] = ([IGNORE_INDEX] * pad +
-                                        batch_label_input_ids[i])
-
-        inputs_embeds = torch.stack(batch_inputs_embeds, dim=0)
-        attention_mask = torch.tensor(batch_attention_mask,
-                                      device=self.llm_device)
-        label_input_ids = torch.tensor(batch_label_input_ids,
-                                       device=self.llm_device)
+        inputs_embeds, attention_mask, label_input_ids = self._get_embeds(
+            question, context, embedding, answer)
 
         with self.autocast_context:
             outputs = self.llm(
@@ -219,52 +292,13 @@ class LLM(torch.nn.Module):
                 LLM, such as textified knowledge graphs. (default: :obj:`None`)
             embedding (list[torch.Tensor], optional): RAG embedding
                 tensors, *i.e.* the embedded form of :obj:`context`. Either
-                :obj:`context` or :obj:`rag_embeddings` should be used, not
+                :obj:`context` or :obj:`embedding` should be used, not
                 both. (default: :obj:`None`)
             max_tokens (int, optional): How many tokens for the LLM to
                 generate. (default: :obj:`32`)
         """
-        if context is not None and embedding is not None:
-            warnings.warn("Using both 'context' and 'embedding' is a waste of "
-                          "compute and memory")
-
-        (batch_size, question, context, eos_user_tokens, bos_embeds,
-         pad_embeds) = self._encode_inputs(question, context)
-
-        batch_inputs_embeds = []
-        batch_attention_mask = []
-        for i in range(batch_size):
-            input_ids: List[int] = []
-            if context is not None:
-                input_ids = context.input_ids[i][:MAX_TXT_LEN]
-            input_ids += question.input_ids[i]
-            input_ids += eos_user_tokens.input_ids
-
-            inputs_embeds = self.word_embedding(
-                torch.tensor(input_ids, device=self.llm_device))
-
-            to_cat = [bos_embeds]
-            if embedding is not None:
-                to_cat.append(embedding[i])
-            to_cat.append(inputs_embeds)
-            inputs_embeds = torch.cat(to_cat, dim=0)
-
-            batch_inputs_embeds.append(inputs_embeds)
-            batch_attention_mask.append([1] * inputs_embeds.size(0))
-
-        # Pad input embeddings:
-        max_length = max([x.size(0) for x in batch_inputs_embeds])
-        for i in range(batch_size):
-            pad = max_length - batch_inputs_embeds[i].size(0)
-            batch_inputs_embeds[i] = torch.cat([
-                pad_embeds.repeat(pad, 1),
-                batch_inputs_embeds[i],
-            ])
-            batch_attention_mask[i] = [0] * pad + batch_attention_mask[i]
-
-        inputs_embeds = torch.stack(batch_inputs_embeds, dim=0)
-        attention_mask = torch.tensor(batch_attention_mask,
-                                      device=self.llm_device)
+        inputs_embeds, attention_mask, _ = self._get_embeds(
+            question, context, embedding)
 
         bos_token = self.tokenizer(
             BOS,
@@ -281,3 +315,6 @@ class LLM(torch.nn.Module):
             )
 
         return self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.model_name})'

torch_geometric/nn/pool/__init__.py

@@ -7,18 +7,19 @@ from torch import Tensor
 import torch_geometric.typing
 from torch_geometric.typing import OptTensor, torch_cluster
 
-from .asap import ASAPooling
 from .avg_pool import avg_pool, avg_pool_neighbor_x, avg_pool_x
-from .edge_pool import EdgePooling
 from .glob import global_add_pool, global_max_pool, global_mean_pool
 from .knn import (KNNIndex, L2KNNIndex, MIPSKNNIndex, ApproxL2KNNIndex,
                   ApproxMIPSKNNIndex)
 from .graclus import graclus
 from .max_pool import max_pool, max_pool_neighbor_x, max_pool_x
-from .mem_pool import MemPooling
-from .pan_pool import PANPooling
-from .sag_pool import SAGPooling
 from .topk_pool import TopKPooling
+from .sag_pool import SAGPooling
+from .edge_pool import EdgePooling
+from .cluster_pool import ClusterPooling
+from .asap import ASAPooling
+from .pan_pool import PANPooling
+from .mem_pool import MemPooling
 from .voxel_grid import voxel_grid
 from .approx_knn import approx_knn, approx_knn_graph
 
@@ -344,6 +345,7 @@ __all__ = [
     'TopKPooling',
     'SAGPooling',
     'EdgePooling',
+    'ClusterPooling',
     'ASAPooling',
     'PANPooling',
     'MemPooling',

torch_geometric/nn/pool/cluster_pool.py

@@ -0,0 +1,145 @@
+from typing import NamedTuple, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from torch_geometric.utils import (
+    dense_to_sparse,
+    one_hot,
+    to_dense_adj,
+    to_scipy_sparse_matrix,
+)
+
+
+class UnpoolInfo(NamedTuple):
+    edge_index: Tensor
+    cluster: Tensor
+    batch: Tensor
+
+
+class ClusterPooling(torch.nn.Module):
+    r"""The cluster pooling operator from the `"Edge-Based Graph Component
+    Pooling" <paper url>`_ paper.
+
+    :class:`ClusterPooling` computes a score for each edge.
+    Based on the selected edges, graph clusters are calculated and compressed
+    to one node using the injective :obj:`"sum"` aggregation function.
+    Edges are remapped based on the nodes created by each cluster and the
+    original edges.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        edge_score_method (str, optional): The function to apply
+            to compute the edge score from raw edge scores (:obj:`"tanh"`,
+            :obj:`"sigmoid"`, :obj:`"log_softmax"`). (default: :obj:`"tanh"`)
+        dropout (float, optional): The probability with
+            which to drop edge scores during training. (default: :obj:`0.0`)
+        threshold (float, optional): The threshold of edge scores. If set to
+            :obj:`None`, will be automatically inferred depending on
+            :obj:`edge_score_method`. (default: :obj:`None`)
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        edge_score_method: str = 'tanh',
+        dropout: float = 0.0,
+        threshold: Optional[float] = None,
+    ):
+        super().__init__()
+        assert edge_score_method in ['tanh', 'sigmoid', 'log_softmax']
+
+        if threshold is None:
+            threshold = 0.5 if edge_score_method == 'sigmoid' else 0.0
+
+        self.in_channels = in_channels
+        self.edge_score_method = edge_score_method
+        self.dropout = dropout
+        self.threshhold = threshold
+
+        self.lin = torch.nn.Linear(2 * in_channels, 1)
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        self.lin.reset_parameters()
+
+    def forward(
+        self,
+        x: Tensor,
+        edge_index: Tensor,
+        batch: Tensor,
+    ) -> Tuple[Tensor, Tensor, Tensor, UnpoolInfo]:
+        r"""Forward pass.
+
+        Args:
+            x (torch.Tensor): The node features.
+            edge_index (torch.Tensor): The edge indices.
+            batch (torch.Tensor): Batch vector
+                :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns
+                each node to a specific example.
+
+        Return types:
+            * **x** *(torch.Tensor)* - The pooled node features.
+            * **edge_index** *(torch.Tensor)* - The coarsened edge indices.
+            * **batch** *(torch.Tensor)* - The coarsened batch vector.
+            * **unpool_info** *(UnpoolInfo)* - Information that can be consumed
+              for unpooling.
+        """
+        mask = edge_index[0] != edge_index[1]
+        edge_index = edge_index[:, mask]
+
+        edge_attr = torch.cat(
+            [x[edge_index[0]], x[edge_index[1]]],
+            dim=-1,
+        )
+        edge_score = self.lin(edge_attr).view(-1)
+        edge_score = F.dropout(edge_score, p=self.dropout,
+                               training=self.training)
+
+        if self.edge_score_method == 'tanh':
+            edge_score = edge_score.tanh()
+        elif self.edge_score_method == 'sigmoid':
+            edge_score = edge_score.sigmoid()
+        else:
+            assert self.edge_score_method == 'log_softmax'
+            edge_score = F.log_softmax(edge_score, dim=0)
+
+        return self._merge_edges(x, edge_index, batch, edge_score)
+
+    def _merge_edges(
+        self,
+        x: Tensor,
+        edge_index: Tensor,
+        batch: Tensor,
+        edge_score: Tensor,
+    ) -> Tuple[Tensor, Tensor, Tensor, UnpoolInfo]:
+
+        from scipy.sparse.csgraph import connected_components
+
+        edge_contract = edge_index[:, edge_score > self.threshhold]
+
+        adj = to_scipy_sparse_matrix(edge_contract, num_nodes=x.size(0))
+        _, cluster_np = connected_components(adj, directed=True,
+                                             connection="weak")
+
+        cluster = torch.tensor(cluster_np, dtype=torch.long, device=x.device)
+        C = one_hot(cluster)
+        A = to_dense_adj(edge_index, max_num_nodes=x.size(0)).squeeze(0)
+        S = to_dense_adj(edge_index, edge_attr=edge_score,
+                         max_num_nodes=x.size(0)).squeeze(0)
+
+        A_contract = to_dense_adj(edge_contract,
+                                  max_num_nodes=x.size(0)).squeeze(0)
+        nodes_single = ((A_contract.sum(dim=-1) +
+                         A_contract.sum(dim=-2)) == 0).nonzero()
+        S[nodes_single, nodes_single] = 1.0
+
+        x_out = (S @ C).t() @ x
+        edge_index_out, _ = dense_to_sparse((C.T @ A @ C).fill_diagonal_(0))
+        batch_out = batch.new_empty(x_out.size(0)).scatter_(0, cluster, batch)
+        unpool_info = UnpoolInfo(edge_index, cluster, batch)
+
+        return x_out, edge_index_out, batch_out, unpool_info
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.in_channels})'

torch_geometric/nn/pool/edge_pool.py

@@ -58,7 +58,7 @@ class EdgePooling(torch.nn.Module):
         self,
         in_channels: int,
         edge_score_method: Optional[Callable] = None,
-        dropout: Optional[float] = 0.0,
+        dropout: float = 0.0,
         add_to_edge_score: float = 0.5,
     ):
         super().__init__()