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

torch_geometric/llm/models/llm_judge.py

@@ -0,0 +1,158 @@
+from math import isnan
+from typing import Optional
+
+from torch_geometric.llm.models.txt2kg import \
+    _chunk_to_triples_str_cloud as call_NIM
+
+# Credit for original "Marlin Accuracy" system goes to:
+# Gilberto Titericz (NVIDIA)
+# This work is an adaptation of his for PyG
+SYSTEM_PROMPT_1 = (
+    "Instruction: You are a world class state of the art " +
+    "assistant for rating " +
+    "a User Answer given a Question. The Question is completely" +
+    " answered by the Reference Answer.\n" +
+    "Say 4, if User Answer is full contained and equivalent to" +
+    " Reference Answer" +
+    "in all terms, topics, numbers, metrics, dates and units.\n" +
+    "Say 2, if User Answer is partially contained and almost " +
+    "equivalent to Reference Answer" +
+    "in all terms, topics, numbers, metrics, dates and units.\n" +
+    "Say 0, if User Answer is not contained in Reference Answer" +
+    " or not accurate in all terms, topics," +
+    "numbers, metrics, dates and units or the User Answer do not" +
+    " answer the question.\n" +
+    "Do not explain or justify your rating. Your rating must be " +
+    "only 4, 2 or 0 according to the instructions above.\n" +
+    "### Question: \"{question}\"\n" + "### User Answer: \"{model_pred}\"\n" +
+    "### Reference Answer: \"{correct_answer}\"\n" + "The rating is:\n")
+
+SYSTEM_PROMPT_2 = (
+    "I will rate the User Answer in comparison to the Reference " +
+    "Answer for a given Question.\n" +
+    "A rating of 4 indicates that the User Answer is entirely " +
+    "consistent with the Reference Answer, covering all aspects," +
+    " topics, numbers, metrics, dates, and units.\n" +
+    "A rating of 2 signifies that the User Answer is mostly " +
+    "aligned with the Reference Answer, with minor discrepancies" +
+    " in some areas.\n" +
+    "A rating of 0 means that the User Answer is either " +
+    "inaccurate, incomplete, or unrelated to the Reference " +
+    "Answer, or it fails to address the Question.\n" +
+    "I will provide the rating without any explanation or " +
+    "justification, adhering to the following scale: " +
+    "0 (no match), 2 (partial match), 4 (exact match).\n" +
+    "Do not explain or justify my rating. My rating must" +
+    " be only 4, 2 or 0 only.\n\n" + "Question: \"{question}\"\n\n" +
+    "Reference Answer: \"{model_pred}\"\n\n" +
+    "User Answer: \"{correct_answer}\"\n\n" + "Rating: ")
+
+
+# TODO: add support for Local LM
+# TODO: add multiproc support like txt2kg
+class LLMJudge():
+    """Uses NIMs to score a triple of (question, model_pred, correct_answer)
+    This whole class is an adaptation of Gilberto's work for PyG.
+
+    Args:
+        NVIDIA_NIM_MODEL : (str, optional)
+            The name of the NVIDIA NIM model to use.
+            (default: "nvidia/llama-3.1-nemotron-70b-instruct").
+        NVIDIA_API_KEY : (str, optional)
+            The API key for accessing NVIDIA's NIM models.
+            (default: "").
+        ENDPOINT_URL : (str, optional)
+            The URL hosting your model, in case you are not using
+            the public NIM.
+            (default: "https://integrate.api.nvidia.com/v1").
+    """
+    def __init__(
+        self,
+        NVIDIA_NIM_MODEL: Optional[
+            str] = "nvidia/llama-3.1-nemotron-70b-instruct",
+        NVIDIA_API_KEY: Optional[str] = "",
+        ENDPOINT_URL: Optional[str] = "https://integrate.api.nvidia.com/v1",
+    ) -> None:
+        self.NVIDIA_API_KEY = NVIDIA_API_KEY
+        self.NIM_MODEL = NVIDIA_NIM_MODEL
+        self.ENDPOINT_URL = ENDPOINT_URL
+
+    def _process_score(self, response: str) -> float:
+        """Uses 3 and 1 even though prompt says only 0, 2, 4.
+        This is because LLMs don't always follow instructions.
+        Credit to Gilberto.
+        """
+        for i in [4, 3, 2, 1, 0]:
+            if str(i) in response:
+                return i / 4
+        return float("nan")
+
+    def _average_scores(self, score0: float, score1: float):
+        """Take the average of score0 and score1.
+        Sometimes the LLM fail to respond or have no score in the response.
+        In those cases the failed score is discarded.
+        Credit to Gilberto.
+
+        Args:
+         score0 (float): judge accuracy score.
+         score1 (float): judge accuracy score by permuting agent answer and
+         ground truth.
+
+        Returns:
+            (float) average of score0 and score1 of both contains scores,
+            otherwise pick the max.
+        """
+        score = float("nan")
+        if score0 >= 0 and score1 >= 0:
+            score = (score0 + score1) / 2
+        else:
+            score = max(score0, score1)
+        return score
+
+    def score(
+        self,
+        question: str,
+        model_pred: str,
+        correct_answer: str,
+    ) -> float:
+        """Args:
+            question (str): The original question asked to the model.
+            model_pred (str): The prediction made by the model.
+            correct_answer (str): The actual correct answer to the question.
+
+        Returns:
+            score (float): score of 0-1, may be nan due to LLM judge failure.
+                Evals should skip nan's when aggregating score.
+        """
+        prompt1 = SYSTEM_PROMPT_1.format(question=question,
+                                         model_pred=model_pred,
+                                         correct_answer=correct_answer)
+        prompt2 = SYSTEM_PROMPT_2.format(question=question,
+                                         model_pred=model_pred,
+                                         correct_answer=correct_answer)
+        score1 = float("nan")
+        score2 = float("nan")
+        for _retry in range(200):
+            try:
+                score1 = self._process_score(
+                    call_NIM(prompt1, self.NVIDIA_API_KEY, self.NIM_MODEL,
+                             self.ENDPOINT_URL, post_text=""))
+                if not isnan(score1):
+                    break
+            except ImportError:
+                raise
+            except:  # noqa
+                pass
+        for _retry in range(20):
+            try:
+                score2 = self._process_score(
+                    call_NIM(prompt2, self.NVIDIA_API_KEY, self.NIM_MODEL,
+                             self.ENDPOINT_URL, post_text=""))
+                if not isnan(score2):
+                    break
+            except ImportError:
+                raise
+            except:  # noqa
+                pass
+
+        return self._average_scores(score1, score2)

torch_geometric/llm/models/molecule_gpt.py

@@ -0,0 +1,222 @@
+from typing import List, Optional
+
+import torch
+from torch import Tensor
+
+from torch_geometric.llm.models.llm import BOS, LLM, MAX_NEW_TOKENS
+from torch_geometric.nn.attention import QFormer
+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')')