pyg-nightly 2.7.0.dev20250905__py3-none-any.whl → 2.7.0.dev20250906__py3-none-any.whl

torch_geometric/{nn/nlp → llm/models}/llm.py

@@ -10,15 +10,17 @@ try:
 except ImportError:
     BatchEncoding = Dict
 
-BOS = '<s>[INST]'
-EOS_USER = '[/INST]'
-EOS = '[/s]'
 IGNORE_INDEX = -100
 MAX_TXT_LEN = 512
-MAX_NEW_TOKENS = 32
+MAX_NEW_TOKENS = 128
 PAD_TOKEN_ID = 0
 PADDING_SIDE = 'left'
 
+# legacy constants - used for Llama 2 style prompting
+BOS = '<s>[INST]'
+EOS_USER = '[/INST]'
+EOS = '[/s]'
+
 
 def get_llm_kwargs(required_memory: int, dtype=torch.dtype) -> Dict[str, Any]:
     torch.cuda.empty_cache()
@@ -50,49 +52,89 @@ class LLM(torch.nn.Module):
     r"""A wrapper around a Large Language Model (LLM) from HuggingFace.
 
     Args:
-        model_name (str): The HuggingFace model name, *e.g.*, :obj:`"llama2"`
-            or :obj:`"gemma"`.
-        num_params (int, optional): An integer representing how many parameters
+        model_name (str): The HuggingFace model name
+        num_params (float, optional): An integer representing how many params
             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. If not specified, the number of
-            parameters is determined using the `huggingface_hub` module.
+            automatically allocate the correct number of GPUs needed (using a
+            rough heuristic), given the available GPU memory of your GPUs.  If
+            not specified, the number of parameters is determined using the
+            `huggingface_hub` module.
+        n_gpus (int, optional): Number of GPUs to use. Designed for advanced
+            users to select how many GPU's they want to set this manually and
+            override the automatic set up mechanism.
         dtype (torch.dtype, optional): The data type to use for the LLM.
             (default :obj: `torch.bfloat16`)
+        sys_prompt (str, optional): A system prompt to use for the LLM.
+            (default: :obj: `None`)
     """
     def __init__(
         self,
         model_name: str,
-        num_params: Optional[int] = None,
+        num_params: Optional[float] = None,
+        n_gpus: Optional[int] = None,
         dtype: Optional[torch.dtype] = torch.bfloat16,
+        sys_prompt: Optional[str] = None,
     ) -> None:
         super().__init__()
 
         self.model_name = model_name
 
         from transformers import AutoModelForCausalLM, AutoTokenizer
