torch-molecule 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (154) hide show
  1. torch_molecule-0.1.0/LICENSE +21 -0
  2. torch_molecule-0.1.0/PKG-INFO +240 -0
  3. torch_molecule-0.1.0/README.md +204 -0
  4. torch_molecule-0.1.0/pyproject.toml +73 -0
  5. torch_molecule-0.1.0/setup.cfg +4 -0
  6. torch_molecule-0.1.0/torch_molecule/__init__.py +70 -0
  7. torch_molecule-0.1.0/torch_molecule/base/__init__.py +8 -0
  8. torch_molecule-0.1.0/torch_molecule/base/base.py +396 -0
  9. torch_molecule-0.1.0/torch_molecule/base/encoder.py +44 -0
  10. torch_molecule-0.1.0/torch_molecule/base/generator.py +24 -0
  11. torch_molecule-0.1.0/torch_molecule/base/old_encoder.py +322 -0
  12. torch_molecule-0.1.0/torch_molecule/base/old_predictor.py +904 -0
  13. torch_molecule-0.1.0/torch_molecule/base/old_predictor2.py +436 -0
  14. torch_molecule-0.1.0/torch_molecule/base/predictor.py +112 -0
  15. torch_molecule-0.1.0/torch_molecule/encoder/attrmask/__init__.py +3 -0
  16. torch_molecule-0.1.0/torch_molecule/encoder/attrmask/model.py +113 -0
  17. torch_molecule-0.1.0/torch_molecule/encoder/attrmask/modeling_attrmask.py +304 -0
  18. torch_molecule-0.1.0/torch_molecule/encoder/constant.py +13 -0
  19. torch_molecule-0.1.0/torch_molecule/encoder/contextpred/__init__.py +3 -0
  20. torch_molecule-0.1.0/torch_molecule/encoder/contextpred/model.py +207 -0
  21. torch_molecule-0.1.0/torch_molecule/encoder/contextpred/modeling_contextpred.py +299 -0
  22. torch_molecule-0.1.0/torch_molecule/encoder/contextpred/utils.py +206 -0
  23. torch_molecule-0.1.0/torch_molecule/encoder/edgepred/__init__.py +3 -0
  24. torch_molecule-0.1.0/torch_molecule/encoder/edgepred/model.py +104 -0
  25. torch_molecule-0.1.0/torch_molecule/encoder/edgepred/modeling_edgepred.py +285 -0
  26. torch_molecule-0.1.0/torch_molecule/encoder/graphmae/__init__.py +3 -0
  27. torch_molecule-0.1.0/torch_molecule/encoder/graphmae/dataloader.py +206 -0
  28. torch_molecule-0.1.0/torch_molecule/encoder/graphmae/model.py +137 -0
  29. torch_molecule-0.1.0/torch_molecule/encoder/graphmae/modeling_graphmae.py +323 -0
  30. torch_molecule-0.1.0/torch_molecule/encoder/infograph/__init__.py +3 -0
  31. torch_molecule-0.1.0/torch_molecule/encoder/infograph/loss_function.py +128 -0
  32. torch_molecule-0.1.0/torch_molecule/encoder/infograph/model.py +92 -0
  33. torch_molecule-0.1.0/torch_molecule/encoder/infograph/modeling_infograph.py +301 -0
  34. torch_molecule-0.1.0/torch_molecule/encoder/infograph/module.py +31 -0
  35. torch_molecule-0.1.0/torch_molecule/encoder/moama/__init__.py +3 -0
  36. torch_molecule-0.1.0/torch_molecule/encoder/moama/model.py +131 -0
  37. torch_molecule-0.1.0/torch_molecule/encoder/moama/modeling_moama.py +307 -0
  38. torch_molecule-0.1.0/torch_molecule/encoder/moama/utils.py +143 -0
  39. torch_molecule-0.1.0/torch_molecule/encoder/pretrained/__init__.py +3 -0
  40. torch_molecule-0.1.0/torch_molecule/encoder/pretrained/modeling_pretrained.py +274 -0
  41. torch_molecule-0.1.0/torch_molecule/encoder/supervised/__init__.py +3 -0
  42. torch_molecule-0.1.0/torch_molecule/encoder/supervised/model.py +97 -0
  43. torch_molecule-0.1.0/torch_molecule/encoder/supervised/modeling_supervised.py +354 -0
  44. torch_molecule-0.1.0/torch_molecule/generator/digress/__init__.py +3 -0
  45. torch_molecule-0.1.0/torch_molecule/generator/digress/diffusion.py +194 -0
  46. torch_molecule-0.1.0/torch_molecule/generator/digress/modeling_digress.py +503 -0
  47. torch_molecule-0.1.0/torch_molecule/generator/digress/transformer.py +360 -0
  48. torch_molecule-0.1.0/torch_molecule/generator/digress/utils.py +196 -0
  49. torch_molecule-0.1.0/torch_molecule/generator/gdss/__init__.py +3 -0
  50. torch_molecule-0.1.0/torch_molecule/generator/gdss/layers.py +257 -0
  51. torch_molecule-0.1.0/torch_molecule/generator/gdss/model.py +204 -0
  52. torch_molecule-0.1.0/torch_molecule/generator/gdss/modeling_gdss.py +474 -0
  53. torch_molecule-0.1.0/torch_molecule/generator/gdss/sde.py +334 -0
  54. torch_molecule-0.1.0/torch_molecule/generator/gdss/solver.py +227 -0
  55. torch_molecule-0.1.0/torch_molecule/generator/gdss/toremove/ScoreNetworkA.py +167 -0
  56. torch_molecule-0.1.0/torch_molecule/generator/gdss/toremove/ScoreNetworkX.py +85 -0
  57. torch_molecule-0.1.0/torch_molecule/generator/gdss/utils.py +240 -0
  58. torch_molecule-0.1.0/torch_molecule/generator/graph_dit/__init__.py +3 -0
  59. torch_molecule-0.1.0/torch_molecule/generator/graph_dit/diffusion.py +188 -0
  60. torch_molecule-0.1.0/torch_molecule/generator/graph_dit/modeling_graph_dit.py +581 -0
  61. torch_molecule-0.1.0/torch_molecule/generator/graph_dit/transformer.py +207 -0
  62. torch_molecule-0.1.0/torch_molecule/generator/graph_dit/utils.py +220 -0
  63. torch_molecule-0.1.0/torch_molecule/generator/graph_ga/__init__.py +3 -0
  64. torch_molecule-0.1.0/torch_molecule/generator/graph_ga/crossover.py +182 -0
  65. torch_molecule-0.1.0/torch_molecule/generator/graph_ga/modeling_graph_ga.py +310 -0
  66. torch_molecule-0.1.0/torch_molecule/generator/graph_ga/mutate.py +136 -0
  67. torch_molecule-0.1.0/torch_molecule/generator/graph_ga/oracle.py +95 -0
  68. torch_molecule-0.1.0/torch_molecule/generator/jtvae/__init__.py +3 -0
  69. torch_molecule-0.1.0/torch_molecule/generator/jtvae/jtnn/chemutils.py +450 -0
  70. torch_molecule-0.1.0/torch_molecule/generator/jtvae/jtnn/datautils.py +127 -0
  71. torch_molecule-0.1.0/torch_molecule/generator/jtvae/jtnn/jtmpn.py +140 -0
  72. torch_molecule-0.1.0/torch_molecule/generator/jtvae/jtnn/jtnn_dec.py +350 -0
  73. torch_molecule-0.1.0/torch_molecule/generator/jtvae/jtnn/jtnn_enc.py +131 -0
  74. torch_molecule-0.1.0/torch_molecule/generator/jtvae/jtnn/mol_tree.py +111 -0
  75. torch_molecule-0.1.0/torch_molecule/generator/jtvae/jtnn/mpn.py +129 -0
  76. torch_molecule-0.1.0/torch_molecule/generator/jtvae/jtnn/nnutils.py +66 -0
  77. torch_molecule-0.1.0/torch_molecule/generator/jtvae/jtnn/vocab.py +31 -0
  78. torch_molecule-0.1.0/torch_molecule/generator/jtvae/jtnn_vae.py +246 -0
  79. torch_molecule-0.1.0/torch_molecule/generator/jtvae/modeling_jtvae.py +247 -0
  80. torch_molecule-0.1.0/torch_molecule/generator/lstm/__init__.py +3 -0
  81. torch_molecule-0.1.0/torch_molecule/generator/lstm/action_sampler.py +91 -0
  82. torch_molecule-0.1.0/torch_molecule/generator/lstm/lstm.py +59 -0
  83. torch_molecule-0.1.0/torch_molecule/generator/lstm/modeling_lstm.py +263 -0
  84. torch_molecule-0.1.0/torch_molecule/generator/lstm/smiles_char_dict.py +123 -0
  85. torch_molecule-0.1.0/torch_molecule/generator/lstm/utils.py +62 -0
  86. torch_molecule-0.1.0/torch_molecule/generator/molgpt/__init__.py +3 -0
  87. torch_molecule-0.1.0/torch_molecule/generator/molgpt/dataset.py +87 -0
  88. torch_molecule-0.1.0/torch_molecule/generator/molgpt/gpt.py +251 -0
  89. torch_molecule-0.1.0/torch_molecule/generator/molgpt/modeling_molgpt.py +427 -0
  90. torch_molecule-0.1.0/torch_molecule/generator/pretrained/modeling_pretrained.py +1 -0
  91. torch_molecule-0.1.0/torch_molecule/nn/__init__.py +16 -0
  92. torch_molecule-0.1.0/torch_molecule/nn/attention.py +63 -0
  93. torch_molecule-0.1.0/torch_molecule/nn/embedder.py +183 -0
  94. torch_molecule-0.1.0/torch_molecule/nn/gnn.py +631 -0
  95. torch_molecule-0.1.0/torch_molecule/nn/mlp.py +34 -0
  96. torch_molecule-0.1.0/torch_molecule/predictor/bfgnn/__init__.py +3 -0
  97. torch_molecule-0.1.0/torch_molecule/predictor/bfgnn/model.py +116 -0
  98. torch_molecule-0.1.0/torch_molecule/predictor/bfgnn/modeling_bfgnn.py +81 -0
  99. torch_molecule-0.1.0/torch_molecule/predictor/dir/__init__.py +3 -0
  100. torch_molecule-0.1.0/torch_molecule/predictor/dir/model.py +235 -0
  101. torch_molecule-0.1.0/torch_molecule/predictor/dir/modeling_dir.py +161 -0
  102. torch_molecule-0.1.0/torch_molecule/predictor/dir/utils.py +40 -0
  103. torch_molecule-0.1.0/torch_molecule/predictor/gnn/__init__.py +3 -0
  104. torch_molecule-0.1.0/torch_molecule/predictor/gnn/model.py +114 -0
  105. torch_molecule-0.1.0/torch_molecule/predictor/gnn/modeling_gnn.py +706 -0
  106. torch_molecule-0.1.0/torch_molecule/predictor/gnn/old_modeling_gnn.py +707 -0
  107. torch_molecule-0.1.0/torch_molecule/predictor/grea/__init__.py +3 -0
  108. torch_molecule-0.1.0/torch_molecule/predictor/grea/model.py +190 -0
  109. torch_molecule-0.1.0/torch_molecule/predictor/grea/modeling_grea.py +107 -0
  110. torch_molecule-0.1.0/torch_molecule/predictor/grea/old1_features.py +229 -0
  111. torch_molecule-0.1.0/torch_molecule/predictor/grea/old1_modeling_grea.py +493 -0
  112. torch_molecule-0.1.0/torch_molecule/predictor/grea/utils.py +97 -0
  113. torch_molecule-0.1.0/torch_molecule/predictor/grin/__init__.py +3 -0
  114. torch_molecule-0.1.0/torch_molecule/predictor/grin/model.py +116 -0
  115. torch_molecule-0.1.0/torch_molecule/predictor/grin/modeling_grin.py +91 -0
  116. torch_molecule-0.1.0/torch_molecule/predictor/irm/__init__.py +3 -0
  117. torch_molecule-0.1.0/torch_molecule/predictor/irm/model.py +127 -0
  118. torch_molecule-0.1.0/torch_molecule/predictor/irm/modeling_irm.py +179 -0
  119. torch_molecule-0.1.0/torch_molecule/predictor/lstm/__init__.py +3 -0
  120. torch_molecule-0.1.0/torch_molecule/predictor/lstm/model.py +145 -0
  121. torch_molecule-0.1.0/torch_molecule/predictor/lstm/modeling_lstm.py +627 -0
  122. torch_molecule-0.1.0/torch_molecule/predictor/lstm/token_from_smiles.py +85 -0
  123. torch_molecule-0.1.0/torch_molecule/predictor/old_modeling_graph.py +650 -0
  124. torch_molecule-0.1.0/torch_molecule/predictor/rpgnn/__init__.py +3 -0
  125. torch_molecule-0.1.0/torch_molecule/predictor/rpgnn/model.py +185 -0
  126. torch_molecule-0.1.0/torch_molecule/predictor/rpgnn/modeling_rpgnn.py +76 -0
  127. torch_molecule-0.1.0/torch_molecule/predictor/sgir/__init__.py +3 -0
  128. torch_molecule-0.1.0/torch_molecule/predictor/sgir/modeling_sgir.py +260 -0
  129. torch_molecule-0.1.0/torch_molecule/predictor/sgir/strategy.py +266 -0
  130. torch_molecule-0.1.0/torch_molecule/predictor/smiles_transformer/__init__.py +3 -0
  131. torch_molecule-0.1.0/torch_molecule/predictor/smiles_transformer/model.py +424 -0
  132. torch_molecule-0.1.0/torch_molecule/predictor/smiles_transformer/modeling_transformer.py +252 -0
  133. torch_molecule-0.1.0/torch_molecule/predictor/ssr/__init__.py +3 -0
  134. torch_molecule-0.1.0/torch_molecule/predictor/ssr/model.py +283 -0
  135. torch_molecule-0.1.0/torch_molecule/predictor/ssr/modeling_ssr.py +267 -0
  136. torch_molecule-0.1.0/torch_molecule/predictor/ssr/utils.py +185 -0
  137. torch_molecule-0.1.0/torch_molecule/utils/__init__.py +31 -0
  138. torch_molecule-0.1.0/torch_molecule/utils/checker.py +158 -0
  139. torch_molecule-0.1.0/torch_molecule/utils/checkpoint.py +294 -0
  140. torch_molecule-0.1.0/torch_molecule/utils/format.py +140 -0
  141. torch_molecule-0.1.0/torch_molecule/utils/generic/metrics.py +530 -0
  142. torch_molecule-0.1.0/torch_molecule/utils/generic/pseudo_tasks.py +11 -0
  143. torch_molecule-0.1.0/torch_molecule/utils/generic/weights.py +30 -0
  144. torch_molecule-0.1.0/torch_molecule/utils/graph/features.py +178 -0
  145. torch_molecule-0.1.0/torch_molecule/utils/graph/graph_from_smiles.py +134 -0
  146. torch_molecule-0.1.0/torch_molecule/utils/graph/graph_to_smiles.py +312 -0
  147. torch_molecule-0.1.0/torch_molecule/utils/hf.py +232 -0
  148. torch_molecule-0.1.0/torch_molecule/utils/search.py +64 -0
  149. torch_molecule-0.1.0/torch_molecule.egg-info/PKG-INFO +240 -0
  150. torch_molecule-0.1.0/torch_molecule.egg-info/SOURCES.txt +152 -0
  151. torch_molecule-0.1.0/torch_molecule.egg-info/dependency_links.txt +1 -0
  152. torch_molecule-0.1.0/torch_molecule.egg-info/not-zip-safe +1 -0
  153. torch_molecule-0.1.0/torch_molecule.egg-info/requires.txt +20 -0
  154. torch_molecule-0.1.0/torch_molecule.egg-info/top_level.txt +1 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2024 Haofei Yu
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,240 @@
1
+ Metadata-Version: 2.4
2
+ Name: torch_molecule
3
+ Version: 0.1.0
4
+ Summary: Deep learning packages for molecular discovery with a simple sklearn-style interface
5
+ Author-email: Gang Liu <gliu7@nd.edu>
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/liugangcode/torch-molecule
8
+ Project-URL: Bug Tracker, https://github.com/liugangcode/torch-molecule
9
+ Project-URL: Documentation, https://liugangcode.github.io/torch-molecule/
10
+ Classifier: Programming Language :: Python :: 3
11
+ Classifier: License :: OSI Approved :: MIT License
12
+ Classifier: Operating System :: OS Independent
13
+ Requires-Python: >=3.8
14
+ Description-Content-Type: text/markdown
15
+ License-File: LICENSE
16
+ Requires-Dist: torch>=2.2.0
17
+ Requires-Dist: torch-geometric>=2.6.1
18
+ Requires-Dist: numpy
19
+ Requires-Dist: pandas>=2.2.3
20
+ Requires-Dist: click
21
+ Requires-Dist: huggingface_hub>=0.22.2
22
+ Requires-Dist: joblib>=1.3.2
23
+ Requires-Dist: networkx>=3.2.1
24
+ Requires-Dist: PyYAML>=6.0.2
25
+ Requires-Dist: rdkit>=2023.9.5
26
+ Requires-Dist: scikit_learn>=1.4.1.post1
27
+ Requires-Dist: scipy>=1.14.1
28
+ Requires-Dist: tqdm>=4.66.2
29
+ Requires-Dist: optuna>=4.0.0
30
+ Provides-Extra: dev
31
+ Requires-Dist: pytest>=7.0; extra == "dev"
32
+ Requires-Dist: black; extra == "dev"
33
+ Requires-Dist: isort; extra == "dev"
34
+ Requires-Dist: flake8; extra == "dev"
35
+ Dynamic: license-file
36
+
37
+ <p align="center">
38
+ <img src="https://github.com/liugangcode/torch-molecule/blob/main/assets/logo.png?raw=true" alt="torch-molecule logo" width="600"/>
39
+ </p>
40
+
41
+ <p align="center">
42
+ <a href="https://github.com/liugangcode/torch-molecule">
43
+ <img src="https://img.shields.io/badge/GitHub-Repository-blue?logo=github" alt="GitHub Repository">
44
+ </a>
45
+ <a href="https://liugangcode.github.io/torch-molecule/">
46
+ <img src="https://img.shields.io/badge/Documentation-Online-brightgreen?logo=readthedocs" alt="Documentation">
47
+ </a>
48
+ </p>
49
+
50
+ <p align="center">
51
+ <b>Deep learning for molecular discovery with a simple sklearn-style interface</b>
52
+ </p>
53
+
54
+ ---
55
+
56
+ `torch-molecule` is a package under active development that facilitates molecular discovery through deep learning, featuring a user-friendly, `sklearn`-style interface. It includes model checkpoints for efficient deployment and benchmarking across a range of molecular tasks. Currently, the package focuses on three main components: **Predictive Models**, **Generative Models**, and **Representation Models**. See the [List of Supported Models](#list-of-supported-models) section for all available models.
57
+
58
+ ## Installation
59
+
60
+ 1. **Create a Conda environment**:
61
+ ```bash
62
+ conda create --name torch_molecule python=3.11.7
63
+ conda activate torch_molecule
64
+ ```
65
+
66
+ 2. **Install using pip (0.1.0)**:
67
+
68
+ ```bash
69
+ pip install torch-molecule
70
+ ```
71
+
72
+ 3. **Install from source for the latest version**:
73
+
74
+ Clone the repository:
75
+
76
+ ```bash
77
+ git clone https://github.com/liugangcode/torch-molecule
78
+ cd torch-molecule
79
+ ```
80
+
81
+ Install in editable mode:
82
+
83
+ ```bash
84
+ pip install .
85
+ ```
86
+
87
+ ### Additional Packages
88
+
89
+ | Model | Required Packages |
90
+ |-------|-------------------|
91
+ | HFPretrainedMolecularEncoder | transformers |
92
+ | BFGNNMolecularPredictor | torch-scatter |
93
+ | GRINMolecularPredictor | torch-scatter |
94
+
95
+ ## Usage
96
+
97
+ Refer to the `tests` folder for more use cases.
98
+
99
+ ### Python API Example
100
+
101
+ The following example demonstrates how to use the `GREAMolecularPredictor` class from `torch_molecule`:
102
+
103
+ More examples could be found in the folders `examples` and `tests`.
104
+
105
+ ```python
106
+ from torch_molecule import GREAMolecularPredictor
107
+
108
+ # Train GREA model
109
+ grea_model = GREAMolecularPredictor(
110
+ num_task=num_task,
111
+ task_type="regression",
112
+ model_name="GREA_multitask",
113
+ evaluate_criterion='r2',
114
+ evaluate_higher_better=True,
115
+ verbose=True
116
+ )
117
+
118
+ # Fit the model
119
+ X_train = ['C1=CC=CC=C1', 'C1=CC=CC=C1']
120
+ y_train = [[0.5], [1.5]]
121
+ X_val = ['C1=CC=CC=C1', 'C1=CC=CC=C1']
122
+ y_val = [[0.5], [1.5]]
123
+ N_trial = 10
124
+
125
+ grea_model.autofit(
126
+ X_train=X_train.tolist(),
127
+ y_train=y_train,
128
+ X_val=X_val.tolist(),
129
+ y_val=y_val,
130
+ n_trials=N_trial,
131
+ )
132
+ ```
133
+
134
+ ### Checkpoints
135
+
136
+ `torch-molecule` provides checkpoint functions that can be interacted with on Hugging Face.
137
+
138
+ ```python
139
+ from torch_molecule import GREAMolecularPredictor
140
+ from sklearn.metrics import mean_absolute_error
141
+
142
+ # Define the repository ID for Hugging Face
143
+ repo_id = "user/repo_id"
144
+
145
+ # Initialize the GREAMolecularPredictor model
146
+ model = GREAMolecularPredictor()
147
+
148
+ # Train the model using autofit
149
+ model.autofit(
150
+ X_train=X.tolist(), # List of SMILES strings for training
151
+ y_train=y_train, # numpy array [n_samples, n_tasks] for training labels
152
+ X_val=X_val.tolist(),# List of SMILES strings for validation
153
+ y_val=y_val, # numpy array [n_samples, n_tasks] for validation labels
154
+ )
155
+
156
+ # Make predictions on the test set
157
+ output = model.predict(X_test.tolist()) # (n_sample, n_task)
158
+
159
+ # Calculate the mean absolute error
160
+ mae = mean_absolute_error(y_test, output['prediction'])
161
+ metrics = {'MAE': mae}
162
+
163
+ # Save the trained model to Hugging Face
164
+ model.save_to_hf(
165
+ repo_id=repo_id,
166
+ task_id=f"{task_name}",
167
+ metrics=metrics,
168
+ commit_message=f"Upload GREA_{task_name} model with metrics: {metrics}",
169
+ private=False
170
+ )
171
+
172
+ # Load a pretrained checkpoint from Hugging Face
173
+ model = GREAMolecularPredictor()
174
+ model.load_from_hf(repo_id=repo_id, local_cache=f"{model_dir}/GREA_{task_name}.pt")
175
+
176
+ # Set model parameters
177
+ model.set_params(verbose=True)
178
+
179
+ # Make predictions using the loaded model
180
+ predictions = model.predict(smiles_list)
181
+ ```
182
+
183
+ ## List of Supported Models
184
+
185
+ ### Predictive Models
186
+
187
+ | Model | Reference |
188
+ |----------------------|---------------------|
189
+ | GRIN | [Learning Repetition-Invariant Representations for Polymer Informatics. May 2025](https://arxiv.org/abs/2505.10726) |
190
+ | BFGNN | [Graph neural networks extrapolate out-of-distribution for shortest paths. March 2025](https://arxiv.org/abs/2503.19173) |
191
+ | SGIR | [Semi-Supervised Graph Imbalanced Regression. KDD 2023](https://dl.acm.org/doi/10.1145/3580305.3599497) |
192
+ | GREA | [Graph Rationalization with Environment-based Augmentations. KDD 2022](https://dl.acm.org/doi/abs/10.1145/3534678.3539347) |
193
+ | DIR | [Discovering Invariant Rationales for Graph Neural Networks. ICLR 2022](https://arxiv.org/abs/2201.12872) |
194
+ | SSR | [SizeShiftReg: a Regularization Method for Improving Size-Generalization in Graph Neural Networks. NeurIPS 2022](https://arxiv.org/abs/2206.07096) |
195
+ | IRM | [Invariant Risk Minimization (2019)](https://arxiv.org/abs/1907.02893) |
196
+ | RPGNN | [Relational Pooling for Graph Representations. ICML 2019](https://arxiv.org/abs/1903.02541) |
197
+ | GNNs | [Graph Convolutional Networks. ICLR 2017](https://arxiv.org/abs/1609.02907) and [Graph Isomorphism Network. ICLR 2019](https://arxiv.org/abs/1810.00826) |
198
+ | Transformer (SMILES) | [Transformer (Attention is All You Need. NeurIPS 2017)](https://arxiv.org/abs/1706.03762) based on SMILES strings |
199
+ | LSTM (SMILES) | [Long short-term memory (Neural Computation 1997)](https://ieeexplore.ieee.org/abstract/document/6795963) based on SMILES strings |
200
+
201
+ ### Generative Models
202
+
203
+ | Model | Reference |
204
+ |------------|---------------------|
205
+ | Graph DiT | [Graph Diffusion Transformers for Multi-Conditional Molecular Generation. NeurIPS 2024](https://openreview.net/forum?id=cfrDLD1wfO) |
206
+ | DiGress | [DiGress: Discrete Denoising Diffusion for Graph Generation. ICLR 2023](https://openreview.net/forum?id=UaAD-Nu86WX) |
207
+ | GDSS | [Score-based Generative Modeling of Graphs via the System of Stochastic Differential Equations. ICML 2022](https://proceedings.mlr.press/v162/jo22a/jo22a.pdf) |
208
+ | MolGPT | [MolGPT: Molecular Generation Using a Transformer-Decoder Model. Journal of Chemical Information and Modeling 2021](https://pubs.acs.org/doi/10.1021/acs.jcim.1c00600) |
209
+ | JTVAE | [Junction Tree Variational Autoencoder for Molecular Graph Generation. ICML 2018.](https://proceedings.mlr.press/v80/jin18a) |
210
+ | GraphGA | [A Graph-Based Genetic Algorithm and Its Application to the Multiobjective Evolution of Median Molecules. Journal of Chemical Information and Computer Sciences 2004](https://pubs.acs.org/doi/10.1021/ci034290p) |
211
+ | LSTM (SMILES) | [Long short-term memory (Neural Computation 1997)](https://ieeexplore.ieee.org/abstract/document/6795963) based on SMILES strings |
212
+
213
+ ### Representation Models
214
+
215
+ | Model | Reference |
216
+ |--------------|---------------------|
217
+ | MoAMa | [Motif-aware Attribute Masking for Molecular Graph Pre-training. LoG 2024](https://arxiv.org/abs/2309.04589) |
218
+ | GraphMAE | [GraphMAE: Self-Supervised Masked Graph Autoencoders. KDD 2022](https://arxiv.org/abs/2205.10803) |
219
+ | AttrMasking | [Strategies for Pre-training Graph Neural Networks. ICLR 2020](https://arxiv.org/abs/1905.12265) |
220
+ | ContextPred | [Strategies for Pre-training Graph Neural Networks. ICLR 2020](https://arxiv.org/abs/1905.12265) |
221
+ | EdgePred | [Strategies for Pre-training Graph Neural Networks. ICLR 2020](https://arxiv.org/abs/1905.12265) |
222
+ | InfoGraph | [InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation Learning via Mutual Information Maximization. ICLR 2020](https://arxiv.org/abs/1908.01000) |
223
+ | Supervised | Supervised pretraining |
224
+ | Pretrained | More than ten pretrained models from [Hugging Face](https://huggingface.co) |
225
+
226
+ ## Project Structure
227
+
228
+ See the structure of `torch_molecule` with the command `tree -L 2 torch_molecule -I '__pycache__|*.pyc|*.pyo|.git|old*'`
229
+
230
+ ## Plan
231
+
232
+ 1. **Predictive Models**: Done: GREA, SGIR, IRM, GIN/GCN w/ virtual, DIR. SMILES-based LSTM/Transformers. TODO more
233
+ 2. **Generative Models**: Done: Graph DiT, GraphGA, DiGress, GDS, MolGPT TODO: more
234
+ 3. **Representation Models**: Done: MoAMa, AttrMasking, ContextPred, EdgePred. Many pretrained models from HF. TODO: checkpoints, more
235
+
236
+ > **Note**: This project is in active development, and features may change.
237
+
238
+ ## Acknowledgements
239
+
240
+ The project template was adapted from [https://github.com/lwaekfjlk/python-project-template](https://github.com/lwaekfjlk/python-project-template). We thank the authors for their contribution to the open-source community.
@@ -0,0 +1,204 @@
1
+ <p align="center">
2
+ <img src="https://github.com/liugangcode/torch-molecule/blob/main/assets/logo.png?raw=true" alt="torch-molecule logo" width="600"/>
3
+ </p>
4
+
5
+ <p align="center">
6
+ <a href="https://github.com/liugangcode/torch-molecule">
7
+ <img src="https://img.shields.io/badge/GitHub-Repository-blue?logo=github" alt="GitHub Repository">
8
+ </a>
9
+ <a href="https://liugangcode.github.io/torch-molecule/">
10
+ <img src="https://img.shields.io/badge/Documentation-Online-brightgreen?logo=readthedocs" alt="Documentation">
11
+ </a>
12
+ </p>
13
+
14
+ <p align="center">
15
+ <b>Deep learning for molecular discovery with a simple sklearn-style interface</b>
16
+ </p>
17
+
18
+ ---
19
+
20
+ `torch-molecule` is a package under active development that facilitates molecular discovery through deep learning, featuring a user-friendly, `sklearn`-style interface. It includes model checkpoints for efficient deployment and benchmarking across a range of molecular tasks. Currently, the package focuses on three main components: **Predictive Models**, **Generative Models**, and **Representation Models**. See the [List of Supported Models](#list-of-supported-models) section for all available models.
21
+
22
+ ## Installation
23
+
24
+ 1. **Create a Conda environment**:
25
+ ```bash
26
+ conda create --name torch_molecule python=3.11.7
27
+ conda activate torch_molecule
28
+ ```
29
+
30
+ 2. **Install using pip (0.1.0)**:
31
+
32
+ ```bash
33
+ pip install torch-molecule
34
+ ```
35
+
36
+ 3. **Install from source for the latest version**:
37
+
38
+ Clone the repository:
39
+
40
+ ```bash
41
+ git clone https://github.com/liugangcode/torch-molecule
42
+ cd torch-molecule
43
+ ```
44
+
45
+ Install in editable mode:
46
+
47
+ ```bash
48
+ pip install .
49
+ ```
50
+
51
+ ### Additional Packages
52
+
53
+ | Model | Required Packages |
54
+ |-------|-------------------|
55
+ | HFPretrainedMolecularEncoder | transformers |
56
+ | BFGNNMolecularPredictor | torch-scatter |
57
+ | GRINMolecularPredictor | torch-scatter |
58
+
59
+ ## Usage
60
+
61
+ Refer to the `tests` folder for more use cases.
62
+
63
+ ### Python API Example
64
+
65
+ The following example demonstrates how to use the `GREAMolecularPredictor` class from `torch_molecule`:
66
+
67
+ More examples could be found in the folders `examples` and `tests`.
68
+
69
+ ```python
70
+ from torch_molecule import GREAMolecularPredictor
71
+
72
+ # Train GREA model
73
+ grea_model = GREAMolecularPredictor(
74
+ num_task=num_task,
75
+ task_type="regression",
76
+ model_name="GREA_multitask",
77
+ evaluate_criterion='r2',
78
+ evaluate_higher_better=True,
79
+ verbose=True
80
+ )
81
+
82
+ # Fit the model
83
+ X_train = ['C1=CC=CC=C1', 'C1=CC=CC=C1']
84
+ y_train = [[0.5], [1.5]]
85
+ X_val = ['C1=CC=CC=C1', 'C1=CC=CC=C1']
86
+ y_val = [[0.5], [1.5]]
87
+ N_trial = 10
88
+
89
+ grea_model.autofit(
90
+ X_train=X_train.tolist(),
91
+ y_train=y_train,
92
+ X_val=X_val.tolist(),
93
+ y_val=y_val,
94
+ n_trials=N_trial,
95
+ )
96
+ ```
97
+
98
+ ### Checkpoints
99
+
100
+ `torch-molecule` provides checkpoint functions that can be interacted with on Hugging Face.
101
+
102
+ ```python
103
+ from torch_molecule import GREAMolecularPredictor
104
+ from sklearn.metrics import mean_absolute_error
105
+
106
+ # Define the repository ID for Hugging Face
107
+ repo_id = "user/repo_id"
108
+
109
+ # Initialize the GREAMolecularPredictor model
110
+ model = GREAMolecularPredictor()
111
+
112
+ # Train the model using autofit
113
+ model.autofit(
114
+ X_train=X.tolist(), # List of SMILES strings for training
115
+ y_train=y_train, # numpy array [n_samples, n_tasks] for training labels
116
+ X_val=X_val.tolist(),# List of SMILES strings for validation
117
+ y_val=y_val, # numpy array [n_samples, n_tasks] for validation labels
118
+ )
119
+
120
+ # Make predictions on the test set
121
+ output = model.predict(X_test.tolist()) # (n_sample, n_task)
122
+
123
+ # Calculate the mean absolute error
124
+ mae = mean_absolute_error(y_test, output['prediction'])
125
+ metrics = {'MAE': mae}
126
+
127
+ # Save the trained model to Hugging Face
128
+ model.save_to_hf(
129
+ repo_id=repo_id,
130
+ task_id=f"{task_name}",
131
+ metrics=metrics,
132
+ commit_message=f"Upload GREA_{task_name} model with metrics: {metrics}",
133
+ private=False
134
+ )
135
+
136
+ # Load a pretrained checkpoint from Hugging Face
137
+ model = GREAMolecularPredictor()
138
+ model.load_from_hf(repo_id=repo_id, local_cache=f"{model_dir}/GREA_{task_name}.pt")
139
+
140
+ # Set model parameters
141
+ model.set_params(verbose=True)
142
+
143
+ # Make predictions using the loaded model
144
+ predictions = model.predict(smiles_list)
145
+ ```
146
+
147
+ ## List of Supported Models
148
+
149
+ ### Predictive Models
150
+
151
+ | Model | Reference |
152
+ |----------------------|---------------------|
153
+ | GRIN | [Learning Repetition-Invariant Representations for Polymer Informatics. May 2025](https://arxiv.org/abs/2505.10726) |
154
+ | BFGNN | [Graph neural networks extrapolate out-of-distribution for shortest paths. March 2025](https://arxiv.org/abs/2503.19173) |
155
+ | SGIR | [Semi-Supervised Graph Imbalanced Regression. KDD 2023](https://dl.acm.org/doi/10.1145/3580305.3599497) |
156
+ | GREA | [Graph Rationalization with Environment-based Augmentations. KDD 2022](https://dl.acm.org/doi/abs/10.1145/3534678.3539347) |
157
+ | DIR | [Discovering Invariant Rationales for Graph Neural Networks. ICLR 2022](https://arxiv.org/abs/2201.12872) |
158
+ | SSR | [SizeShiftReg: a Regularization Method for Improving Size-Generalization in Graph Neural Networks. NeurIPS 2022](https://arxiv.org/abs/2206.07096) |
159
+ | IRM | [Invariant Risk Minimization (2019)](https://arxiv.org/abs/1907.02893) |
160
+ | RPGNN | [Relational Pooling for Graph Representations. ICML 2019](https://arxiv.org/abs/1903.02541) |
161
+ | GNNs | [Graph Convolutional Networks. ICLR 2017](https://arxiv.org/abs/1609.02907) and [Graph Isomorphism Network. ICLR 2019](https://arxiv.org/abs/1810.00826) |
162
+ | Transformer (SMILES) | [Transformer (Attention is All You Need. NeurIPS 2017)](https://arxiv.org/abs/1706.03762) based on SMILES strings |
163
+ | LSTM (SMILES) | [Long short-term memory (Neural Computation 1997)](https://ieeexplore.ieee.org/abstract/document/6795963) based on SMILES strings |
164
+
165
+ ### Generative Models
166
+
167
+ | Model | Reference |
168
+ |------------|---------------------|
169
+ | Graph DiT | [Graph Diffusion Transformers for Multi-Conditional Molecular Generation. NeurIPS 2024](https://openreview.net/forum?id=cfrDLD1wfO) |
170
+ | DiGress | [DiGress: Discrete Denoising Diffusion for Graph Generation. ICLR 2023](https://openreview.net/forum?id=UaAD-Nu86WX) |
171
+ | GDSS | [Score-based Generative Modeling of Graphs via the System of Stochastic Differential Equations. ICML 2022](https://proceedings.mlr.press/v162/jo22a/jo22a.pdf) |
172
+ | MolGPT | [MolGPT: Molecular Generation Using a Transformer-Decoder Model. Journal of Chemical Information and Modeling 2021](https://pubs.acs.org/doi/10.1021/acs.jcim.1c00600) |
173
+ | JTVAE | [Junction Tree Variational Autoencoder for Molecular Graph Generation. ICML 2018.](https://proceedings.mlr.press/v80/jin18a) |
174
+ | GraphGA | [A Graph-Based Genetic Algorithm and Its Application to the Multiobjective Evolution of Median Molecules. Journal of Chemical Information and Computer Sciences 2004](https://pubs.acs.org/doi/10.1021/ci034290p) |
175
+ | LSTM (SMILES) | [Long short-term memory (Neural Computation 1997)](https://ieeexplore.ieee.org/abstract/document/6795963) based on SMILES strings |
176
+
177
+ ### Representation Models
178
+
179
+ | Model | Reference |
180
+ |--------------|---------------------|
181
+ | MoAMa | [Motif-aware Attribute Masking for Molecular Graph Pre-training. LoG 2024](https://arxiv.org/abs/2309.04589) |
182
+ | GraphMAE | [GraphMAE: Self-Supervised Masked Graph Autoencoders. KDD 2022](https://arxiv.org/abs/2205.10803) |
183
+ | AttrMasking | [Strategies for Pre-training Graph Neural Networks. ICLR 2020](https://arxiv.org/abs/1905.12265) |
184
+ | ContextPred | [Strategies for Pre-training Graph Neural Networks. ICLR 2020](https://arxiv.org/abs/1905.12265) |
185
+ | EdgePred | [Strategies for Pre-training Graph Neural Networks. ICLR 2020](https://arxiv.org/abs/1905.12265) |
186
+ | InfoGraph | [InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation Learning via Mutual Information Maximization. ICLR 2020](https://arxiv.org/abs/1908.01000) |
187
+ | Supervised | Supervised pretraining |
188
+ | Pretrained | More than ten pretrained models from [Hugging Face](https://huggingface.co) |
189
+
190
+ ## Project Structure
191
+
192
+ See the structure of `torch_molecule` with the command `tree -L 2 torch_molecule -I '__pycache__|*.pyc|*.pyo|.git|old*'`
193
+
194
+ ## Plan
195
+
196
+ 1. **Predictive Models**: Done: GREA, SGIR, IRM, GIN/GCN w/ virtual, DIR. SMILES-based LSTM/Transformers. TODO more
197
+ 2. **Generative Models**: Done: Graph DiT, GraphGA, DiGress, GDS, MolGPT TODO: more
198
+ 3. **Representation Models**: Done: MoAMa, AttrMasking, ContextPred, EdgePred. Many pretrained models from HF. TODO: checkpoints, more
199
+
200
+ > **Note**: This project is in active development, and features may change.
201
+
202
+ ## Acknowledgements
203
+
204
+ The project template was adapted from [https://github.com/lwaekfjlk/python-project-template](https://github.com/lwaekfjlk/python-project-template). We thank the authors for their contribution to the open-source community.
@@ -0,0 +1,73 @@
1
+ [build-system]
2
+ requires = ["setuptools>=64", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "torch_molecule"
7
+ version = "0.1.0"
8
+ description = "Deep learning packages for molecular discovery with a simple sklearn-style interface"
9
+ authors = [{name = "Gang Liu", email = "gliu7@nd.edu"}]
10
+ readme = "README.md"
11
+ requires-python = ">=3.8"
12
+ license = {text = "MIT"}
13
+ classifiers = [
14
+ "Programming Language :: Python :: 3",
15
+ "License :: OSI Approved :: MIT License",
16
+ "Operating System :: OS Independent",
17
+ ]
18
+ dependencies = [
19
+ "torch>=2.2.0",
20
+ "torch-geometric>=2.6.1",
21
+ "numpy",
22
+ "pandas>=2.2.3",
23
+ "click",
24
+ "huggingface_hub>=0.22.2",
25
+ "joblib>=1.3.2",
26
+ "networkx>=3.2.1",
27
+ "PyYAML>=6.0.2",
28
+ "rdkit>=2023.9.5",
29
+ "scikit_learn>=1.4.1.post1",
30
+ "scipy>=1.14.1",
31
+ "tqdm>=4.66.2",
32
+ "optuna>=4.0.0",
33
+ ]
34
+
35
+ [project.urls]
36
+ "Homepage" = "https://github.com/liugangcode/torch-molecule"
37
+ "Bug Tracker" = "https://github.com/liugangcode/torch-molecule"
38
+ "Documentation" = "https://liugangcode.github.io/torch-molecule/"
39
+
40
+ [tool.setuptools.packages.find]
41
+ include = ["torch_molecule*"]
42
+ exclude = [
43
+ "private*",
44
+ "tests*",
45
+ "old_*.py",
46
+ "**/old_*.py"
47
+ ]
48
+
49
+ [tool.setuptools.package-data]
50
+ "torch_molecule" = ["**/*.yaml", "**/*.json"]
51
+
52
+ [tool.setuptools]
53
+ zip-safe = false
54
+
55
+ [tool.pytest.ini_options]
56
+ addopts = "--verbose"
57
+ testpaths = ["tests"]
58
+
59
+ [project.optional-dependencies]
60
+ dev = [
61
+ "pytest>=7.0",
62
+ "black",
63
+ "isort",
64
+ "flake8",
65
+ ]
66
+
67
+ [tool.black]
68
+ line-length = 100
69
+ target-version = ['py38']
70
+
71
+ [tool.isort]
72
+ profile = "black"
73
+ line_length = 100
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,70 @@
1
+ __version__ = "0.1.0"
2
+
3
+ """
4
+ predictor module
5
+ """
6
+ from .predictor.grea import GREAMolecularPredictor
7
+ from .predictor.gnn import GNNMolecularPredictor
8
+ from .predictor.sgir import SGIRMolecularPredictor
9
+ from .predictor.irm import IRMMolecularPredictor
10
+ from .predictor.ssr import SSRMolecularPredictor
11
+ from .predictor.dir import DIRMolecularPredictor
12
+ from .predictor.rpgnn import RPGNNMolecularPredictor
13
+ from .predictor.bfgnn import BFGNNMolecularPredictor
14
+ from .predictor.grin import GRINMolecularPredictor
15
+ from .predictor.lstm import LSTMMolecularPredictor
16
+ from .predictor.smiles_transformer import SMILESTransformerMolecularPredictor
17
+ """
18
+ encoder module
19
+ """
20
+ from .encoder.attrmask import AttrMaskMolecularEncoder
21
+ from .encoder.moama import MoamaMolecularEncoder
22
+ from .encoder.graphmae import GraphMAEMolecularEncoder
23
+ from .encoder.supervised import SupervisedMolecularEncoder
24
+ from .encoder.contextpred import ContextPredMolecularEncoder
25
+ from .encoder.edgepred import EdgePredMolecularEncoder
26
+ from .encoder.infograph import InfoGraphMolecularEncoder
27
+ from .encoder.pretrained import HFPretrainedMolecularEncoder
28
+ """
29
+ generator module
30
+ """
31
+ from .generator.graph_dit import GraphDITMolecularGenerator
32
+ from .generator.digress import DigressMolecularGenerator
33
+ from .generator.gdss import GDSSMolecularGenerator
34
+ from .generator.graph_ga import GraphGAMolecularGenerator
35
+ from .generator.jtvae import JTVAEMolecularGenerator
36
+ from .generator.lstm import LSTMMolecularGenerator
37
+ from .generator.molgpt import MolGPTMolecularGenerator
38
+
39
+ __all__ = [
40
+ # 'BaseMolecularPredictor',
41
+ # predictors
42
+ 'SGIRMolecularPredictor',
43
+ 'GREAMolecularPredictor',
44
+ 'GNNMolecularPredictor',
45
+ 'IRMMolecularPredictor',
46
+ 'SSRMolecularPredictor',
47
+ 'DIRMolecularPredictor',
48
+ 'RPGNNMolecularPredictor',
49
+ 'BFGNNMolecularPredictor',
50
+ 'GRINMolecularPredictor',
51
+ 'LSTMMolecularPredictor',
52
+ 'SMILESTransformerMolecularPredictor',
53
+ # encoders
54
+ 'SupervisedMolecularEncoder',
55
+ 'AttrMaskMolecularEncoder',
56
+ 'MoamaMolecularEncoder',
57
+ 'GraphMAEMolecularEncoder',
58
+ 'ContextPredMolecularEncoder',
59
+ 'EdgePredMolecularEncoder',
60
+ 'InfoGraphMolecularEncoder',
61
+ 'HFPretrainedMolecularEncoder',
62
+ # generators
63
+ 'GraphDITMolecularGenerator',
64
+ 'DigressMolecularGenerator',
65
+ 'GDSSMolecularGenerator',
66
+ 'GraphGAMolecularGenerator',
67
+ 'JTVAEMolecularGenerator',
68
+ 'MolGPTMolecularGenerator',
69
+ 'LSTMMolecularGenerator',
70
+ ]
@@ -0,0 +1,8 @@
1
+ from .predictor import BaseMolecularPredictor
2
+ from .encoder import BaseMolecularEncoder
3
+ from .generator import BaseMolecularGenerator
4
+ __all__ = [
5
+ "BaseMolecularPredictor",
6
+ "BaseMolecularEncoder",
7
+ "BaseMolecularGenerator",
8
+ ]