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

torch_geometric/nn/models/molecule_gpt.py

@@ -0,0 +1,222 @@
+from typing import List, Optional
+
+import torch
+from torch import Tensor
+
+from torch_geometric.nn.attention import QFormer
+from torch_geometric.nn.nlp.llm import BOS, LLM, MAX_NEW_TOKENS
+from torch_geometric.utils import to_dense_batch
+
+
+def pad_or_truncate(embeddings: Tensor, max_seq_len: int,
+                    padding_value: int = 0) -> Tensor:
+    batch_size, current_seq_len, d = embeddings.size()
+
+    if current_seq_len > max_seq_len:
+        return embeddings[:, :max_seq_len, :]
+    elif current_seq_len < max_seq_len:
+        pad_tensor = torch.full((batch_size, max_seq_len - current_seq_len, d),
+                                padding_value, dtype=embeddings.dtype,
+                                device=embeddings.device)
+        return torch.cat([embeddings, pad_tensor], dim=1)
+    else:
+        return embeddings
+
+
+class MoleculeGPT(torch.nn.Module):
+    r"""The MoleculeGPT model from the `"MoleculeGPT: Instruction
+    Following Large Language Models for Molecular Property Prediction"
+    <https://ai4d3.github.io/papers/34.pdf>`_ paper.
+
+    Args:
+        llm (LLM): The LLM to use.
+        graph_encoder (torch.nn.Module): Encode 2D molecule graph.
+        smiles_encoder (torch.nn.Module): Encode 1D SMILES.
+        mlp_out_channels (int, optional): The size of each embedding
+            after qformer encoding. (default: :obj:`32`)
+        max_tokens (int, optional): Max output tokens of 1D/2D encoder.
+            (default: :obj:`20`)
+
+    .. warning::
+        This module has been tested with the following HuggingFace models
+
+        * :obj:`llm_to_use="lmsys/vicuna-7b-v1.5"`
+
+        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:`MoleculeGPT`, see
+        `examples/llm/molecule_gpt.py <https://github.com/pyg-team/
+        pytorch_geometric/blob/master/examples/llm/molecule_gpt.py>`_.
+    """
+    def __init__(
+        self,
+        llm: LLM,
+        graph_encoder: torch.nn.Module,
+        smiles_encoder: torch.nn.Module,
+        mlp_out_channels: int = 32,
+        max_tokens: Optional[int] = 20,
+    ) -> None:
+        super().__init__()
+        self.llm = llm
+        self.graph_encoder = graph_encoder.to(self.llm.device)
+        self.smiles_encoder = smiles_encoder.to(self.llm.device)
+
+        self.graph_qformer = QFormer(
+            input_dim=self.graph_encoder.nn[-1].out_features,
+            hidden_dim=mlp_out_channels,
+            output_dim=mlp_out_channels,
+            num_heads=4,
+            num_layers=2,
+        ).to(self.llm.device)
+
+        self.smiles_qformer = QFormer(
+            input_dim=self.smiles_encoder.model.pooler.dense.out_features,
+            hidden_dim=mlp_out_channels,
+            output_dim=mlp_out_channels,
+            num_heads=4,
+            num_layers=2,
+        ).to(self.llm.device)
+
+        self.max_tokens = max_tokens
+
+        self.word_embedding = self.llm.word_embedding
+        self.llm_generator = self.llm.llm
+
+        # LLMs
+        in_dim = 2 * mlp_out_channels * max_tokens
+        out_dim = self.llm.llm.model.embed_tokens.embedding_dim
+        self.projector = torch.nn.Sequential(
+            torch.nn.Linear(in_dim, in_dim),
+            torch.nn.Sigmoid(),
+            torch.nn.Linear(in_dim, out_dim),
+        ).to(self.llm.device)
+
+    def encode(
+        self,
+        x: Tensor,
+        edge_index: Tensor,
+        batch: Tensor,
+        edge_attr: Optional[Tensor],
+        smiles: List[str],
+    ) -> Tensor:
+        batch_size = len(smiles)
+        # 2D Graph Branch: [bs, node_len, d]
+        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)
+
+        x_graph = self.graph_encoder(x, edge_index, edge_attr=edge_attr)
+        x_graph = to_dense_batch(x_graph, batch)[0]
+        out_graph = self.graph_qformer(x_graph)
+        out_graph = pad_or_truncate(out_graph, max_seq_len=self.max_tokens,
+                                    padding_value=0)
+        out_graph = out_graph.view(batch_size, -1)
+
+        # 1D SMILES Branch: [bs, seq_len, d]
+        x_smiles = self.smiles_encoder.encode(smiles,
+                                              output_device=self.llm.device)
+        out_smiles = self.smiles_qformer(x_smiles)
+        out_smiles = pad_or_truncate(out_smiles, max_seq_len=self.max_tokens,
+                                     padding_value=0)
+        out_smiles = out_smiles.view(batch_size, -1)
+
+        # Merge into LLMs
+        x_cat = torch.cat([out_graph, out_smiles], dim=1)
+        return x_cat
+
+    def forward(
+        self,
+        x: Tensor,
+        edge_index: Tensor,
+        batch: Tensor,
+        edge_attr: Optional[Tensor],
+        smiles: List[str],
+        instructions: List[str],
+        label: List[str],
+        additional_text_context: Optional[List[str]] = None,
+    ):
+        x = self.encode(x, edge_index, batch, edge_attr, smiles)
+        x = self.projector(x)
+        xs = x.split(1, dim=0)
+
+        batch_unique = batch.unique()
+        batch_size = len(instructions)
+        if len(batch_unique) < batch_size:
+            xs = [
+                xs[i] if i in batch_unique else None for i in range(batch_size)
+            ]
+
+        (
+            inputs_embeds,
+            attention_mask,
+            label_input_ids,
+        ) = self.llm._get_embeds(instructions, 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,
+        x: Tensor,
+        edge_index: Tensor,
+        batch: Tensor,
+        edge_attr: Optional[Tensor],
+        smiles: List[str],
+        instructions: List[str],
+        additional_text_context: Optional[List[str]] = None,
+        max_out_tokens: Optional[int] = MAX_NEW_TOKENS,
+    ):
+        x = self.encode(x, edge_index, batch, edge_attr, smiles)
+        x = self.projector(x)
+        xs = x.split(1, dim=0)
+
+        # Handle questions without node features:
+        batch_unique = batch.unique()
+        batch_size = len(instructions)
+        if len(batch_unique) < batch_size:
+            xs = [
+                xs[i] if i in batch_unique else None for i in range(batch_size)
+            ]
+
+        inputs_embeds, attention_mask, _ = self.llm._get_embeds(
+            instructions, 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'  graph={self.graph_encoder.__class__.__name__},\n'
+                f'  smiles={self.smiles_encoder},\n'
+                f')')

