pyg-nightly 2.7.0.dev20250124__py3-none-any.whl → 2.7.0.dev20250125__py3-none-any.whl

{pyg_nightly-2.7.0.dev20250124.dist-info → pyg_nightly-2.7.0.dev20250125.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pyg-nightly
-Version: 2.7.0.dev20250124
+Version: 2.7.0.dev20250125
 Summary: Graph Neural Network Library for PyTorch
 Keywords: deep-learning,pytorch,geometric-deep-learning,graph-neural-networks,graph-convolutional-networks
 Author-email: Matthias Fey <matthias@pyg.org>

{pyg_nightly-2.7.0.dev20250124.dist-info → pyg_nightly-2.7.0.dev20250125.dist-info}/RECORD

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=woO2qQnJ5h-sALP2KdvQ69AVAzlveBh2Qs87DiEu7A4,1904
+torch_geometric/__init__.py,sha256=ruVTuh8ilC7szVlSM3RmBoqMIi7XgJ-VsGT6lSLsJqw,1904
 torch_geometric/_compile.py,sha256=f-WQeH4VLi5Hn9lrgztFUCSrN_FImjhQa6BxFzcYC38,1338
 torch_geometric/_onnx.py,sha256=V9ffrIKSqhDw6xUZ12lkuSfNs48cQp2EeJ6Z19GfnVw,349
 torch_geometric/backend.py,sha256=lVaf7aLoVaB3M-UcByUJ1G4T4FOK6LXAg0CF4W3E8jo,1575
@@ -54,7 +54,7 @@ torch_geometric/data/temporal.py,sha256=WOJ6gFrTLikaLhUvotyUF5ql14FkE5Ox3hNkdSp6
 torch_geometric/data/view.py,sha256=XjkVSc-UWZFCT4DlXLShZtO8duhFQkS9gq88zZXANsk,1089
 torch_geometric/data/lightning/__init__.py,sha256=w3En1tJfy3kSqe1MycpOyZpHFO3fxBCgNCUOznPA3YU,178
 torch_geometric/data/lightning/datamodule.py,sha256=Bn9iaIfE4NWDDWWMqCvBeZ4bIW1Silx_Ol5CPJCliaQ,29242
-torch_geometric/datasets/__init__.py,sha256=eqVmuffZnc-O7KBdXO98SNVwSGehT5uy2LAC86MxGO4,6107
+torch_geometric/datasets/__init__.py,sha256=d9nuTCytBvg60lm_WYRAQwjoZxR1H_7JsW8een1k1No,6186
 torch_geometric/datasets/actor.py,sha256=oUxgJIX8bi5hJr1etWNYIFyVQNDDXi1nyVpHGGMEAGQ,4304
 torch_geometric/datasets/airfrans.py,sha256=212gYsk7PvF-qcmvM2YXaOBhFrS79evAGg_sPHXih4w,5439
 torch_geometric/datasets/airports.py,sha256=b3gkv3gY2JkUpmGiz36Z-g7EcnSfU8lBG1YsCOWdJ6k,3758
@@ -104,6 +104,7 @@ torch_geometric/datasets/icews.py,sha256=Vdlk-PD10AU68xq8X5IOgrK0wgIBFq8A0D6_Wtr
 torch_geometric/datasets/igmc_dataset.py,sha256=pMiOoXjvqhfsDDNw51WT_IVi6wGJ0cUNwTdpEprPh3E,4611
 torch_geometric/datasets/imdb.py,sha256=QVJbtPPkcLznyvzuxDCxmqO5xXocVG59KhrjXi1qXg0,4232
 torch_geometric/datasets/infection_dataset.py,sha256=jIYqX0vkCE-3fNjaijzCSmY1RVMFiX3gnmLwkqDXRkI,7293
+torch_geometric/datasets/instruct_mol_dataset.py,sha256=EK_3lRflFYS6KHfPM1AcYtB7FRs2We3GgTu39H9vVKI,4990
 torch_geometric/datasets/jodie.py,sha256=8CW43ZepM26dk2HMGvXDDF-4BorBeegqegViWyeYOks,3643
 torch_geometric/datasets/karate.py,sha256=khCcCUEaw7FuYBKwEsOoogpTShKYnx5nXrRtCOAoEAU,3462
 torch_geometric/datasets/last_fm.py,sha256=jKM3gw7T5x4AlUtmA0TXB2iWpNMi-S-ME2bP37kzE3Q,4581
@@ -629,6 +630,6 @@ torch_geometric/utils/undirected.py,sha256=H_nfpI0_WluOG6VfjPyldvcjL4w5USAKWu2x5
 torch_geometric/visualization/__init__.py,sha256=PyR_4K5SafsJrBr6qWrkjKr6GBL1b7FtZybyXCDEVwY,154
 torch_geometric/visualization/graph.py,sha256=ZuLPL92yGRi7lxlqsUPwL_EVVXF7P2kMcveTtW79vpA,4784
 torch_geometric/visualization/influence.py,sha256=CWMvuNA_Nf1sfbJmQgn58yS4OFpeKXeZPe7kEuvkUBw,477
-pyg_nightly-2.7.0.dev20250124.dist-info/WHEEL,sha256=CpUCUxeHQbRN5UGRQHYRJorO5Af-Qy_fHMctcQ8DSGI,82
-pyg_nightly-2.7.0.dev20250124.dist-info/METADATA,sha256=Rzd4j6F9Pjy0ArOM92QFdu_jUeoFnqyDAbIlh136kzk,62977
-pyg_nightly-2.7.0.dev20250124.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250125.dist-info/WHEEL,sha256=CpUCUxeHQbRN5UGRQHYRJorO5Af-Qy_fHMctcQ8DSGI,82
+pyg_nightly-2.7.0.dev20250125.dist-info/METADATA,sha256=aLdh0yAjnLp2dOVP4uP8PHdBesMh-ed4zj7HUkFQuJs,62977
+pyg_nightly-2.7.0.dev20250125.dist-info/RECORD,,

torch_geometric/__init__.py

@@ -30,7 +30,7 @@ from .lazy_loader import LazyLoader
 contrib = LazyLoader('contrib', globals(), 'torch_geometric.contrib')
 graphgym = LazyLoader('graphgym', globals(), 'torch_geometric.graphgym')
 
-__version__ = '2.7.0.dev20250124'
+__version__ = '2.7.0.dev20250125'
 
 __all__ = [
     'Index',

torch_geometric/datasets/__init__.py

@@ -79,6 +79,7 @@ from .neurograph import NeuroGraphDataset
 from .web_qsp_dataset import WebQSPDataset, CWQDataset
 from .git_mol_dataset import GitMolDataset
 from .molecule_gpt_dataset import MoleculeGPTDataset
+from .instruct_mol_dataset import InstructMolDataset
 from .tag_dataset import TAGDataset
 
 from .dbp15k import DBP15K
@@ -196,6 +197,7 @@ homo_datasets = [
     'CWQDataset',
     'GitMolDataset',
     'MoleculeGPTDataset',
+    'InstructMolDataset',
     'TAGDataset',
 ]
 

torch_geometric/datasets/instruct_mol_dataset.py

@@ -0,0 +1,134 @@
+import json
+import sys
+from typing import Callable, List, Optional
+
+import torch
+from tqdm import tqdm
+
+from torch_geometric.data import Data, InMemoryDataset
+from torch_geometric.io import fs
+from torch_geometric.utils import one_hot
+
+
+class InstructMolDataset(InMemoryDataset):
+    r"""The dataset from the `"InstructMol: Multi-Modal Integration for
+    Building a Versatile and Reliable Molecular Assistant in Drug Discovery"
+    <https://arxiv.org/pdf/2311.16208>`_ paper.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        transform (callable, optional): A function/transform that takes in an
+            :obj:`torch_geometric.data.Data` object and returns a transformed
+            version. The data object will be transformed before every access.
+            (default: :obj:`None`)
+        pre_transform (callable, optional): A function/transform that takes in
+            an :obj:`torch_geometric.data.Data` object and returns a
+            transformed version. The data object will be transformed before
+            being saved to disk. (default: :obj:`None`)
+        pre_filter (callable, optional): A function that takes in an
+            :obj:`torch_geometric.data.Data` object and returns a boolean
+            value, indicating whether the data object should be included in the
+            final dataset. (default: :obj:`None`)
+        force_reload (bool, optional): Whether to re-process the dataset.
+            (default: :obj:`False`)
+    """
+    raw_url = 'https://huggingface.co/datasets/OpenMol/PubChemSFT/blob/main'
+
+    def __init__(
+        self,
+        root: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        pre_filter: Optional[Callable] = None,
+        force_reload: bool = False,
+    ):
+        super().__init__(root, transform, pre_transform, pre_filter,
+                         force_reload=force_reload)
+        self.load(self.processed_paths[0])
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return ['all_clean.json']
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        return ['data.pt']
+
+    def download(self) -> None:
+        print('downloading dataset...')
+        fs.cp(f'{self.raw_url}/all_clean.json', self.raw_dir)
+
+    def process(self) -> None:
+        try:
+            from rdkit import Chem
+            from rdkit.Chem.rdchem import BondType as BT
+            WITH_RDKIT = True
+
+        except ImportError:
+            WITH_RDKIT = False
+
+        if not WITH_RDKIT:
+            print(("Using a pre-processed version of the dataset. Please "
+                   "install 'rdkit' to alternatively process the raw data."),
+                  file=sys.stderr)
+
+            data_list = fs.torch_load(self.raw_paths[0])
+            data_list = [Data(**data_dict) for data_dict in data_list]
+
+            if self.pre_filter is not None:
+                data_list = [d for d in data_list if self.pre_filter(d)]
+
+            if self.pre_transform is not None:
+                data_list = [self.pre_transform(d) for d in data_list]
+
+            self.save(data_list, self.processed_paths[0])
+            return
+
+        # types of atom and bond
+        types = {'H': 0, 'C': 1, 'N': 2, 'O': 3, 'F': 4, 'Unknow': 5}
+        bonds = {BT.SINGLE: 0, BT.DOUBLE: 1, BT.TRIPLE: 2, BT.AROMATIC: 3}
+
+        # load data
+        mols = json.load(open(f'{self.raw_dir}/all_clean.json'))
+
+        data_list = []
+        for smiles, qa_pairs in tqdm(mols.items(), total=len(mols)):
+            mol = Chem.MolFromSmiles(smiles)
+            if mol is None:
+                continue
+
+            x: torch.Tensor = torch.tensor([
+                types[atom.GetSymbol()] if atom.GetSymbol() in types else 5
+                for atom in mol.GetAtoms()
+            ])
+            x = one_hot(x, num_classes=len(types), dtype=torch.float)
+
+            rows, cols, edge_types = [], [], []
+            for bond in mol.GetBonds():
+                i, j = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
+                edge_types += [bonds[bond.GetBondType()]] * 2
+                rows += [i, j]
+                cols += [j, i]
+
+            edge_index = torch.tensor([rows, cols], dtype=torch.long)
+            edge_type = torch.tensor(edge_types, dtype=torch.long)
+            edge_attr = one_hot(edge_type, num_classes=len(bonds))
+
+            for question, answer in qa_pairs:
+                data = Data(
+                    x=x,
+                    edge_index=edge_index,
+                    edge_attr=edge_attr,
+                    smiles=smiles,
+                    instruction=question,
+                    y=answer,
+                )
+
+                if self.pre_filter is not None and not self.pre_filter(data):
+                    continue
+                if self.pre_transform is not None:
+                    data = self.pre_transform(data)
+
+                data_list.append(data)
+
+        self.save(data_list, self.processed_paths[0])