-
-        if num_params is None:
-            from huggingface_hub import get_safetensors_metadata
-            safetensors_metadata = get_safetensors_metadata(model_name)
-            param_count = safetensors_metadata.parameter_count
-            num_params = list(param_count.values())[0] // 10**9
-
-        # 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)
+        if n_gpus is None:
+            if num_params is None:
+                from huggingface_hub import get_safetensors_metadata
+                safetensors_metadata = get_safetensors_metadata(model_name)
+                param_count = safetensors_metadata.parameter_count
+                num_params = float(list(param_count.values())[0] // 10**9)
+
+            # 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)
+        else:
+            gpu_memory: List[int] = []
+            for i in range(n_gpus):
+                gpu_memory.append(torch.cuda.mem_get_info(i)[0] // 1024**3)
+            kwargs = dict(revision='main')
+            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
 
         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
+        if self.tokenizer.chat_template and self.tokenizer.bos_token is None:
+            dummy_convo = [
+                {
+                    "role": "system",
+                    "content": "dummy"
+                },
+                {
+                    "role": "user",
+                    "content": "convo"
+                },
+            ]
+            text = self.tokenizer.apply_chat_template(
+                dummy_convo,
+                tokenize=True,
+            )
+            self.tokenizer.bos_token = self.tokenizer.decode(text[0])
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token_id = PAD_TOKEN_ID
+        if self.tokenizer.padding_side is None:
+            self.tokenizer.padding_side = PADDING_SIDE
         self.llm = AutoModelForCausalLM.from_pretrained(model_name, **kwargs)
         self.word_embedding = self.llm.model.get_input_embeddings()
-
+        if sys_prompt is not None:
+            self.sys_prompt = sys_prompt
+        else:
+            self.sys_prompt = ""
         if 'max_memory' not in kwargs:  # Pure CPU:
             warnings.warn("LLM is being used on CPU, which may be slow",
                           stacklevel=2)
@@ -100,8 +142,12 @@ class LLM(torch.nn.Module):
             self.autocast_context = nullcontext()
         else:
             self.device = self.llm.device
-            self.autocast_context = torch.amp.autocast('cuda', dtype=dtype)
+            if dtype == torch.float32:
+                self.autocast_context = nullcontext()
+            else:
+                self.autocast_context = torch.amp.autocast('cuda', dtype=dtype)
 
+    # legacy function - used for Llama 2 style prompting
     def _encode_inputs(
         self,
         question: List[str],
@@ -135,6 +181,7 @@ class LLM(torch.nn.Module):
         label_input_ids = label_input_ids + eos_tokens.input_ids
         return label_input_ids
 
+    # legacy function - used for Llama 2 style prompting
     def _input_ids(
         self,
         i: int,
@@ -149,6 +196,7 @@ class LLM(torch.nn.Module):
         input_ids += eos_user_tokens.input_ids
         return input_ids
 
+    # legacy function - used for Llama 2 style prompting
     def _inputs_embeds(
         self,
         i: int,
@@ -208,7 +256,8 @@ class LLM(torch.nn.Module):
                                            device=self.device)
         return inputs_embeds, attention_mask, label_input_ids
 
-    def _get_embeds(
+    # legacy function - used for Llama 2 style prompting
+    def _get_embeds_old(
         self,
         question: List[str],
         context: Optional[List[str]] = None,
@@ -255,6 +304,95 @@ class LLM(torch.nn.Module):
 
         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:
+        if not self.tokenizer.chat_template or not self.sys_prompt:
+            warnings.warn(
+                f"HuggingFace model {self.model_name} is not using a "
+                "chat template, using Llama 2 style prompting. Please "
+                "consider using a more recent model and initialize the "
+                "LLM with `sys_prompt`.", stacklevel=2)
+            return self._get_embeds_old(question, context, embedding, answer)
+        batch_label_input_ids = None
+        if answer is not None:
+            label = self.tokenizer(answer, add_special_tokens=False)
+            eos_tokens = self.tokenizer(self.tokenizer.eos_token,
+                                        add_special_tokens=False)
+            batch_label_input_ids = []
+
+        batch_inputs_embeds = []
+        batch_attention_mask = []
+        for i in range(len(question)):
+            ctx = f"{context[i]} - " if context else ""
+            messages = [
+                {
+                    "role": "system",
+                    "content": self.sys_prompt
+                },
+                {
+                    "role": "user",
+                    "content": f"{ctx} - {question[i]}"
+                },
+            ]
+            text = self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+                enable_thinking=True,
+            )
+            text = text[len(self.tokenizer.bos_token):]
+            input_ids = self.tokenizer(text,
+                                       add_special_tokens=False).input_ids
+            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
+
+            bos_token = self.tokenizer(
+                self.tokenizer.bos_token,
+                add_special_tokens=False,
+                return_tensors='pt',
+            ).input_ids[0].to(self.device)
+
+            bos_embeds = self.word_embedding(bos_token)
+
+            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)
+            inputs_embeds = torch.cat(to_cat, dim=0).to(self.device)
+
+            (
+                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,
+            )
+
+        pad_token = torch.tensor(self.tokenizer.pad_token_id,
+                                 device=self.device)
+        pad_embeds = self.word_embedding(pad_token).unsqueeze(0)
+
+        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],
@@ -311,17 +449,13 @@ class LLM(torch.nn.Module):
         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,
+                bos_token_id=self.tokenizer.bos_token_id,
                 max_new_tokens=max_tokens,
                 attention_mask=attention_mask,
+                pad_token_id=self.tokenizer.eos_token_id,
                 use_cache=True,
             )
 

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/{nn → llm}/models/molecule_gpt.py

@@ -3,8 +3,8 @@ 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.nn.nlp.llm import BOS, LLM, MAX_NEW_TOKENS
 from torch_geometric.utils import to_dense_batch
 
 

torch_geometric/{nn/nlp → llm/models}/sentence_transformer.py

@@ -4,6 +4,7 @@ from typing import List, Optional, Union
 import torch
 import torch.nn.functional as F
 from torch import Tensor
+from tqdm import tqdm
 
 
 class PoolingStrategy(Enum):
@@ -31,6 +32,19 @@ class SentenceTransformer(torch.nn.Module):
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
 
+        # Maximum sequence length from the model configuration (e.g. 8192 for
+        # models like ModernBERT)
+        self.max_seq_length = self.model.config.max_position_embeddings
+        """
+        Some models define a max sequence length in their configuration. Others
+        only in the tokenizer. This is a hacky heuristic to find the max
+        sequence length that works for the model.
+        """
+        probe_tokens = self.tokenizer("hacky heuristic", padding='max_length',
+                                      return_tensors='pt')
+        self.max_seq_length = min(self.max_seq_length,
+                                  probe_tokens.input_ids.shape[1])
+
     def forward(self, input_ids: Tensor, attention_mask: Tensor) -> Tensor:
         out = self.model(input_ids=input_ids, attention_mask=attention_mask)
 
@@ -67,6 +81,7 @@ class SentenceTransformer(torch.nn.Module):
                 padding=True,
                 truncation=True,
                 return_tensors='pt',
+                max_length=self.max_seq_length,
             )
             input_ids.append(token.input_ids.to(self.device))
             attention_masks.append(token.attention_mask.to(self.device))
@@ -88,6 +103,7 @@ class SentenceTransformer(torch.nn.Module):
         text: List[str],
         batch_size: Optional[int] = None,
         output_device: Optional[Union[torch.device, str]] = None,
+        verbose=False,
     ) -> Tensor:
         is_empty = len(text) == 0
         text = ['dummy'] if is_empty else text