torch_geometric/nn/models/node2vec.py

@@ -173,7 +173,6 @@ class Node2Vec(torch.nn.Module):
         test_z: Tensor,
         test_y: Tensor,
         solver: str = 'lbfgs',
-        multi_class: str = 'auto',
         *args,
         **kwargs,
     ) -> float:
@@ -182,7 +181,7 @@ class Node2Vec(torch.nn.Module):
         """
         from sklearn.linear_model import LogisticRegression
 
-        clf = LogisticRegression(solver=solver, multi_class=multi_class, *args,
+        clf = LogisticRegression(solver=solver, *args,
                                  **kwargs).fit(train_z.detach().cpu().numpy(),
                                                train_y.detach().cpu().numpy())
         return clf.score(test_z.detach().cpu().numpy(),

torch_geometric/nn/models/schnet.py

@@ -11,6 +11,7 @@ from torch import Tensor
 from torch.nn import Embedding, Linear, ModuleList, Sequential
 
 from torch_geometric.data import Dataset, download_url, extract_zip
+from torch_geometric.io import fs
 from torch_geometric.nn import MessagePassing, SumAggregation, radius_graph
 from torch_geometric.nn.resolver import aggregation_resolver as aggr_resolver
 from torch_geometric.typing import OptTensor
@@ -216,7 +217,7 @@ class SchNet(torch.nn.Module):
 
         with warnings.catch_warnings():
             warnings.simplefilter('ignore')
-            state = torch.load(path, map_location='cpu')
+            state = fs.torch_load(path, map_location='cpu')
 
         net = SchNet(
             hidden_channels=128,

torch_geometric/nn/models/signed_gcn.py

@@ -1,6 +1,5 @@
 from typing import Optional, Tuple
 
-import scipy.sparse
 import torch
 import torch.nn.functional as F
 from torch import Tensor
@@ -98,6 +97,7 @@ class SignedGCN(torch.nn.Module):
                 :obj:`max_val + 1` of :attr:`pos_edge_index` and
                 :attr:`neg_edge_index`. (default: :obj:`None`)
         """
