torch-molecule 0.1.0.dev2__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 (155) hide show
  1. torch_molecule-0.1.0.dev2/LICENSE +21 -0
  2. torch_molecule-0.1.0.dev2/PKG-INFO +243 -0
  3. torch_molecule-0.1.0.dev2/README.md +207 -0
  4. torch_molecule-0.1.0.dev2/pyproject.toml +73 -0
  5. torch_molecule-0.1.0.dev2/setup.cfg +4 -0
  6. torch_molecule-0.1.0.dev2/tests/test_io.py +319 -0
  7. torch_molecule-0.1.0.dev2/torch_molecule/__init__.py +70 -0
  8. torch_molecule-0.1.0.dev2/torch_molecule/base/__init__.py +8 -0
  9. torch_molecule-0.1.0.dev2/torch_molecule/base/base.py +396 -0
  10. torch_molecule-0.1.0.dev2/torch_molecule/base/encoder.py +44 -0
  11. torch_molecule-0.1.0.dev2/torch_molecule/base/generator.py +24 -0
  12. torch_molecule-0.1.0.dev2/torch_molecule/base/old_encoder.py +322 -0
  13. torch_molecule-0.1.0.dev2/torch_molecule/base/old_predictor.py +904 -0
  14. torch_molecule-0.1.0.dev2/torch_molecule/base/old_predictor2.py +436 -0
  15. torch_molecule-0.1.0.dev2/torch_molecule/base/predictor.py +112 -0
  16. torch_molecule-0.1.0.dev2/torch_molecule/encoder/attrmask/__init__.py +3 -0
  17. torch_molecule-0.1.0.dev2/torch_molecule/encoder/attrmask/model.py +113 -0
  18. torch_molecule-0.1.0.dev2/torch_molecule/encoder/attrmask/modeling_attrmask.py +304 -0
  19. torch_molecule-0.1.0.dev2/torch_molecule/encoder/constant.py +13 -0
  20. torch_molecule-0.1.0.dev2/torch_molecule/encoder/contextpred/__init__.py +3 -0
  21. torch_molecule-0.1.0.dev2/torch_molecule/encoder/contextpred/model.py +207 -0
  22. torch_molecule-0.1.0.dev2/torch_molecule/encoder/contextpred/modeling_contextpred.py +299 -0
  23. torch_molecule-0.1.0.dev2/torch_molecule/encoder/contextpred/utils.py +206 -0
  24. torch_molecule-0.1.0.dev2/torch_molecule/encoder/edgepred/__init__.py +3 -0
  25. torch_molecule-0.1.0.dev2/torch_molecule/encoder/edgepred/model.py +104 -0
  26. torch_molecule-0.1.0.dev2/torch_molecule/encoder/edgepred/modeling_edgepred.py +285 -0
  27. torch_molecule-0.1.0.dev2/torch_molecule/encoder/graphmae/__init__.py +3 -0
  28. torch_molecule-0.1.0.dev2/torch_molecule/encoder/graphmae/dataloader.py +206 -0
  29. torch_molecule-0.1.0.dev2/torch_molecule/encoder/graphmae/model.py +137 -0
  30. torch_molecule-0.1.0.dev2/torch_molecule/encoder/graphmae/modeling_graphmae.py +323 -0
  31. torch_molecule-0.1.0.dev2/torch_molecule/encoder/infograph/__init__.py +3 -0
  32. torch_molecule-0.1.0.dev2/torch_molecule/encoder/infograph/loss_function.py +128 -0
  33. torch_molecule-0.1.0.dev2/torch_molecule/encoder/infograph/model.py +92 -0
  34. torch_molecule-0.1.0.dev2/torch_molecule/encoder/infograph/modeling_infograph.py +301 -0
  35. torch_molecule-0.1.0.dev2/torch_molecule/encoder/infograph/module.py +31 -0
  36. torch_molecule-0.1.0.dev2/torch_molecule/encoder/moama/__init__.py +3 -0
  37. torch_molecule-0.1.0.dev2/torch_molecule/encoder/moama/model.py +131 -0
  38. torch_molecule-0.1.0.dev2/torch_molecule/encoder/moama/modeling_moama.py +307 -0
  39. torch_molecule-0.1.0.dev2/torch_molecule/encoder/moama/utils.py +143 -0
  40. torch_molecule-0.1.0.dev2/torch_molecule/encoder/pretrained/__init__.py +3 -0
  41. torch_molecule-0.1.0.dev2/torch_molecule/encoder/pretrained/modeling_pretrained.py +274 -0
  42. torch_molecule-0.1.0.dev2/torch_molecule/encoder/supervised/__init__.py +3 -0
  43. torch_molecule-0.1.0.dev2/torch_molecule/encoder/supervised/model.py +97 -0
  44. torch_molecule-0.1.0.dev2/torch_molecule/encoder/supervised/modeling_supervised.py +354 -0
  45. torch_molecule-0.1.0.dev2/torch_molecule/generator/digress/__init__.py +3 -0
  46. torch_molecule-0.1.0.dev2/torch_molecule/generator/digress/diffusion.py +194 -0
  47. torch_molecule-0.1.0.dev2/torch_molecule/generator/digress/modeling_digress.py +503 -0
  48. torch_molecule-0.1.0.dev2/torch_molecule/generator/digress/transformer.py +360 -0
  49. torch_molecule-0.1.0.dev2/torch_molecule/generator/digress/utils.py +196 -0
  50. torch_molecule-0.1.0.dev2/torch_molecule/generator/gdss/__init__.py +3 -0
  51. torch_molecule-0.1.0.dev2/torch_molecule/generator/gdss/layers.py +257 -0
  52. torch_molecule-0.1.0.dev2/torch_molecule/generator/gdss/model.py +204 -0
  53. torch_molecule-0.1.0.dev2/torch_molecule/generator/gdss/modeling_gdss.py +474 -0
  54. torch_molecule-0.1.0.dev2/torch_molecule/generator/gdss/sde.py +334 -0
  55. torch_molecule-0.1.0.dev2/torch_molecule/generator/gdss/solver.py +227 -0
  56. torch_molecule-0.1.0.dev2/torch_molecule/generator/gdss/toremove/ScoreNetworkA.py +167 -0
  57. torch_molecule-0.1.0.dev2/torch_molecule/generator/gdss/toremove/ScoreNetworkX.py +85 -0
  58. torch_molecule-0.1.0.dev2/torch_molecule/generator/gdss/utils.py +240 -0
  59. torch_molecule-0.1.0.dev2/torch_molecule/generator/graph_dit/__init__.py +3 -0
  60. torch_molecule-0.1.0.dev2/torch_molecule/generator/graph_dit/diffusion.py +188 -0
  61. torch_molecule-0.1.0.dev2/torch_molecule/generator/graph_dit/modeling_graph_dit.py +581 -0
  62. torch_molecule-0.1.0.dev2/torch_molecule/generator/graph_dit/transformer.py +207 -0
  63. torch_molecule-0.1.0.dev2/torch_molecule/generator/graph_dit/utils.py +220 -0
  64. torch_molecule-0.1.0.dev2/torch_molecule/generator/graph_ga/__init__.py +3 -0
  65. torch_molecule-0.1.0.dev2/torch_molecule/generator/graph_ga/crossover.py +182 -0
  66. torch_molecule-0.1.0.dev2/torch_molecule/generator/graph_ga/modeling_graph_ga.py +310 -0
  67. torch_molecule-0.1.0.dev2/torch_molecule/generator/graph_ga/mutate.py +136 -0
  68. torch_molecule-0.1.0.dev2/torch_molecule/generator/graph_ga/oracle.py +95 -0
  69. torch_molecule-0.1.0.dev2/torch_molecule/generator/jtvae/__init__.py +3 -0
  70. torch_molecule-0.1.0.dev2/torch_molecule/generator/jtvae/jtnn/chemutils.py +450 -0
  71. torch_molecule-0.1.0.dev2/torch_molecule/generator/jtvae/jtnn/datautils.py +127 -0
  72. torch_molecule-0.1.0.dev2/torch_molecule/generator/jtvae/jtnn/jtmpn.py +140 -0
  73. torch_molecule-0.1.0.dev2/torch_molecule/generator/jtvae/jtnn/jtnn_dec.py +350 -0
  74. torch_molecule-0.1.0.dev2/torch_molecule/generator/jtvae/jtnn/jtnn_enc.py +131 -0
  75. torch_molecule-0.1.0.dev2/torch_molecule/generator/jtvae/jtnn/mol_tree.py +111 -0
  76. torch_molecule-0.1.0.dev2/torch_molecule/generator/jtvae/jtnn/mpn.py +129 -0
  77. torch_molecule-0.1.0.dev2/torch_molecule/generator/jtvae/jtnn/nnutils.py +66 -0
  78. torch_molecule-0.1.0.dev2/torch_molecule/generator/jtvae/jtnn/vocab.py +31 -0
  79. torch_molecule-0.1.0.dev2/torch_molecule/generator/jtvae/jtnn_vae.py +246 -0
  80. torch_molecule-0.1.0.dev2/torch_molecule/generator/jtvae/modeling_jtvae.py +247 -0
  81. torch_molecule-0.1.0.dev2/torch_molecule/generator/lstm/__init__.py +3 -0
  82. torch_molecule-0.1.0.dev2/torch_molecule/generator/lstm/action_sampler.py +91 -0
  83. torch_molecule-0.1.0.dev2/torch_molecule/generator/lstm/lstm.py +59 -0
  84. torch_molecule-0.1.0.dev2/torch_molecule/generator/lstm/modeling_lstm.py +263 -0
  85. torch_molecule-0.1.0.dev2/torch_molecule/generator/lstm/smiles_char_dict.py +123 -0
  86. torch_molecule-0.1.0.dev2/torch_molecule/generator/lstm/utils.py +62 -0
  87. torch_molecule-0.1.0.dev2/torch_molecule/generator/molgpt/__init__.py +3 -0
  88. torch_molecule-0.1.0.dev2/torch_molecule/generator/molgpt/dataset.py +87 -0
  89. torch_molecule-0.1.0.dev2/torch_molecule/generator/molgpt/gpt.py +251 -0
  90. torch_molecule-0.1.0.dev2/torch_molecule/generator/molgpt/modeling_molgpt.py +427 -0
  91. torch_molecule-0.1.0.dev2/torch_molecule/generator/pretrained/modeling_pretrained.py +1 -0
  92. torch_molecule-0.1.0.dev2/torch_molecule/nn/__init__.py +16 -0
  93. torch_molecule-0.1.0.dev2/torch_molecule/nn/attention.py +63 -0
  94. torch_molecule-0.1.0.dev2/torch_molecule/nn/embedder.py +183 -0
  95. torch_molecule-0.1.0.dev2/torch_molecule/nn/gnn.py +631 -0
  96. torch_molecule-0.1.0.dev2/torch_molecule/nn/mlp.py +34 -0
  97. torch_molecule-0.1.0.dev2/torch_molecule/predictor/bfgnn/__init__.py +3 -0
  98. torch_molecule-0.1.0.dev2/torch_molecule/predictor/bfgnn/model.py +116 -0
  99. torch_molecule-0.1.0.dev2/torch_molecule/predictor/bfgnn/modeling_bfgnn.py +81 -0
  100. torch_molecule-0.1.0.dev2/torch_molecule/predictor/dir/__init__.py +3 -0
  101. torch_molecule-0.1.0.dev2/torch_molecule/predictor/dir/model.py +235 -0
  102. torch_molecule-0.1.0.dev2/torch_molecule/predictor/dir/modeling_dir.py +161 -0
  103. torch_molecule-0.1.0.dev2/torch_molecule/predictor/dir/utils.py +40 -0
  104. torch_molecule-0.1.0.dev2/torch_molecule/predictor/gnn/__init__.py +3 -0
  105. torch_molecule-0.1.0.dev2/torch_molecule/predictor/gnn/model.py +114 -0
  106. torch_molecule-0.1.0.dev2/torch_molecule/predictor/gnn/modeling_gnn.py +706 -0
  107. torch_molecule-0.1.0.dev2/torch_molecule/predictor/gnn/old_modeling_gnn.py +707 -0
  108. torch_molecule-0.1.0.dev2/torch_molecule/predictor/grea/__init__.py +3 -0
  109. torch_molecule-0.1.0.dev2/torch_molecule/predictor/grea/model.py +190 -0
  110. torch_molecule-0.1.0.dev2/torch_molecule/predictor/grea/modeling_grea.py +107 -0
  111. torch_molecule-0.1.0.dev2/torch_molecule/predictor/grea/old1_features.py +229 -0
  112. torch_molecule-0.1.0.dev2/torch_molecule/predictor/grea/old1_modeling_grea.py +493 -0
  113. torch_molecule-0.1.0.dev2/torch_molecule/predictor/grea/utils.py +97 -0
  114. torch_molecule-0.1.0.dev2/torch_molecule/predictor/grin/__init__.py +3 -0
  115. torch_molecule-0.1.0.dev2/torch_molecule/predictor/grin/model.py +116 -0
  116. torch_molecule-0.1.0.dev2/torch_molecule/predictor/grin/modeling_grin.py +91 -0
  117. torch_molecule-0.1.0.dev2/torch_molecule/predictor/irm/__init__.py +3 -0
  118. torch_molecule-0.1.0.dev2/torch_molecule/predictor/irm/model.py +127 -0
  119. torch_molecule-0.1.0.dev2/torch_molecule/predictor/irm/modeling_irm.py +179 -0
  120. torch_molecule-0.1.0.dev2/torch_molecule/predictor/lstm/__init__.py +3 -0
  121. torch_molecule-0.1.0.dev2/torch_molecule/predictor/lstm/model.py +145 -0
  122. torch_molecule-0.1.0.dev2/torch_molecule/predictor/lstm/modeling_lstm.py +627 -0
  123. torch_molecule-0.1.0.dev2/torch_molecule/predictor/lstm/token_from_smiles.py +85 -0
  124. torch_molecule-0.1.0.dev2/torch_molecule/predictor/old_modeling_graph.py +650 -0
  125. torch_molecule-0.1.0.dev2/torch_molecule/predictor/rpgnn/__init__.py +3 -0
  126. torch_molecule-0.1.0.dev2/torch_molecule/predictor/rpgnn/model.py +185 -0
  127. torch_molecule-0.1.0.dev2/torch_molecule/predictor/rpgnn/modeling_rpgnn.py +76 -0
  128. torch_molecule-0.1.0.dev2/torch_molecule/predictor/sgir/__init__.py +3 -0
  129. torch_molecule-0.1.0.dev2/torch_molecule/predictor/sgir/modeling_sgir.py +260 -0
  130. torch_molecule-0.1.0.dev2/torch_molecule/predictor/sgir/strategy.py +266 -0
  131. torch_molecule-0.1.0.dev2/torch_molecule/predictor/smiles_transformer/__init__.py +3 -0
  132. torch_molecule-0.1.0.dev2/torch_molecule/predictor/smiles_transformer/model.py +424 -0
  133. torch_molecule-0.1.0.dev2/torch_molecule/predictor/smiles_transformer/modeling_transformer.py +252 -0
  134. torch_molecule-0.1.0.dev2/torch_molecule/predictor/ssr/__init__.py +3 -0
  135. torch_molecule-0.1.0.dev2/torch_molecule/predictor/ssr/model.py +283 -0
  136. torch_molecule-0.1.0.dev2/torch_molecule/predictor/ssr/modeling_ssr.py +267 -0
  137. torch_molecule-0.1.0.dev2/torch_molecule/predictor/ssr/utils.py +185 -0
  138. torch_molecule-0.1.0.dev2/torch_molecule/utils/__init__.py +31 -0
  139. torch_molecule-0.1.0.dev2/torch_molecule/utils/checker.py +158 -0
  140. torch_molecule-0.1.0.dev2/torch_molecule/utils/checkpoint.py +294 -0
  141. torch_molecule-0.1.0.dev2/torch_molecule/utils/format.py +140 -0
  142. torch_molecule-0.1.0.dev2/torch_molecule/utils/generic/metrics.py +530 -0
  143. torch_molecule-0.1.0.dev2/torch_molecule/utils/generic/pseudo_tasks.py +11 -0
  144. torch_molecule-0.1.0.dev2/torch_molecule/utils/generic/weights.py +30 -0
  145. torch_molecule-0.1.0.dev2/torch_molecule/utils/graph/features.py +178 -0
  146. torch_molecule-0.1.0.dev2/torch_molecule/utils/graph/graph_from_smiles.py +134 -0
  147. torch_molecule-0.1.0.dev2/torch_molecule/utils/graph/graph_to_smiles.py +312 -0
  148. torch_molecule-0.1.0.dev2/torch_molecule/utils/hf.py +232 -0
  149. torch_molecule-0.1.0.dev2/torch_molecule/utils/search.py +64 -0
  150. torch_molecule-0.1.0.dev2/torch_molecule.egg-info/PKG-INFO +243 -0
  151. torch_molecule-0.1.0.dev2/torch_molecule.egg-info/SOURCES.txt +153 -0
  152. torch_molecule-0.1.0.dev2/torch_molecule.egg-info/dependency_links.txt +1 -0
  153. torch_molecule-0.1.0.dev2/torch_molecule.egg-info/not-zip-safe +1 -0
  154. torch_molecule-0.1.0.dev2/torch_molecule.egg-info/requires.txt +20 -0
  155. torch_molecule-0.1.0.dev2/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,243 @@
