pyg-nightly 2.6.0.dev20240704__py3-none-any.whl → 2.8.0.dev20251207__py3-none-any.whl

torch_geometric/llm/models/sentence_transformer.py

@@ -0,0 +1,188 @@
+from enum import Enum
+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):
+    MEAN = 'mean'
+    LAST = 'last'
+    CLS = 'cls'
+    LAST_HIDDEN_STATE = 'last_hidden_state'
+
+
+class SentenceTransformer(torch.nn.Module):
+    r"""A wrapper around a Sentence-Transformer from HuggingFace.
+
+    Args:
+        model_name (str): The HuggingFace model name, *e.g.*, :obj:`"BERT"`.
+        pooling_strategy (str, optional): The pooling strategy to use
+            for generating node embeddings. (default: :obj:`"mean"`)
+    """
+    def __init__(
+        self,
+        model_name: str,
+        pooling_strategy: Union[PoolingStrategy, str] = 'mean',
+    ) -> None:
+        super().__init__()
+
+        self.model_name = model_name
+        self.pooling_strategy = PoolingStrategy(pooling_strategy)
+
+        from transformers import AutoModel, AutoTokenizer
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModel.from_pretrained(model_name)
+        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)
+
+        emb = out[0]  # First element contains all token embeddings.
+        if self.pooling_strategy == PoolingStrategy.MEAN:
+            emb = mean_pooling(emb, attention_mask)
+        elif self.pooling_strategy == PoolingStrategy.LAST:
+            emb = last_pooling(emb, attention_mask)
+        elif self.pooling_strategy == PoolingStrategy.LAST_HIDDEN_STATE:
+            emb = out.last_hidden_state
+        else:
+            assert self.pooling_strategy == PoolingStrategy.CLS
+            emb = emb[:, 0, :]
+
+        emb = F.normalize(emb, p=2, dim=1)
+        return emb
+
+    def get_input_ids(
+        self,
+        text: List[str],
+        batch_size: Optional[int] = None,
+        output_device: Optional[Union[torch.device, str]] = None,
+    ) -> Tensor:
+        is_empty = len(text) == 0
+        text = ['dummy'] if is_empty else text
+
+        batch_size = len(text) if batch_size is None else batch_size
+
+        input_ids: List[Tensor] = []
+        attention_masks: List[Tensor] = []
+        for start in range(0, len(text), batch_size):
+            token = self.tokenizer(
+                text[start:start + batch_size],
+                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))
+
+        def _out(x: List[Tensor]) -> Tensor:
+            out = torch.cat(x, dim=0) if len(x) > 1 else x[0]
+            out = out[:0] if is_empty else out
+            return out.to(output_device)
+
+        return _out(input_ids), _out(attention_masks)
+
+    @property
+    def device(self) -> torch.device:
+        return next(iter(self.model.parameters())).device
+
+    @torch.no_grad()
+    def encode(
+        self,
+        text: List[str],
+        batch_size: Optional[int] = None,
+        output_device: Optional[Union[torch.device, str]] = None,
+        verbose=False,
+    ) -> Tensor:
+        r"""Main function for users. Converts strings to embeddings.
+
+        Args:
+            text (List[str]): List of strings to embed.
+            batch_size (int, optional): How many strings to process.
+                Defaults to processing all at once, but this may lead to
+                OOM errors. (default: obj:`None`)
+            output_device (Union[torch.device, str], optional):
+                By default outputs cpu pytorch tensor, but can choose
+                to output to specific cuda devices. (default: obj:`None`)
+            verbose (bool, optional): Controls the verbosity of outputs.
+                (default: obj:`False`)
+        """
+        is_empty = len(text) == 0
+        text = ['dummy'] if is_empty else text
+
+        batch_size = len(text) if batch_size is None else batch_size
+
+        embs: List[Tensor] = []
+        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,
+            )
+            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
+        return out
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(model_name={self.model_name})'
+
+
+def mean_pooling(emb: Tensor, attention_mask: Tensor) -> Tensor:
+    mask = attention_mask.unsqueeze(-1).expand(emb.size()).to(emb.dtype)
+    return (emb * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
+
+
+def last_pooling(emb: Tensor, attention_mask: Tensor) -> Tensor:
+    # Check whether language model uses left padding,
+    # which is always used for decoder LLMs
+    left_padding = attention_mask[:, -1].sum() == attention_mask.size(0)
+    if left_padding:
+        return emb[:, -1]
+
+    seq_indices = attention_mask.sum(dim=1) - 1
+    return emb[torch.arange(emb.size(0), device=emb.device), seq_indices]

torch_geometric/llm/models/txt2kg.py

@@ -0,0 +1,353 @@
+import os
+import time
+from typing import List, Optional, Tuple
+
+import torch
+import torch.multiprocessing as mp
+
+CLIENT_INITD = False
+
+CLIENT = None
+GLOBAL_NIM_KEY = ""
+SYSTEM_PROMPT = "Please convert the above text into a list of knowledge triples with the form ('entity', 'relation', 'entity'). Separate each with a new line. Do not output anything else. Try to focus on key triples that form a connected graph."  # noqa
+
+
+class TXT2KG():
+    """A class to convert text data into a Knowledge Graph (KG) format.
+    Uses NVIDIA NIMs + Prompt engineering by default.
+    Default model `nvidia/llama-3.1-nemotron-70b-instruct`
+    is on par or better than GPT4o in benchmarks.
+    We need a high quality model to ensure high quality KG.
+    Otherwise we have garbage in garbage out for the rest of the
+    GNN+LLM RAG pipeline.
+
+    Use local_lm flag for local debugging/dev. You still need to be able to
+    inference a 14B param LLM, 'VAGOsolutions/SauerkrautLM-v2-14b-DPO'.
+    Smaller LLMs did not work at all in testing.
+    Note this 14B model requires a considerable amount of GPU memory.
+    See examples/llm/txt2kg_rag.py for an example.
+
+    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").
+        local_LM : bool, optional
+            A flag indicating whether a local Language Model (LM)
+            should be used. This uses HuggingFace and will be slower
+            than deploying your own private NIM endpoint. This flag
+            is mainly recommended for dev/debug.
+            (default: False).
+        chunk_size : int, optional
+            The size of the chunks in which the text data is processed
+            (default: 512).
+    """
+    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",
+        local_LM: bool = False,
+        chunk_size: int = 512,
+    ) -> None:
+        self.local_LM = local_LM
+        # Initialize the local LM flag and the NIM model info accordingly
+        if self.local_LM:
+            # If using a local LM, set the initd_LM flag to False
+            self.initd_LM = False
+        else:
+            # If not using a local LM, store the provided NIM model info
+            self.NVIDIA_API_KEY = NVIDIA_API_KEY
+            self.NIM_MODEL = NVIDIA_NIM_MODEL
+            self.ENDPOINT_URL = ENDPOINT_URL
+
+        # Set the chunk size for processing text data
+        self.chunk_size = chunk_size
+
+        # Initialize counters and storage for parsing results
+        self.doc_id_counter = 0
+        self.relevant_triples = {}
+        self.total_chars_parsed = 0
+        self.time_to_parse = 0.0
+
+    def save_kg(self, path: str) -> None:
+        """Saves the relevant triples in the knowledge graph (KG) to a file.
+
+        Args:
+            path (str): The file path where the KG will be saved.
+
+        Returns:
+            None
+        """
+        torch.save(self.relevant_triples, path)
+
+    def _chunk_to_triples_str_local(self, txt: str) -> str:
+        # call LLM on text
+        chunk_start_time = time.time()
+        if not self.initd_LM:
+            from torch_geometric.nn.nlp import LLM
+            LM_name = "VAGOsolutions/SauerkrautLM-v2-14b-DPO"
+            self.model = LLM(LM_name).eval()
+            self.initd_LM = True
+        out_str = self.model.inference(question=[txt + '\n' + SYSTEM_PROMPT],
+                                       max_tokens=self.chunk_size)[0]
+        # for debug
+        self.total_chars_parsed += len(txt)
+        self.time_to_parse += round(time.time() - chunk_start_time, 2)
+        self.avg_chars_parsed_per_sec = self.total_chars_parsed / self.time_to_parse  # noqa
+        return out_str
+
+    def add_doc_2_KG(
+        self,
+        txt: str,
+        QA_pair: Optional[Tuple[str, str]] = None,
+    ) -> None:
+        """Add a document to the Knowledge Graph (KG).
+
+        Args:
+            txt (str): The text to extract triples from.
+            QA_pair (Tuple[str, str]], optional):
+                A QA pair to associate with the extracted triples.
+                Useful for downstream evaluation.
+
+        Returns:
+        - None
+        """
+        if not self.local_LM:
+            # Ensure NVIDIA_API_KEY is set before proceeding
+            assert self.NVIDIA_API_KEY != '', \
+                "Please init TXT2KG w/ NVIDIA_API_KEY or set local_lm=True"
+        if QA_pair:
+            # QA_pairs should be unique keys, check if already exists in KG
+            if QA_pair in self.relevant_triples.keys():
+                print("Warning: QA_Pair was already added to the set")
+                print("Q=", QA_pair[0])
+                print("A=", QA_pair[1])
+                print("Previously parsed triples=",
+                      self.relevant_triples[QA_pair])
+                print("Skipping...")
+            key = QA_pair
+        else:
+            # If no QA_pair, use the current doc_id_counter as the key
+            key = self.doc_id_counter
+
+        # Handle empty text (context-less QA pairs)
+        if txt == "":
+            self.relevant_triples[key] = []
+        else:
+            # Chunk the text into smaller pieces for processing
+            chunks = _chunk_text(txt, chunk_size=self.chunk_size)
+
+            if self.local_LM:
+                # For debugging purposes...
+                # process chunks sequentially on the local LM
+                self.relevant_triples[key] = _llm_then_python_parse(
+                    chunks, _parse_n_check_triples,
+                    self._chunk_to_triples_str_local)
+            else:
+                # Process chunks in parallel using multiple processes
+                num_procs = min(len(chunks), _get_num_procs())
+                meta_chunk_size = int(len(chunks) / num_procs)
+                in_chunks_per_proc = {
+                    j:
+                    chunks[j *
+                           meta_chunk_size:min((j + 1) *
+                                               meta_chunk_size, len(chunks))]
+                    for j in range(num_procs)
+                }
+                for _retry_j in range(5):
+                    try:
+                        for _retry_i in range(200):
+                            try:
+                                # Spawn multiple processes
+                                # process chunks in parallel
+                                mp.spawn(
+                                    _multiproc_helper,
+                                    args=(in_chunks_per_proc,
+                                          _parse_n_check_triples,
+                                          _chunk_to_triples_str_cloud,
+                                          self.NVIDIA_API_KEY, self.NIM_MODEL,
+                                          self.ENDPOINT_URL), nprocs=num_procs)
+                                break
+                            except:  # noqa
+                                # keep retrying...
+                                # txt2kg is costly -> stoppage is costly
+                                pass
+
+                        # Collect the results from each process
+                        self.relevant_triples[key] = []
+                        for rank in range(num_procs):
+                            self.relevant_triples[key] += torch.load(
+                                "/tmp/outs_for_proc_" + str(rank))
+                            os.remove("/tmp/outs_for_proc_" + str(rank))
+                        break
+                    except:  # noqa
+                        pass
+        # Increment the doc_id_counter for the next document
+        self.doc_id_counter += 1
+
+
+known_reasoners = [
+    "llama-3.1-nemotron-ultra-253b-v1",
+    "kimi-k2-instruct",
+    "nemotron-super-49b-v1_5",
+    "gpt-oss",
+]
+
+
+def _chunk_to_triples_str_cloud(
+        txt: str, GLOBAL_NIM_KEY='',
+        NIM_MODEL="nvidia/llama-3.1-nemotron-ultra-253b-v1",
+        ENDPOINT_URL="https://integrate.api.nvidia.com/v1",
+        post_text=SYSTEM_PROMPT) -> str:
+    global CLIENT_INITD
+    if not CLIENT_INITD:
+        # We use NIMs since most PyG users may not be able to run a 70B+ model
+        try:
+            from openai import OpenAI
+        except ImportError:
+            quit(
+                "Failed to import `openai` package, please install it and rerun the script"  # noqa
+            )
+        global CLIENT
+        CLIENT = OpenAI(base_url=ENDPOINT_URL, api_key=GLOBAL_NIM_KEY)
+        CLIENT_INITD = True
+    txt_input = txt
+    if post_text != "":
+        txt_input += '\n' + post_text
+    messages = []
+    if any([model_name_str in NIM_MODEL
+            for model_name_str in known_reasoners]):
+        messages.append({"role": "system", "content": "detailed thinking on"})
+    messages.append({"role": "user", "content": txt_input})
+    completion = CLIENT.chat.completions.create(model=NIM_MODEL,
+                                                messages=messages,
+                                                temperature=0, top_p=1,
+                                                max_tokens=1024, stream=True)
+    out_str = ""
+    for chunk in completion:
+        if chunk.choices[0].delta.content is not None:
+            out_str += chunk.choices[0].delta.content
+    return out_str
+
+
+def _parse_n_check_triples(triples_str: str) -> List[Tuple[str, str, str]]:
+    # use pythonic checks for triples
+    processed = []
+    split_by_newline = triples_str.split("\n")
+    # sometimes LLM fails to obey the prompt
+    if len(split_by_newline) > 1:
+        split_triples = split_by_newline
+        llm_obeyed = True
+    else:
+        # handles form "(e, r, e) (e, r, e) ... (e, r, e)""
+        split_triples = triples_str[1:-1].split(") (")
+        llm_obeyed = False
+    for triple_str in split_triples:
+        try:
+            if llm_obeyed:
+                # remove parenthesis and single quotes for parsing
+                triple_str = triple_str.replace("(", "").replace(")",
+                                                                 "").replace(
+                                                                     "'", "")
+            split_trip = triple_str.split(',')
+            # remove blank space at beginning or end
+            split_trip = [(i[1:] if i[0] == " " else i) for i in split_trip]
+            split_trip = [(i[:-1].lower() if i[-1] == " " else i)
+                          for i in split_trip]
+            potential_trip = tuple(split_trip)
+        except:  # noqa
+            continue
+        if 'tuple' in str(type(potential_trip)) and len(
+                potential_trip
+        ) == 3 and "note:" not in potential_trip[0].lower():
+            # additional check for empty node/edge attrs
+            if potential_trip[0] != '' and potential_trip[
+                    1] != '' and potential_trip[2] != '':
+                processed.append(potential_trip)
+    return processed
+
+
+def _llm_then_python_parse(chunks, py_fn, llm_fn, **kwargs):
+    relevant_triples = []
+    for chunk in chunks:
+        relevant_triples += py_fn(llm_fn(chunk, **kwargs))
+    return relevant_triples
+
+
+def _multiproc_helper(rank, in_chunks_per_proc, py_fn, llm_fn, NIM_KEY,
+                      NIM_MODEL, ENDPOINT_URL):
+    out = _llm_then_python_parse(in_chunks_per_proc[rank], py_fn, llm_fn,
+                                 GLOBAL_NIM_KEY=NIM_KEY, NIM_MODEL=NIM_MODEL,
+                                 ENDPOINT_URL=ENDPOINT_URL)
+    torch.save(out, "/tmp/outs_for_proc_" + str(rank))
+
+
+def _get_num_procs():
+    if hasattr(os, "sched_getaffinity"):
+        try:
+            num_proc = len(os.sched_getaffinity(0)) / (2)
+        except Exception:
+            pass
+    if num_proc is None:
+        num_proc = os.cpu_count() / (2)
+    return int(num_proc)
+
+
+def _chunk_text(text: str, chunk_size: int = 512) -> list[str]:
+    """Function to chunk text into sentence-based segments.
+    Co-authored with Claude AI.
+    """
+    # If the input text is empty or None, return an empty list
+    if not text:
+        return []
+
+    # List of punctuation marks that typically end sentences
+    sentence_endings = '.!?'
+
+    # List to store the resulting chunks
+    chunks = []
+
+    # Continue processing the entire text
+    while text:
+        # If the remaining text is shorter than chunk_size, add it and break
+        if len(text) <= chunk_size:
+            chunks.append(text.strip())
+            break
+
+        # Start with the maximum possible chunk
+        chunk = text[:chunk_size]
+
+        # Try to find the last sentence ending within the chunk
+        best_split = chunk_size
+        for ending in sentence_endings:
+            # Find the last occurrence of the ending punctuation
+            last_ending = chunk.rfind(ending)
+            if last_ending != -1:
+                # Ensure we include the punctuation and any following space
+                best_split = min(
+                    best_split, last_ending + 1 +
+                    (1 if last_ending + 1 < len(chunk)
+                     and chunk[last_ending + 1].isspace() else 0))
+
+        # Adjust to ensure we don't break words
+        # If the next character is a letter, find the last space
+        if best_split < len(text) and text[best_split].isalpha():
+            # Find the last space before the current split point
+            space_split = text[:best_split].rfind(' ')
+            if space_split != -1:
+                best_split = space_split
+
+        # Append the chunk, ensuring it's stripped
+        chunks.append(text[:best_split].strip())
+
+        # Remove the processed part from the text
+        text = text[best_split:].lstrip()
+
+    return chunks

torch_geometric/llm/models/vision_transformer.py

@@ -0,0 +1,38 @@
+from typing import Optional, Union
+
+import torch
+from torch import Tensor
+
+
+class VisionTransformer(torch.nn.Module):
+    r"""A wrapper around a Vision-Transformer from HuggingFace.
+
+    Args:
+        model_name (str): The HuggingFace model name, *e.g.*, :obj:`"ViT"`.
+    """
+    def __init__(
+        self,
+        model_name: str,
+    ) -> None:
+        super().__init__()
+        self.model_name = model_name
+
+        from transformers import SwinConfig, SwinModel
+
+        self.config = SwinConfig.from_pretrained(model_name)
+        self.model = SwinModel(self.config)
+
+    @torch.no_grad()
+    def forward(
+        self,
+        images: Tensor,
+        output_device: Optional[Union[torch.device, str]] = None,
+    ) -> Tensor:
+        return self.model(images).last_hidden_state.to(output_device)
+
+    @property
+    def device(self) -> torch.device:
+        return next(iter(self.model.parameters())).device
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(model_name={self.model_name})'