+        import scipy.sparse as sp
         from sklearn.decomposition import TruncatedSVD
 
         edge_index = torch.cat([pos_edge_index, neg_edge_index], dim=1)
@@ -119,7 +119,7 @@ class SignedGCN(torch.nn.Module):
         # https://github.com/benedekrozemberczki/SGCN/blob/master/src/utils.py
         edge_index = edge_index.detach().numpy()
         val = val.detach().numpy()
-        A = scipy.sparse.coo_matrix((val, edge_index), shape=(N, N))
+        A = sp.coo_matrix((val, edge_index), shape=(N, N))
         svd = TruncatedSVD(n_components=self.in_channels, n_iter=128)
         svd.fit(A)
         x = svd.components_.T
@@ -256,7 +256,7 @@ class SignedGCN(torch.nn.Module):
             neg_p = self.discriminate(z, neg_edge_index)[:, :2].max(dim=1)[1]
         pred = (1 - torch.cat([pos_p, neg_p])).cpu()
         y = torch.cat(
-            [pred.new_ones((pos_p.size(0))),
+            [pred.new_ones(pos_p.size(0)),
              pred.new_zeros(neg_p.size(0))])
         pred, y = pred.numpy(), y.numpy()
 

torch_geometric/nn/module_dict.py

@@ -1,4 +1,4 @@
-from typing import Final, Iterable, Mapping, Optional, Set, Tuple, Union
+from typing import Final, Iterable, Mapping, Optional, Tuple, Union
 
 import torch
 from torch.nn import Module
@@ -11,7 +11,7 @@ Key = Union[str, Tuple[str, ...]]
 # internal representation and converts it back to `.` in the external
 # representation. It also allows passing tuples as keys.
 class ModuleDict(torch.nn.ModuleDict):
-    CLASS_ATTRS: Final[Set[str]] = set(dir(torch.nn.ModuleDict))
+    CLASS_ATTRS: Final[Tuple[str, ...]] = tuple(dir(torch.nn.ModuleDict))
 
     def __init__(
         self,

torch_geometric/nn/nlp/__init__.py

@@ -0,0 +1,9 @@
+from .sentence_transformer import SentenceTransformer
+from .vision_transformer import VisionTransformer
+from .llm import LLM
+
+__all__ = classes = [
+    'SentenceTransformer',
+    'VisionTransformer',
+    'LLM',
+]

torch_geometric/nn/nlp/llm.py

@@ -0,0 +1,322 @@
+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]'
+IGNORE_INDEX = -100
+MAX_TXT_LEN = 512
+MAX_NEW_TOKENS = 32
+PAD_TOKEN_ID = 0
+PADDING_SIDE = 'left'
+
+
+def get_llm_kwargs(required_memory: int, dtype=torch.dtype) -> Dict[str, Any]:
+    torch.cuda.empty_cache()
+
+    gpu_memory: List[int] = []
+    for i in range(torch.cuda.device_count()):
+        gpu_memory.append(torch.cuda.mem_get_info(i)[0] // 1024**3)
+        # Use the minimum number of GPUs to fit the LLM on.
+        if sum(gpu_memory) >= required_memory:
+            break
+
+    if sum(gpu_memory) < required_memory:
+        gpu_memory = []  # If not enough VRAM, use pure CPU.
+
+    kwargs = dict(revision='main')
+    if len(gpu_memory) > 0:
+        kwargs['max_memory'] = {
+            i: f'{memory}GiB'
+            for i, memory in enumerate(gpu_memory)
+        }
+        kwargs['low_cpu_mem_usage'] = True
+        kwargs['device_map'] = 'auto'
+        kwargs['torch_dtype'] = dtype
+
+    return kwargs
+
+
+class LLM(torch.nn.Module):
+    r"""A wrapper around a Large Language Model (LLM) from HuggingFace.
+
+    model_name (str): The HuggingFace model name, *e.g.*, :obj:`"llama2"` or
+        :obj:`"gemma"`.
+    num_params (int): An integer representing how many parameters the
+        HuggingFace model has, in billions. This is used to automatically
+        allocate the correct number of GPUs needed, given the available GPU
+        memory of your GPUs.
+    dtype (torch.dtype, optional): The data type to use for the LLM.
+        (default :obj: `torch.bfloat16`)
+    """
+    def __init__(
+        self,
+        model_name: str,
+        num_params: int,
+        dtype=torch.bfloat16,
+    ) -> None:
+        super().__init__()
+
+        self.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 '{model_name}' with configuration: {kwargs}")
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            use_fast=False,
+        )
+        self.tokenizer.pad_token_id = PAD_TOKEN_ID
+        self.tokenizer.padding_side = PADDING_SIDE
+        self.llm = AutoModelForCausalLM.from_pretrained(model_name, **kwargs)
+        self.word_embedding = self.llm.model.get_input_embeddings()
+
+        if 'max_memory' not in kwargs:  # Pure CPU:
+            warnings.warn("LLM is being used on CPU, which may be slow")
+            self.device = torch.device('cpu')
+            self.autocast_context = nullcontext()
+        else:
+            self.device = self.llm.device
+            self.autocast_context = torch.amp.autocast('cuda', dtype=dtype)
+
+    def _encode_inputs(
+        self,
+        question: List[str],
+        context: Optional[List[str]] = None,
+    ) -> tuple:
+        batch_size = len(question)
+        questions = self.tokenizer(question, add_special_tokens=False)
+        if context is not None:
+            context = self.tokenizer(context, add_special_tokens=False)
+
+        eos_user_tokens = self.tokenizer(EOS_USER, add_special_tokens=False)
+        bos_token = self.tokenizer(
+            BOS,
+            add_special_tokens=False,
+            return_tensors='pt',
+        ).input_ids[0].to(self.device)
+        bos_embeds = self.word_embedding(bos_token)
+        pad_token = torch.tensor(self.tokenizer.pad_token_id,
+                                 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],
+        answer: List[str],
+        context: Optional[List[str]] = None,
+        embedding: Optional[List[Tensor]] = None,
+    ) -> Tensor:
+        r"""The forward pass.
+
+        Args:
+            question (list[str]): The questions/prompts.
+            answer (list[str]): The answers/labels.
+            context (list[str], optional): Additional context to give to the
+                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:`embedding` should be used, not
+                both. (default: :obj:`None`)
+        """
+        inputs_embeds, attention_mask, label_input_ids = self._get_embeds(
+            question, context, embedding, answer)
+
+        with self.autocast_context:
+            outputs = self.llm(
+                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],
+        context: Optional[List[str]] = None,
+        embedding: Optional[List[Tensor]] = None,
+        max_tokens: Optional[int] = MAX_NEW_TOKENS,
+    ) -> List[str]:
+        r"""The inference pass.
+
+        Args:
+            question (list[str]): The questions/prompts.
+            answer (list[str]): The answers/labels.
+            context (list[str], optional): Additional context to give to the
+                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:`embedding` should be used, not
+                both. (default: :obj:`None`)
+            max_tokens (int, optional): How many tokens for the LLM to
+                generate. (default: :obj:`32`)
+        """
+        inputs_embeds, attention_mask, _ = self._get_embeds(
+            question, context, embedding)
+
+        bos_token = self.tokenizer(
+            BOS,
+            add_special_tokens=False,
+        ).input_ids[0]
+
+        with self.autocast_context:
+            outputs = self.llm.generate(
+                inputs_embeds=inputs_embeds,
+                bos_token_id=bos_token,
+                max_new_tokens=max_tokens,
+                attention_mask=attention_mask,
+                use_cache=True,
+            )
+
+        return self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.model_name})'