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

torch_geometric/nn/models/g_retriever.py

@@ -0,0 +1,230 @@
+from typing import List, Optional
+
+import torch
+from torch import Tensor
+
+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`)
+        mlp_out_tokens (int, optional): Number of LLM prefix tokens to
+            reserve for GNN output. (default: :obj:`1`)
+
+    .. 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,
+        mlp_out_channels: int = 4096,
+        mlp_out_tokens: int = 1,
+    ) -> 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 * mlp_out_tokens),
+            torch.nn.Unflatten(-1, (mlp_out_tokens, 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(1, dim=0)
+
+        # Handle case where theres more than one embedding for each sample
+        xs = [x.squeeze(0) for x in xs]
+
+        # Handle questions without node features:
+        batch_unique = batch.unique()
+        batch_size = len(question)
+        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(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(1, dim=0)
+
+        # Handle case where theres more than one embedding for each sample
+        xs = [x.squeeze(0) for x in xs]
+
+        # Handle questions without node features:
+        batch_unique = batch.unique()
+        batch_size = len(question)
+        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(
+            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/models/git_mol.py

@@ -0,0 +1,336 @@
+from typing import List, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn import BatchNorm1d, LayerNorm, Linear, ReLU, Sequential
+
+from torch_geometric.nn import GINEConv
+from torch_geometric.nn.nlp import SentenceTransformer, VisionTransformer
+from torch_geometric.utils import add_self_loops, to_dense_batch
+
+
+class GraphEncoder(torch.nn.Module):
+    def __init__(
+        self,
+        num_layers: int,
+        in_channels: int,
+        dropout: float = 0.,
+        num_atom_type: int = 120,
+        num_chirality_tag: int = 3,
+        num_bond_type: int = 6,
+        num_bond_direction: int = 3,
+    ) -> None:
+        super().__init__()
+
+        self.num_layers = num_layers
+        self.dropout = dropout
+
+        self.x_embed1 = torch.nn.Embedding(num_atom_type, in_channels)
+        self.x_embed2 = torch.nn.Embedding(num_chirality_tag, in_channels)
+        self.edge_embed1 = torch.nn.Embedding(num_bond_type, in_channels)
+        self.edge_embed2 = torch.nn.Embedding(num_bond_direction, in_channels)
+
+        self.gnns = torch.nn.ModuleList()
+        self.batch_norms = torch.nn.ModuleList()
+        for _ in range(num_layers):
+            self.gnns.append(
+                GINEConv(
+                    nn=Sequential(
+                        Linear(in_channels, in_channels * 2),
+                        ReLU(),
+                        Linear(in_channels * 2, in_channels),
+                    ),
+                    train_eps=True,
+                    edge_dim=in_channels,
+                ))
+            self.batch_norms.append(BatchNorm1d(in_channels))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        torch.nn.init.xavier_uniform_(self.x_embed1.weight.data)
+        torch.nn.init.xavier_uniform_(self.x_embed2.weight.data)
+        torch.nn.init.xavier_uniform_(self.edge_embed1.weight.data)
+        torch.nn.init.xavier_uniform_(self.edge_embed2.weight.data)
+
+    def forward(
+        self,
+        x: Tensor,
+        edge_index: Tensor,
+        batch: Tensor,
+        edge_attr: Tensor,
+    ) -> Tensor:
+        x = self.x_embed1(x[:, 0].long()) + self.x_embed2(x[:, 1].long())
+        edge_index, edge_attr = add_self_loops(
+            edge_index,
+            edge_attr,
+            fill_value=0,
+            num_nodes=x.size(0),
+        )
+        edge_attr = self.edge_embed1(edge_attr[:, 0]) + self.edge_embed2(
+            edge_attr[:, 1])
+        for i, (gnn, bn) in enumerate(zip(self.gnns, self.batch_norms)):
+            x = gnn(x, edge_index, edge_attr)
+            x = bn(x)
+            if i < self.num_layers - 1:
+                x = F.relu(x)
+            x = F.dropout(x, self.dropout, training=self.training)
+
+        x, mask = to_dense_batch(x, batch)
+        return x, mask
+
+
+class GITFormer(torch.nn.Module):
+    def __init__(
+        self,
+        num_query_token: int,
+        vision_graph_width: int,
+        cross_attention_freq: int = 2,
+    ):
+        super().__init__()
+        from transformers import AutoConfig, AutoModel
+
+        config = AutoConfig.from_pretrained("allenai/scibert_scivocab_uncased")
+        config.encoder_width = vision_graph_width
+        # insert cross-attention layer every other block
+        config.add_cross_attention = True
+        config.is_decoder = True
+        config.cross_attention_freq = cross_attention_freq
+        config.query_length = num_query_token
+        self.Qformer = AutoModel.from_pretrained(
+            "allenai/scibert_scivocab_uncased", config=config)
+        self.query_tokens = torch.nn.Parameter(
+            torch.zeros(1, num_query_token, config.hidden_size))
+        self.query_tokens.data.normal_(mean=0.0, std=config.initializer_range)
+
+
+class GITMol(torch.nn.Module):
+    r"""The GITMol model from the `"GIT-Mol: A Multi-modal Large Language
+    Model for Molecular Science with Graph, Image, and Text"
+    <https://arxiv.org/pdf/2308.06911>`_ paper.
+
+    .. note::
+        For an example of using :class:`GITMol`, see
+        `examples/llm/git_mol.py <https://github.com/pyg-team/
+        pytorch_geometric/blob/master/examples/llm/git_mol.py>`_.
+    """
+    def __init__(self) -> None:
+        super().__init__()
+        # graph
+        self.graph_encoder = GraphEncoder(num_layers=2, in_channels=16)
+        self.graph_proj = Linear(16, 768)
+        self.ln_graph = LayerNorm(768)
+        # text
+        self.text_encoder = SentenceTransformer(
+            model_name='allenai/scibert_scivocab_uncased',
+            pooling_strategy='last_hidden_state',
+        )
+        self.text_proj = Linear(768, 768)
+        self.ln_text = LayerNorm(768)
+        # vision
+        self.vision_encoder = VisionTransformer(
+            model_name='microsoft/swin-base-patch4-window7-224', )
+        self.vision_proj = Linear(1024, 768)
+        self.ln_vision = LayerNorm(768)
+        # cross-attention
+        self.gitformer = GITFormer(384, 768)
+
+        self.xtm_head = torch.nn.ModuleDict({
+            'image':
+            Linear(self.gitformer.Qformer.config.hidden_size, 2),
+            'graph':
+            Linear(self.gitformer.Qformer.config.hidden_size, 2),
+            'cs_text':
+            Linear(self.gitformer.Qformer.config.hidden_size, 2),
+        })
+
+        self.xtc_proj = torch.nn.ModuleDict({
+            'image':
+            Linear(self.gitformer.Qformer.config.hidden_size, 768),
+            'graph':
+            Linear(self.gitformer.Qformer.config.hidden_size, 768),
+            'cs_text':
+            Linear(self.gitformer.Qformer.config.hidden_size, 768),
+        })
+        self.temp = torch.nn.Parameter(0.07 * torch.ones([]))
+        self.model_freeze()
+
+    def model_freeze(self) -> None:
+        for param in self.graph_encoder.parameters():
+            param.requires_grad = False
+
+        for param in self.vision_encoder.parameters():
+            param.requires_grad = False
+
+    def forward(
+        self,
+        x: Tensor,
+        edge_index: Tensor,
+        batch: Tensor,
+        edge_attr: Optional[Tensor],
+        smiles: List[str],
+        images: Tensor,
+        captions: List[str],
+    ) -> Tensor:
+        batch_size = len(smiles)
+
+        x_vision = self.vision_encoder(images)
+        x_vision = self.vision_proj(x_vision)
+        x_vision = self.ln_vision(x_vision)  # [bs, patch_len, d]
+        vision_atts = torch.ones(x_vision.size()[:-1],
+                                 dtype=torch.long).to(x_vision.device)
+        vision_targets = torch.arange(batch_size).to(x_vision.device)
+
+        x_graph, graph_atts = self.graph_encoder(x, edge_index, batch,
+                                                 edge_attr)
+        x_graph = self.graph_proj(x_graph)
+        x_graph = self.ln_graph(x_graph)  # [bs, node_len, d]
+        graph_targets = torch.arange(batch_size).to(x_graph.device)
+
+        x_smiles = self.text_encoder.encode(smiles)  # [bs, seq_len, d]
+        smiles_atts = torch.ones(x_smiles.size()[:-1],
+                                 dtype=torch.long).to(x_smiles.device)
+        smiles_targets = torch.arange(batch_size).to(x_smiles.device)
+
+        caption_input_ids, caption_attention_masks = self.text_encoder.get_input_ids(  # noqa: E501
+            captions)
+
+        text_output = self.gitformer.Qformer(
+            caption_input_ids,
+            attention_mask=caption_attention_masks,
+            return_dict=True,
+        )
+        text_feat = F.normalize(
+            self.text_proj(text_output.last_hidden_state[:, 0, :]), dim=-1)
+
+        loss = 0
+        for x_embed, x_atts, x_targets, modal in zip(
+            [x_graph, x_smiles, x_vision],
+            [graph_atts, smiles_atts, vision_atts],
+            [graph_targets, smiles_targets, vision_targets],
+            ['graph', 'cs_text', 'image'],
+        ):
+            loss += self._calc_xtc_loss(x_embed, x_atts, x_targets, text_feat,
+                                        modal)
+            loss += self._calc_xtm_loss(x_embed, caption_input_ids,
+                                        caption_attention_masks, modal)
+
+        return loss / 6
+
+    def _calc_xtm_loss(
+        self,
+        x_embeds: Tensor,
+        input_ids: Tensor,
+        attention_mask: Tensor,
+        modal: str,
+    ) -> Tensor:
+        # Initializing lists to hold the original and negative samples
+        x_embeds_list = []
+        text_input_ids_list = []
+        text_attention_mask_list = []
+
+        batch_size = x_embeds.size(0)
+        for i in range(batch_size):
+            # Original samples
+            x_embeds_list.append(x_embeds[i])
+            text_input_ids_list.append(input_ids[i, :])
+            text_attention_mask_list.append(attention_mask[i, :])
+
+            if batch_size > 1:
+                # Negative samples (neg_text_input_ids corresponds to x_embeds)
+                neg_text_input_ids = input_ids[i - 1 if i == batch_size -
+                                               1 else i + 1, :]
+                neg_text_attention_mask = attention_mask[i -
+                                                         1 if i == batch_size -
+                                                         1 else i + 1, :]
+                text_input_ids_list.append(neg_text_input_ids)
+                text_attention_mask_list.append(neg_text_attention_mask)
+                x_embeds_list.append(x_embeds[i, :])
+
+                # Negative samples (text_input_ids corresponds to neg_x_embeds)
+                neg_x_embeds = x_embeds[i - 1 if i == batch_size - 1 else i +
+                                        1, :]
+                x_embeds_list.append(neg_x_embeds)
+                text_input_ids_list.append(input_ids[i, :])
+                text_attention_mask_list.append(attention_mask[i, :])
+
+        # Stack all samples into two large tensors
+        x_embeds_all = torch.stack(x_embeds_list, dim=1) \
+            .reshape(-1, x_embeds.size(1), x_embeds.size(2))
+        text_input_ids_all = torch.stack(text_input_ids_list, dim=1) \
+            .reshape(-1, input_ids.size(1))
+        # Create image attention masks for the concatenated tensor
+        image_attns_all = torch.ones(x_embeds_all.size()[:-1],
+                                     dtype=torch.long).to(x_embeds_all.device)
+        query_tokens_xtm = self.gitformer.query_tokens.expand(
+            text_input_ids_all.shape[0], -1, -1)
+        query_attns_xtm = torch.ones(query_tokens_xtm.size()[:-1],
+                                     dtype=torch.long).to(x_embeds_all.device)
+
+        output_xtm = self.gitformer.Qformer(
+            inputs_embeds=query_tokens_xtm,
+            attention_mask=query_attns_xtm,
+            encoder_hidden_states=x_embeds_all,
+            encoder_attention_mask=image_attns_all,
+            return_dict=True,
+        ).last_hidden_state
+
+        xtm_embeddings = output_xtm[:, :query_tokens_xtm.size(1), :]
+
+        xtm_logit = self.xtm_head[modal](xtm_embeddings).mean(dim=1)
+        # Create labels: 1 for the original samples, 0 for the negative samples
+        if batch_size > 1:
+            labels = torch.cat(
+                [torch.ones(batch_size),
+                 torch.zeros(batch_size * 2)], dim=0)
+        else:
+            labels = torch.ones(batch_size)
+        labels = labels.long().to(xtm_logit.device)
+
+        # Calculate cross entropy loss
+        return F.cross_entropy(xtm_logit, labels)
+
+    def _calc_xtc_loss(
+        self,
+        x_embeds: Tensor,
+        x_atts: Tensor,
+        x_targets: Tensor,
+        text_feat: Tensor,
+        modal: str,
+    ) -> Tensor:
+        query_tokens = self.gitformer.query_tokens.expand(
+            x_embeds.shape[0], -1, -1)
+
+        query_output = self.gitformer.Qformer(
+            inputs_embeds=query_tokens,
+            encoder_hidden_states=x_embeds,
+            encoder_attention_mask=x_atts,
+            return_dict=True,
+        ).last_hidden_state
+
+        x_feats = F.normalize(self.xtc_proj[modal](query_output), dim=-1)
+
+        sim_q2t = torch.matmul(
+            x_feats.unsqueeze(1),
+            text_feat.unsqueeze(-1),
+        ).squeeze(-1)
+
+        # modal-text similarity: aggregate across all query tokens
+        sim_x2t, _ = sim_q2t.max(-1)
+        sim_x2t = sim_x2t / self.temp
+
+        # text-query similarity
+        sim_t2q = torch.matmul(
+            text_feat.unsqueeze(1).unsqueeze(1),
+            x_feats.permute(0, 2, 1),
+        ).squeeze(-2)
+
+        # text-modal similarity: aggregate across all query tokens
+        sim_t2x, _ = sim_t2q.max(-1)
+        sim_t2x = sim_t2x / self.temp
+
+        loss_itc = (
+            F.cross_entropy(sim_x2t, x_targets, label_smoothing=0.1) +
+            F.cross_entropy(sim_t2x, x_targets, label_smoothing=0.1)) / 2
+
+        return loss_itc