1
+ Metadata-Version: 2.4
2
+ Name: torch_molecule
3
+ Version: 0.1.0.dev2
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="assets/logo.png" 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
+ > **Note**: This project is in active development, and features may change.
59
+
60
+ ## Installation
61
+
62
+ 1. **Create a Conda environment**:
63
+ ```bash
64
+ conda create --name torch_molecule python=3.11.7
65
+ conda activate torch_molecule
66
+ ```
67
+
68
+ 2. **Install `torch_molecule` from GitHub**:
69
+
70
+ Clone the repository:
71
+ ```bash
72
+ git clone https://github.com/liugangcode/torch-molecule
73
+ ```
74
+
75
+ Install the requirements:
76
+ ```bash
77
+ pip install -r requirements.txt
78
+ ```
79
+
80
+ Editable install:
81
+ ```bash
82
+ pip install -e .
83
+ ```
84
+
85
+ 3. **Install `torch_molecule` from PyPI** (Legacy):
86
+ ```bash
87
+ pip install -i https://test.pypi.org/simple/ torch-molecule
88
+ ```
89
+
90
+ ### Additional Packages
91
+
92
+ | Model | Required Packages |
93
+ |-------|-------------------|
94
+ | HFPretrainedMolecularEncoder | transformers |
95
+ | BFGNNMolecularPredictor | torch-scatter |
96
+ | GRINMolecularPredictor | torch-scatter |
97
+
98
+ ## Usage
99
+
100
+ Refer to the `tests` folder for more use cases.
101
+
102
+ ### Python API Example
103
+
104
+ The following example demonstrates how to use the `GREAMolecularPredictor` class from `torch_molecule`:
105
+
106
+ More examples could be found in the folders `examples` and `tests`.
107
+
108
+ ```python
109
+ from torch_molecule import GREAMolecularPredictor
110
+
111
+ # Train GREA model
112
+ grea_model = GREAMolecularPredictor(
113
+ num_task=num_task,
114
+ task_type="regression",
115
+ model_name="GREA_multitask",
116
+ evaluate_criterion='r2',
117
+ evaluate_higher_better=True,
118
+ verbose=True
119
+ )
120
+
121
+ # Fit the model
122
+ X_train = ['C1=CC=CC=C1', 'C1=CC=CC=C1']
123
+ y_train = [[0.5], [1.5]]
124
+ X_val = ['C1=CC=CC=C1', 'C1=CC=CC=C1']
125
+ y_val = [[0.5], [1.5]]
126
+ N_trial = 10
127
+
128
+ grea_model.autofit(
129
+ X_train=X_train.tolist(),
130
+ y_train=y_train,
131
+ X_val=X_val.tolist(),
132
+ y_val=y_val,
133
+ n_trials=N_trial,
134
+ )
135
+ ```
136
+
137
+ ### Checkpoints
138
+
139
+ `torch-molecule` provides checkpoint functions that can be interacted with on Hugging Face.
140
+
141
+ ```python
142
+ from torch_molecule import GREAMolecularPredictor
143
+ from sklearn.metrics import mean_absolute_error
144
+
145
+ # Define the repository ID for Hugging Face
146
+ repo_id = "user/repo_id"
147
+
148
+ # Initialize the GREAMolecularPredictor model
149
+ model = GREAMolecularPredictor()
150
+
151
+ # Train the model using autofit
152
+ model.autofit(
153
+ X_train=X.tolist(), # List of SMILES strings for training
154
+ y_train=y_train, # numpy array [n_samples, n_tasks] for training labels
155
+ X_val=X_val.tolist(),# List of SMILES strings for validation
156
+ y_val=y_val, # numpy array [n_samples, n_tasks] for validation labels
157
+ )
158
+
159
+ # Make predictions on the test set
160
+ output = model.predict(X_test.tolist()) # (n_sample, n_task)
161
+
162
+ # Calculate the mean absolute error
163
+ mae = mean_absolute_error(y_test, output['prediction'])
164
+ metrics = {'MAE': mae}
165
+
166
+ # Save the trained model to Hugging Face
167
+ model.save_to_hf(
168
+ repo_id=repo_id,
169
+ task_id=f"{task_name}",
170
+ metrics=metrics,
171
+ commit_message=f"Upload GREA_{task_name} model with metrics: {metrics}",
172
+ private=False
173
+ )
174
+
175
+ # Load a pretrained checkpoint from Hugging Face
176
+ model = GREAMolecularPredictor()
177
+ model.load_from_hf(repo_id=repo_id, local_cache=f"{model_dir}/GREA_{task_name}.pt")
178
+
179
+ # Set model parameters
180
+ model.set_params(verbose=True)
181
+
182
+ # Make predictions using the loaded model
183
+ predictions = model.predict(smiles_list)
184
+ ```
185
+
186
+ ## List of Supported Models
187
+
188
+ ### Predictive Models
189
+
190
+ | Model | Reference |
191
+ |----------------------|---------------------|
192
+ | GRIN | [Learning Repetition-Invariant Representations for Polymer Informatics. May 2025](https://arxiv.org/abs/2505.10726) |
193
+ | BFGNN | [Graph neural networks extrapolate out-of-distribution for shortest paths. March 2025](https://arxiv.org/abs/2503.19173) |
194
+ | SGIR | [Semi-Supervised Graph Imbalanced Regression. KDD 2023](https://dl.acm.org/doi/10.1145/3580305.3599497) |
195
+ | GREA | [Graph Rationalization with Environment-based Augmentations. KDD 2022](https://dl.acm.org/doi/abs/10.1145/3534678.3539347) |
196
+ | DIR | [Discovering Invariant Rationales for Graph Neural Networks. ICLR 2022](https://arxiv.org/abs/2201.12872) |
197
+ | SSR | [SizeShiftReg: a Regularization Method for Improving Size-Generalization in Graph Neural Networks. NeurIPS 2022](https://arxiv.org/abs/2206.07096) |
198
+ | IRM | [Invariant Risk Minimization (2019)](https://arxiv.org/abs/1907.02893) |
199
+ | RPGNN | [Relational Pooling for Graph Representations. ICML 2019](https://arxiv.org/abs/1903.02541) |
200
+ | GNNs | [Graph Convolutional Networks. ICLR 2017](https://arxiv.org/abs/1609.02907) and [Graph Isomorphism Network. ICLR 2019](https://arxiv.org/abs/1810.00826) |
201
+ | Transformer (SMILES) | [Transformer (Attention is All You Need. NeurIPS 2017)](https://arxiv.org/abs/1706.03762) based on SMILES strings |
202
+ | LSTM (SMILES) | [Long short-term memory (Neural Computation 1997)](https://ieeexplore.ieee.org/abstract/document/6795963) based on SMILES strings |
203
+
204
+ ### Generative Models
205
+
206
+ | Model | Reference |
207
+ |------------|---------------------|
208
+ | Graph DiT | [Graph Diffusion Transformers for Multi-Conditional Molecular Generation. NeurIPS 2024](https://openreview.net/forum?id=cfrDLD1wfO) |
209
+ | DiGress | [DiGress: Discrete Denoising Diffusion for Graph Generation. ICLR 2023](https://openreview.net/forum?id=UaAD-Nu86WX) |
210
+ | GDSS | [Score-based Generative Modeling of Graphs via the System of Stochastic Differential Equations. ICML 2022](https://proceedings.mlr.press/v162/jo22a/jo22a.pdf) |
211
+ | 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) |
212
+ | JTVAE | [Junction Tree Variational Autoencoder for Molecular Graph Generation. ICML 2018.](https://proceedings.mlr.press/v80/jin18a) |
213
+ | 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) |
214
+ | LSTM (SMILES) | [Long short-term memory (Neural Computation 1997)](https://ieeexplore.ieee.org/abstract/document/6795963) based on SMILES strings |
215
+
216
+ ### Representation Models
217
+
218
+ | Model | Reference |
219
+ |--------------|---------------------|
220
+ | MoAMa | [Motif-aware Attribute Masking for Molecular Graph Pre-training. LoG 2024](https://arxiv.org/abs/2309.04589) |
221
+ | GraphMAE | [GraphMAE: Self-Supervised Masked Graph Autoencoders. KDD 2022](https://arxiv.org/abs/2205.10803) |
222
+ | AttrMasking | [Strategies for Pre-training Graph Neural Networks. ICLR 2020](https://arxiv.org/abs/1905.12265) |
223
+ | ContextPred | [Strategies for Pre-training Graph Neural Networks. ICLR 2020](https://arxiv.org/abs/1905.12265) |
224
+ | EdgePred | [Strategies for Pre-training Graph Neural Networks. ICLR 2020](https://arxiv.org/abs/1905.12265) |
225
+ | InfoGraph | [InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation Learning via Mutual Information Maximization. ICLR 2020](https://arxiv.org/abs/1908.01000) |
226
+ | Supervised | Supervised pretraining |
227
+ | Pretrained | More than ten pretrained models from [Hugging Face](https://huggingface.co) |
228
+
229
+ ## Project Structure
230
+
231
+ See the structure of `torch_molecule` with the command `tree -L 2 torch_molecule -I '__pycache__|*.pyc|*.pyo|.git|old*'`
232
+
233
+ ## Plan
234
+
235
+ 1. **Predictive Models**: Done: GREA, SGIR, IRM, GIN/GCN w/ virtual, DIR. SMILES-based LSTM/Transformers. TODO more
236
+ 2. **Generative Models**: Done: Graph DiT, GraphGA, DiGress, GDS, MolGPT TODO: more
237
+ 3. **Representation Models**: Done: MoAMa, AttrMasking, ContextPred, EdgePred. Many pretrained models from HF. TODO: checkpoints, more
238
+
239
+ ## Acknowledgements
240
+
241
+ This project is under active development, and some features may change over time.
242
+
243
+ 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,207 @@
1
+ <p align="center">
2
+ <img src="assets/logo.png" 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
+ > **Note**: This project is in active development, and features may change.
23
+
24
+ ## Installation
25
+
26
+ 1. **Create a Conda environment**:
27
+ ```bash
28
+ conda create --name torch_molecule python=3.11.7
29
+ conda activate torch_molecule
30
+ ```
31
+
32
+ 2. **Install `torch_molecule` from GitHub**:
33
+
34
+ Clone the repository:
35
+ ```bash
36
+ git clone https://github.com/liugangcode/torch-molecule
37
+ ```
38
+
39
+ Install the requirements:
40
+ ```bash
41
+ pip install -r requirements.txt
42
+ ```
43
+
44
+ Editable install:
45
+ ```bash
46
+ pip install -e .
47
+ ```
48
+
49
+ 3. **Install `torch_molecule` from PyPI** (Legacy):
50
+ ```bash
51
+ pip install -i https://test.pypi.org/simple/ torch-molecule
52
+ ```
53
+
54
+ ### Additional Packages
55
+
56
+ | Model | Required Packages |
57
+ |-------|-------------------|
58
+ | HFPretrainedMolecularEncoder | transformers |
59
+ | BFGNNMolecularPredictor | torch-scatter |
60
+ | GRINMolecularPredictor | torch-scatter |
61
+
62
+ ## Usage
63
+
64
+ Refer to the `tests` folder for more use cases.
65
+
66
+ ### Python API Example
67
+
68
+ The following example demonstrates how to use the `GREAMolecularPredictor` class from `torch_molecule`:
69
+
70
+ More examples could be found in the folders `examples` and `tests`.
71
+
72
+ ```python
73
+ from torch_molecule import GREAMolecularPredictor
74
+
75
+ # Train GREA model
76
+ grea_model = GREAMolecularPredictor(
77
+ num_task=num_task,
78
+ task_type="regression",
79
+ model_name="GREA_multitask",
80
+ evaluate_criterion='r2',
81
+ evaluate_higher_better=True,
82
+ verbose=True
83
+ )
84
+
85
+ # Fit the model
86
+ X_train = ['C1=CC=CC=C1', 'C1=CC=CC=C1']
87
+ y_train = [[0.5], [1.5]]
88
+ X_val = ['C1=CC=CC=C1', 'C1=CC=CC=C1']
89
+ y_val = [[0.5], [1.5]]
90
+ N_trial = 10
91
+
92
+ grea_model.autofit(
93
+ X_train=X_train.tolist(),
94
+ y_train=y_train,
95
+ X_val=X_val.tolist(),
96
+ y_val=y_val,
97
+ n_trials=N_trial,
98
+ )
99
+ ```
100
+
101
+ ### Checkpoints
102
+
103
+ `torch-molecule` provides checkpoint functions that can be interacted with on Hugging Face.
104
+
105
+ ```python
106
+ from torch_molecule import GREAMolecularPredictor
107
+ from sklearn.metrics import mean_absolute_error
108
+
109
+ # Define the repository ID for Hugging Face
110
+ repo_id = "user/repo_id"
111
+
112
+ # Initialize the GREAMolecularPredictor model
113
+ model = GREAMolecularPredictor()
114
+
115
+ # Train the model using autofit
116
+ model.autofit(
117
+ X_train=X.tolist(), # List of SMILES strings for training
118
+ y_train=y_train, # numpy array [n_samples, n_tasks] for training labels
119
+ X_val=X_val.tolist(),# List of SMILES strings for validation
120
+ y_val=y_val, # numpy array [n_samples, n_tasks] for validation labels
121
+ )
122
+
123
+ # Make predictions on the test set
124
+ output = model.predict(X_test.tolist()) # (n_sample, n_task)
125
+
126
+ # Calculate the mean absolute error
127
+ mae = mean_absolute_error(y_test, output['prediction'])
128
+ metrics = {'MAE': mae}
129
+
130
+ # Save the trained model to Hugging Face
131
+ model.save_to_hf(
132
+ repo_id=repo_id,
133
+ task_id=f"{task_name}",
134
+ metrics=metrics,
135
+ commit_message=f"Upload GREA_{task_name} model with metrics: {metrics}",
136
+ private=False
137
+ )
138
+
139
+ # Load a pretrained checkpoint from Hugging Face
140
+ model = GREAMolecularPredictor()
141
+ model.load_from_hf(repo_id=repo_id, local_cache=f"{model_dir}/GREA_{task_name}.pt")
142
+
143
+ # Set model parameters
144
+ model.set_params(verbose=True)
145
+
146
+ # Make predictions using the loaded model
147
+ predictions = model.predict(smiles_list)
148
+ ```
149
+
150
+ ## List of Supported Models
151
+
152
+ ### Predictive Models
153
+
154
+ | Model | Reference |
155
+ |----------------------|---------------------|
156
+ | GRIN | [Learning Repetition-Invariant Representations for Polymer Informatics. May 2025](https://arxiv.org/abs/2505.10726) |
157
+ | BFGNN | [Graph neural networks extrapolate out-of-distribution for shortest paths. March 2025](https://arxiv.org/abs/2503.19173) |
158
+ | SGIR | [Semi-Supervised Graph Imbalanced Regression. KDD 2023](https://dl.acm.org/doi/10.1145/3580305.3599497) |
159
+ | GREA | [Graph Rationalization with Environment-based Augmentations. KDD 2022](https://dl.acm.org/doi/abs/10.1145/3534678.3539347) |
160
+ | DIR | [Discovering Invariant Rationales for Graph Neural Networks. ICLR 2022](https://arxiv.org/abs/2201.12872) |
161
+ | SSR | [SizeShiftReg: a Regularization Method for Improving Size-Generalization in Graph Neural Networks. NeurIPS 2022](https://arxiv.org/abs/2206.07096) |
162
+ | IRM | [Invariant Risk Minimization (2019)](https://arxiv.org/abs/1907.02893) |
163
+ | RPGNN | [Relational Pooling for Graph Representations. ICML 2019](https://arxiv.org/abs/1903.02541) |
164
+ | GNNs | [Graph Convolutional Networks. ICLR 2017](https://arxiv.org/abs/1609.02907) and [Graph Isomorphism Network. ICLR 2019](https://arxiv.org/abs/1810.00826) |
165
+ | Transformer (SMILES) | [Transformer (Attention is All You Need. NeurIPS 2017)](https://arxiv.org/abs/1706.03762) based on SMILES strings |
166
+ | LSTM (SMILES) | [Long short-term memory (Neural Computation 1997)](https://ieeexplore.ieee.org/abstract/document/6795963) based on SMILES strings |
167
+
168
+ ### Generative Models
169
+
170
+ | Model | Reference |
171
+ |------------|---------------------|
172
+ | Graph DiT | [Graph Diffusion Transformers for Multi-Conditional Molecular Generation. NeurIPS 2024](https://openreview.net/forum?id=cfrDLD1wfO) |
173
+ | DiGress | [DiGress: Discrete Denoising Diffusion for Graph Generation. ICLR 2023](https://openreview.net/forum?id=UaAD-Nu86WX) |
174
+ | GDSS | [Score-based Generative Modeling of Graphs via the System of Stochastic Differential Equations. ICML 2022](https://proceedings.mlr.press/v162/jo22a/jo22a.pdf) |
175
+ | 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) |
176
+ | JTVAE | [Junction Tree Variational Autoencoder for Molecular Graph Generation. ICML 2018.](https://proceedings.mlr.press/v80/jin18a) |
177
+ | 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) |
178
+ | LSTM (SMILES) | [Long short-term memory (Neural Computation 1997)](https://ieeexplore.ieee.org/abstract/document/6795963) based on SMILES strings |
179
+
180
+ ### Representation Models
181
+
182
+ | Model | Reference |
183
+ |--------------|---------------------|
184
+ | MoAMa | [Motif-aware Attribute Masking for Molecular Graph Pre-training. LoG 2024](https://arxiv.org/abs/2309.04589) |
185
+ | GraphMAE | [GraphMAE: Self-Supervised Masked Graph Autoencoders. KDD 2022](https://arxiv.org/abs/2205.10803) |
186
+ | AttrMasking | [Strategies for Pre-training Graph Neural Networks. ICLR 2020](https://arxiv.org/abs/1905.12265) |
187
+ | ContextPred | [Strategies for Pre-training Graph Neural Networks. ICLR 2020](https://arxiv.org/abs/1905.12265) |
188
+ | EdgePred | [Strategies for Pre-training Graph Neural Networks. ICLR 2020](https://arxiv.org/abs/1905.12265) |
189
+ | InfoGraph | [InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation Learning via Mutual Information Maximization. ICLR 2020](https://arxiv.org/abs/1908.01000) |
190
+ | Supervised | Supervised pretraining |
191
+ | Pretrained | More than ten pretrained models from [Hugging Face](https://huggingface.co) |
192
+
193
+ ## Project Structure
194
+
195
+ See the structure of `torch_molecule` with the command `tree -L 2 torch_molecule -I '__pycache__|*.pyc|*.pyo|.git|old*'`
196
+
197
+ ## Plan
198
+
199
+ 1. **Predictive Models**: Done: GREA, SGIR, IRM, GIN/GCN w/ virtual, DIR. SMILES-based LSTM/Transformers. TODO more
200
+ 2. **Generative Models**: Done: Graph DiT, GraphGA, DiGress, GDS, MolGPT TODO: more
201
+ 3. **Representation Models**: Done: MoAMa, AttrMasking, ContextPred, EdgePred. Many pretrained models from HF. TODO: checkpoints, more
202
+
203
+ ## Acknowledgements
204
+
205
+ This project is under active development, and some features may change over time.
206
+
207
+ 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.dev2"
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
+