@@ -95,20 +111,38 @@ class SentenceTransformer(torch.nn.Module):
         batch_size = len(text) if batch_size is None else batch_size
 
         embs: List[Tensor] = []
-        for start in range(0, len(text), batch_size):
+        loader = range(0, len(text), batch_size)
+        if verbose:
+            loader = tqdm(
+                loader, desc="Encoding " + str(len(text)) +
+                " strings w/ SentenceTransformer")
+        for start in loader:
             token = self.tokenizer(
                 text[start:start + batch_size],
                 padding=True,
                 truncation=True,
                 return_tensors='pt',
+                max_length=self.max_seq_length,
             )
-
-            emb = self(
-                input_ids=token.input_ids.to(self.device),
-                attention_mask=token.attention_mask.to(self.device),
-            ).to(output_device)
-
-            embs.append(emb)
+            try:
+                emb = self(
+                    input_ids=token.input_ids.to(self.device),
+                    attention_mask=token.attention_mask.to(self.device),
+                ).to(output_device)
+
+                embs.append(emb)
+            except:  # noqa
+                # fallback to using CPU for huge strings that cause OOMs
+                print("Sentence Transformer failed on cuda, trying w/ cpu...")
+                previous_device = self.device
+                self.model = self.model.to("cpu")
+                emb = self(
+                    input_ids=token.input_ids.to(self.device),
+                    attention_mask=token.attention_mask.to(self.device),
+                ).to(output_device)
+
+                embs.append(emb)
+                self.model = self.model.to(previous_device)
 
         out = torch.cat(embs, dim=0) if len(embs) > 1 else embs[0]
         out = out[:0] if is_empty else out