AITW_viscosity_surrogate 1.0.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 (26) hide show
  1. aitw_viscosity_surrogate-1.0.0/LICENSE.txt +21 -0
  2. aitw_viscosity_surrogate-1.0.0/PKG-INFO +254 -0
  3. aitw_viscosity_surrogate-1.0.0/README.md +218 -0
  4. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/__init__.py +3 -0
  5. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/chemberta.py +219 -0
  6. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/chemberta_model_base.pt +0 -0
  7. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/chemberta_model_residuals.pt +0 -0
  8. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/chemberta_residuals.py +253 -0
  9. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/cli/__init__.py +7 -0
  10. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/cli/chemberta.py +97 -0
  11. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/cli/chemberta_residuals.py +101 -0
  12. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/cli/ensemble.py +124 -0
  13. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/cli/main.py +22 -0
  14. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/cli/utils.py +7 -0
  15. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/cli/xgboost_baseline.py +94 -0
  16. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/dataset.py +262 -0
  17. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/defaults.py +26 -0
  18. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/ensemble.py +231 -0
  19. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/ensemble_config.json +5 -0
  20. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/loggers.py +92 -0
  21. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/models.py +186 -0
  22. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/utils.py +34 -0
  23. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/xgboost.py +161 -0
  24. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/xgboost_ensemble.json +1 -0
  25. aitw_viscosity_surrogate-1.0.0/aitw_viscosity_surrogate/xgboost_model.json +1 -0
  26. aitw_viscosity_surrogate-1.0.0/pyproject.toml +104 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2023 Aalto University
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,254 @@
1
+ Metadata-Version: 2.5
2
+ Name: AITW_viscosity_surrogate
3
+ Version: 1.0.0
4
+ Summary: Run model training and prediction for predicting the viscosity of molecules based on their SMILES representation.
5
+ Keywords: viscosity,molecules,SMILES,surrogate model,machine learning,deep learning,neural networks
6
+ Author-email: Fahime Seyedheydari <fahime.seyedheydari@aalto.fi>, Marcin Mińkowski <marcin.minkowski@aalto.fi>, Davide Grassano <davide.grassano@epfl.ch>
7
+ Requires-Python: >=3.11
8
+ Description-Content-Type: text/markdown
9
+ Classifier: Development Status :: 2 - Pre-Alpha
10
+ Classifier: Operating System :: Microsoft :: Windows
11
+ Classifier: Operating System :: POSIX :: Linux
12
+ Classifier: Operating System :: MacOS
13
+ Classifier: Programming Language :: Python
14
+ License-File: LICENSE.txt
15
+ Requires-Dist: numpy
16
+ Requires-Dist: scipy
17
+ Requires-Dist: pandas
18
+ Requires-Dist: matplotlib
19
+ Requires-Dist: rdkit>=2024,<2027
20
+ Requires-Dist: xgboost
21
+ Requires-Dist: torch
22
+ Requires-Dist: transformers
23
+ Requires-Dist: scikit-learn
24
+ Requires-Dist: rich_click
25
+ Requires-Dist: trogon
26
+ Requires-Dist: pre-commit ; extra == "pre-commit"
27
+ Requires-Dist: pylint ; extra == "pre-commit"
28
+ Requires-Dist: flit ; extra == "release"
29
+ Requires-Dist: pytest ; extra == "tests"
30
+ Project-URL: Source, https://github.com/AI-TranspWood/viscosity-surrogate
31
+ Provides-Extra: pre-commit
32
+ Provides-Extra: release
33
+ Provides-Extra: tests
34
+ Import-Name: aitw_viscosity_surrogate
35
+
36
+ # AITW_viscosity_surrogate
37
+
38
+ Four approaches used to predict the viscosity of chemical compounds from their SMILES representations.
39
+ Together they span classical machine learning, pure deep learning on molecular language, a hybrid architecture combining both, and a final ensemble that unifies all three.
40
+
41
+ ## Dataset
42
+
43
+ - Viscosities taken from [Katritzky et al. (2000)](https://onlinelibrary.wiley.com/doi/10.1002/(SICI)1099-1395(200001)13:1%3C80::AID-POC179%3E3.0.CO;2-8)
44
+ - SMILES strings taken from [PubChem](https://pubchem.ncbi.nlm.nih.gov/) using the API https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/{}/property/CanonicalSMILES/TXT
45
+
46
+ ## Installation
47
+
48
+ ```bash
49
+ cd <PATH to folder with pyproject.toml>
50
+ pip install .
51
+ ```
52
+
53
+ ## Usage
54
+
55
+ ### CLI
56
+
57
+ The installation will make available a `aitw-viscosity-surrogate` command line interface.
58
+
59
+ - Run `aitw-viscosity-surrogate --help` to see the available commands.
60
+
61
+ More information on the CLI can be found in the approach-specific subsections
62
+
63
+ #### Tab autocompletion
64
+
65
+ Enabling tab autocompletion https://click.palletsprojects.com/en/stable/shell-completion/
66
+
67
+ E.G for `bash` run the command
68
+
69
+ ```bash
70
+ eval "$(_AITW_INFILTRATION_SURROGATE_COMPLETE=bash_source aitw-viscosity-surrogate)"
71
+ ```
72
+
73
+ You can also add it to either `~/.bashrc` or, if you are using a virtual environment, to `bin/activate` of the virtual environment to avoid running the command for every new shell.
74
+
75
+ ## Approaches
76
+
77
+ ### Dataset format
78
+
79
+ For all approaches, the input dataset must be a CSV file containing the following columns:
80
+
81
+ - `compound_name`: The name of the compound.
82
+ - `SMILES`: The SMILES representation of the compound.
83
+ - `log_visc_calc`: The calculated log-viscosity of the compound.
84
+ - `log_visc_exp`: The experimentally measured log-viscosity of the compound.
85
+ - `visc_mPas_calc`: The calculated viscosity of the compound in mPa·s.
86
+ - `visc_mPas_exp`: The experimentally measured viscosity of the compound in mPa·
87
+
88
+
89
+ ### 1. Baseline: XGBoost with Molecular Fingerprints and Descriptors
90
+
91
+ The first approach represents a classical cheminformatics pipeline in which each molecule is converted into a fixed-length numerical vector composed of a 2048-bit Morgan fingerprint and the full set of RDKit-computed physicochemical descriptors, including molecular weight, polarity, hydrogen-bonding capacity, and topological complexity.
92
+ These features are fed into an XGBoost gradient-boosted tree regressor, which learns nonlinear relationships between molecular structure and log-viscosity without requiring any deep learning.
93
+ This approach serves as a strong, interpretable baseline against which more complex models can be compared.
94
+
95
+ #### Training
96
+
97
+ ```bash
98
+ aitw-viscosity-surrogate xgboost train --help
99
+ aitw-viscosity-surrogate xgboost train \
100
+ --input-file examples/dataset.csv
101
+ ```
102
+
103
+ The output of the training command contains the following files:
104
+
105
+ - `{output_dir}/xgboost_model.json`: The trained XGBoost model in JSON format.
106
+ - `{output_dir}/xgboost_performance.png`: A plot of predicted vs. experimental log-viscosity values for the training set.
107
+ - `{output_dir}/xgboost_predictions.csv`: A CSV file containing the predicted log-viscosity values and errors for the training set.
108
+ - `{output_dir}/xgboost_training.log`: A log file containing the training process.
109
+
110
+ #### Inference
111
+
112
+ Inference can be ran using a pre-trained model included in the package or a custom model trained using the `train` command.
113
+ The input file must be a CSV file containing the following columns:
114
+ - `compound_name`: The name of the compound.
115
+ - `SMILES`: The SMILES representation of the compound.
116
+ - `visc_mPas_exp`: OPTIONAL. The reference viscosity of the compound in mPa·s. If provided, the error between predicted and reference values will be calculated and included in the output.
117
+
118
+ ```bash
119
+ aitw-viscosity-surrogate xgboost inference --help
120
+ aitw-viscosity-surrogate xgboost inference \
121
+ --input-file examples/dataset.csv \
122
+ --output-file xgboost_predictions.csv
123
+ ```
124
+
125
+ It will produce a CSV file containing the predicted log-viscosity values for the input compounds.
126
+
127
+ The `xgboost_model.json` file included in the package was trained on the dataset provided in `examples/dataset.csv`.
128
+ <!-- It's validation is shown [in this plot](./images/xgboost_train.png). -->
129
+
130
+ ![XGBoost training performance](./images/xgboost_train.png)
131
+
132
+ ### 2. Pure LLM: Fine-Tuned ChemBERTa
133
+
134
+ The second approach treats each molecule's SMILES string as a sentence in a chemical ``language'' and fine-tunes ChemBERTa, a transformer model pretrained on millions of molecules, to predict viscosity directly from this textual representation. No hand-crafted descriptors are used; instead, the model relies entirely on the contextual embeddings it has learned during pretraining, with a small feed-forward regression head mapping the resulting molecular embedding to a single viscosity value. This approach tests whether a language model can capture the structure-property relationships needed for viscosity prediction without any explicit chemical feature engineering.
135
+
136
+ #### Training
137
+
138
+ ```bash
139
+ aitw-viscosity-surrogate chemberta train --help
140
+ aitw-viscosity-surrogate chemberta train \
141
+ --input-file examples/dataset.csv
142
+ ```
143
+
144
+ The output of the training command contains the following files:
145
+
146
+ - `{output_dir}/chemberta_model_base.pt`: The trained ChemBERTa model in PyTorch format. This will be a dictionary containing the `model_state_dict` and the ChemBERTa model name used for training (`model_name`).
147
+ - `{output_dir}/chemberta_base_performance.png`: A plot of predicted vs. experimental log-viscosity values for the training set.
148
+ - `{output_dir}/chemberta_base_predictions.csv`: A CSV file containing the predicted log-viscosity values and errors for the training set.
149
+ - `{output_dir}/chemberta_base_training.log`: A log file containing the training process.
150
+ - `{output_dir}/chemberta_base_indexes.json`: A JSON file containing the indexes of the training and test sets used for training and validation.
151
+
152
+ #### Inference
153
+
154
+ Inference can be ran using a pre-trained model included in the package or a custom model trained using the `train` command.
155
+ The input file must be a CSV file containing the following columns:
156
+ - `compound_name`: The name of the compound.
157
+ - `SMILES`: The SMILES representation of the compound.
158
+ - `visc_mPas_exp`: OPTIONAL. The reference viscosity of the compound in mPa·s. If provided, the error between predicted and reference values will be calculated and included in the output.
159
+
160
+ ```bash
161
+ aitw-viscosity-surrogate chemberta inference --help
162
+ aitw-viscosity-surrogate chemberta inference \
163
+ --input-file examples/dataset.csv \
164
+ --output-file chemberta_predictions.csv
165
+ ```
166
+
167
+ It will produce a CSV file containing the predicted log-viscosity values for the input compounds.
168
+
169
+ The `chemberta_model_base.pt` file included in the package was trained on the dataset provided in `examples/dataset.csv`.
170
+
171
+ ![ChemBERTa base training performance](./images/chemberta_train.png)
172
+
173
+
174
+ ### 3. Hybrid: ChemBERTa with Residual Descriptor Fusion
175
+
176
+ The third approach combines the strengths of the previous two methods by fusing ChemBERTa's learned molecular embedding with ten classical RDKit descriptors, projecting the combined representation through a residual neural network with skip connections to preserve information across layers. This architecture allows the model to leverage both the rich, pretrained chemical knowledge encoded in ChemBERTa and the well-established physical intuition captured by classical descriptors, resulting in improved predictive accuracy over either approach used alone.
177
+
178
+ #### Training
179
+
180
+ ```bash
181
+ aitw-viscosity-surrogate chemberta-residuals train --help
182
+ aitw-viscosity-surrogate chemberta-residuals train \
183
+ --input-file examples/dataset.csv
184
+ ```
185
+
186
+ The output of the training command contains the following files:
187
+
188
+ - `{output_dir}/chemberta_model_residuals.pt`: The trained ChemBERTa model in PyTorch format. This will be a dictionary containing the `model_state_dict`, the ChemBERTa model name used for training (`model_name`) and the scaler (`scaler`).
189
+ - `{output_dir}/chemberta_residuals_performance.png`: A plot of predicted vs. experimental log-viscosity values for the training set.
190
+ - `{output_dir}/chemberta_residuals_predictions.csv`: A CSV file containing the predicted log-viscosity values and errors for the training set.
191
+ - `{output_dir}/chemberta_residuals_training.log`: A log file containing the training process.
192
+ - `{output_dir}/chemberta_residuals_indexes.json`: A JSON file containing the indexes of the training and test sets used for training and validation.
193
+
194
+ #### Inference
195
+
196
+ Inference can be ran using a pre-trained model included in the package or a custom model trained using the `train` command.
197
+ The input file must be a CSV file containing the following columns:
198
+ - `compound_name`: The name of the compound.
199
+ - `SMILES`: The SMILES representation of the compound.
200
+ - `visc_mPas_exp`: OPTIONAL. The reference viscosity of the compound in mPa·s. If provided, the error between predicted and reference values will be calculated and included in the output.
201
+
202
+ ```bash
203
+ aitw-viscosity-surrogate chemberta-residuals inference --help
204
+ aitw-viscosity-surrogate chemberta-residuals inference \
205
+ --input-file examples/dataset.csv \
206
+ --output-file chemberta_residuals_predictions.csv
207
+ ```
208
+
209
+ It will produce a CSV file containing the predicted log-viscosity values for the input compounds.
210
+
211
+ The `chemberta_model_residuals.pt` file included in the package was trained on the dataset provided in `examples/dataset.csv`.
212
+
213
+ ![ChemBERTa residual training performance](./images/chemberta_res_train.png)
214
+
215
+
216
+ ### 4. Ensemble: Combining XGBoost and Two ChemBERTa Variants
217
+
218
+ The final approach trains three separate models, an XGBoost regressor, a simple fine-tuned ChemBERTa model, and the residual ChemBERTa-descriptor hybrid, and then searches over weighted combinations of their predictions to find the optimal blend. By systematically testing different weightings across all three models on a held-out test set, this ensemble approach demonstrates that combining diverse modeling strategies, each capturing different aspects of molecular structure, yields more robust and accurate viscosity predictions than any single model in isolation.
219
+
220
+ #### Training
221
+
222
+ The training of all the three previous approaches will be performed in a single command.
223
+ The only difference is that the XGBoost model will be trained using only the 10 descriptors used in the ChemBERTa residuals approach, instead of the full set of descriptors.
224
+
225
+ ```bash
226
+ aitw-viscosity-surrogate chemberta train --help
227
+ aitw-viscosity-surrogate chemberta train \
228
+ --input-file examples/dataset.csv
229
+ ```
230
+
231
+ The output of the training command contains the files from the three previous approaches, plus the following files:
232
+
233
+ - `{output_dir}/ensemble_config.json`: A JSON file containing the optimal weights for each model in the ensemble.
234
+ - `{output_dir}/ensemble_performance.png`: A plot of predicted vs. experimental log-viscosity values for the ensemble on the training set.
235
+ - `{output_dir}/ensemble_predictions.csv`: A CSV file containing the predicted log-viscosity values and errors for the ensemble on the training set.
236
+
237
+ #### Inference
238
+
239
+ Inference can be ran using pre-trained models included in the package or custom models trained using the `train` command.
240
+ The input file must be a CSV file containing the following columns:
241
+ - `compound_name`: The name of the compound.
242
+ - `SMILES`: The SMILES representation of the compound.
243
+
244
+ ```bash
245
+ aitw-viscosity-surrogate ensemble inference --help
246
+ aitw-viscosity-surrogate ensemble inference \
247
+ --input-file examples/dataset.csv \
248
+ --output-file ensemble_predictions.csv
249
+ ```
250
+
251
+ The output CSV will contain the predicted log-viscosity and viscosity values with all models plus the ensemble prediction denoted with a suffix of `_ensemble`.
252
+
253
+ ![Ensemble training performance](./images/ensemble_train.png)
254
+
@@ -0,0 +1,218 @@
1
+ # AITW_viscosity_surrogate
2
+
3
+ Four approaches used to predict the viscosity of chemical compounds from their SMILES representations.
4
+ Together they span classical machine learning, pure deep learning on molecular language, a hybrid architecture combining both, and a final ensemble that unifies all three.
5
+
6
+ ## Dataset
7
+
8
+ - Viscosities taken from [Katritzky et al. (2000)](https://onlinelibrary.wiley.com/doi/10.1002/(SICI)1099-1395(200001)13:1%3C80::AID-POC179%3E3.0.CO;2-8)
9
+ - SMILES strings taken from [PubChem](https://pubchem.ncbi.nlm.nih.gov/) using the API https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/{}/property/CanonicalSMILES/TXT
10
+
11
+ ## Installation
12
+
13
+ ```bash
14
+ cd <PATH to folder with pyproject.toml>
15
+ pip install .
16
+ ```
17
+
18
+ ## Usage
19
+
20
+ ### CLI
21
+
22
+ The installation will make available a `aitw-viscosity-surrogate` command line interface.
23
+
24
+ - Run `aitw-viscosity-surrogate --help` to see the available commands.
25
+
26
+ More information on the CLI can be found in the approach-specific subsections
27
+
28
+ #### Tab autocompletion
29
+
30
+ Enabling tab autocompletion https://click.palletsprojects.com/en/stable/shell-completion/
31
+
32
+ E.G for `bash` run the command
33
+
34
+ ```bash
35
+ eval "$(_AITW_INFILTRATION_SURROGATE_COMPLETE=bash_source aitw-viscosity-surrogate)"
36
+ ```
37
+
38
+ You can also add it to either `~/.bashrc` or, if you are using a virtual environment, to `bin/activate` of the virtual environment to avoid running the command for every new shell.
39
+
40
+ ## Approaches
41
+
42
+ ### Dataset format
43
+
44
+ For all approaches, the input dataset must be a CSV file containing the following columns:
45
+
46
+ - `compound_name`: The name of the compound.
47
+ - `SMILES`: The SMILES representation of the compound.
48
+ - `log_visc_calc`: The calculated log-viscosity of the compound.
49
+ - `log_visc_exp`: The experimentally measured log-viscosity of the compound.
50
+ - `visc_mPas_calc`: The calculated viscosity of the compound in mPa·s.
51
+ - `visc_mPas_exp`: The experimentally measured viscosity of the compound in mPa·
52
+
53
+
54
+ ### 1. Baseline: XGBoost with Molecular Fingerprints and Descriptors
55
+
56
+ The first approach represents a classical cheminformatics pipeline in which each molecule is converted into a fixed-length numerical vector composed of a 2048-bit Morgan fingerprint and the full set of RDKit-computed physicochemical descriptors, including molecular weight, polarity, hydrogen-bonding capacity, and topological complexity.
57
+ These features are fed into an XGBoost gradient-boosted tree regressor, which learns nonlinear relationships between molecular structure and log-viscosity without requiring any deep learning.
58
+ This approach serves as a strong, interpretable baseline against which more complex models can be compared.
59
+
60
+ #### Training
61
+
62
+ ```bash
63
+ aitw-viscosity-surrogate xgboost train --help
64
+ aitw-viscosity-surrogate xgboost train \
65
+ --input-file examples/dataset.csv
66
+ ```
67
+
68
+ The output of the training command contains the following files:
69
+
70
+ - `{output_dir}/xgboost_model.json`: The trained XGBoost model in JSON format.
71
+ - `{output_dir}/xgboost_performance.png`: A plot of predicted vs. experimental log-viscosity values for the training set.
72
+ - `{output_dir}/xgboost_predictions.csv`: A CSV file containing the predicted log-viscosity values and errors for the training set.
73
+ - `{output_dir}/xgboost_training.log`: A log file containing the training process.
74
+
75
+ #### Inference
76
+
77
+ Inference can be ran using a pre-trained model included in the package or a custom model trained using the `train` command.
78
+ The input file must be a CSV file containing the following columns:
79
+ - `compound_name`: The name of the compound.
80
+ - `SMILES`: The SMILES representation of the compound.
81
+ - `visc_mPas_exp`: OPTIONAL. The reference viscosity of the compound in mPa·s. If provided, the error between predicted and reference values will be calculated and included in the output.
82
+
83
+ ```bash
84
+ aitw-viscosity-surrogate xgboost inference --help
85
+ aitw-viscosity-surrogate xgboost inference \
86
+ --input-file examples/dataset.csv \
87
+ --output-file xgboost_predictions.csv
88
+ ```
89
+
90
+ It will produce a CSV file containing the predicted log-viscosity values for the input compounds.
91
+
92
+ The `xgboost_model.json` file included in the package was trained on the dataset provided in `examples/dataset.csv`.
93
+ <!-- It's validation is shown [in this plot](./images/xgboost_train.png). -->
94
+
95
+ ![XGBoost training performance](./images/xgboost_train.png)
96
+
97
+ ### 2. Pure LLM: Fine-Tuned ChemBERTa
98
+
99
+ The second approach treats each molecule's SMILES string as a sentence in a chemical ``language'' and fine-tunes ChemBERTa, a transformer model pretrained on millions of molecules, to predict viscosity directly from this textual representation. No hand-crafted descriptors are used; instead, the model relies entirely on the contextual embeddings it has learned during pretraining, with a small feed-forward regression head mapping the resulting molecular embedding to a single viscosity value. This approach tests whether a language model can capture the structure-property relationships needed for viscosity prediction without any explicit chemical feature engineering.
100
+
101
+ #### Training
102
+
103
+ ```bash
104
+ aitw-viscosity-surrogate chemberta train --help
105
+ aitw-viscosity-surrogate chemberta train \
106
+ --input-file examples/dataset.csv
107
+ ```
108
+
109
+ The output of the training command contains the following files:
110
+
111
+ - `{output_dir}/chemberta_model_base.pt`: The trained ChemBERTa model in PyTorch format. This will be a dictionary containing the `model_state_dict` and the ChemBERTa model name used for training (`model_name`).
112
+ - `{output_dir}/chemberta_base_performance.png`: A plot of predicted vs. experimental log-viscosity values for the training set.
113
+ - `{output_dir}/chemberta_base_predictions.csv`: A CSV file containing the predicted log-viscosity values and errors for the training set.
114
+ - `{output_dir}/chemberta_base_training.log`: A log file containing the training process.
115
+ - `{output_dir}/chemberta_base_indexes.json`: A JSON file containing the indexes of the training and test sets used for training and validation.
116
+
117
+ #### Inference
118
+
119
+ Inference can be ran using a pre-trained model included in the package or a custom model trained using the `train` command.
120
+ The input file must be a CSV file containing the following columns:
121
+ - `compound_name`: The name of the compound.
122
+ - `SMILES`: The SMILES representation of the compound.
123
+ - `visc_mPas_exp`: OPTIONAL. The reference viscosity of the compound in mPa·s. If provided, the error between predicted and reference values will be calculated and included in the output.
124
+
125
+ ```bash
126
+ aitw-viscosity-surrogate chemberta inference --help
127
+ aitw-viscosity-surrogate chemberta inference \
128
+ --input-file examples/dataset.csv \
129
+ --output-file chemberta_predictions.csv
130
+ ```
131
+
132
+ It will produce a CSV file containing the predicted log-viscosity values for the input compounds.
133
+
134
+ The `chemberta_model_base.pt` file included in the package was trained on the dataset provided in `examples/dataset.csv`.
135
+
136
+ ![ChemBERTa base training performance](./images/chemberta_train.png)
137
+
138
+
139
+ ### 3. Hybrid: ChemBERTa with Residual Descriptor Fusion
140
+
141
+ The third approach combines the strengths of the previous two methods by fusing ChemBERTa's learned molecular embedding with ten classical RDKit descriptors, projecting the combined representation through a residual neural network with skip connections to preserve information across layers. This architecture allows the model to leverage both the rich, pretrained chemical knowledge encoded in ChemBERTa and the well-established physical intuition captured by classical descriptors, resulting in improved predictive accuracy over either approach used alone.
142
+
143
+ #### Training
144
+
145
+ ```bash
146
+ aitw-viscosity-surrogate chemberta-residuals train --help
147
+ aitw-viscosity-surrogate chemberta-residuals train \
148
+ --input-file examples/dataset.csv
149
+ ```
150
+
151
+ The output of the training command contains the following files:
152
+
153
+ - `{output_dir}/chemberta_model_residuals.pt`: The trained ChemBERTa model in PyTorch format. This will be a dictionary containing the `model_state_dict`, the ChemBERTa model name used for training (`model_name`) and the scaler (`scaler`).
154
+ - `{output_dir}/chemberta_residuals_performance.png`: A plot of predicted vs. experimental log-viscosity values for the training set.
155
+ - `{output_dir}/chemberta_residuals_predictions.csv`: A CSV file containing the predicted log-viscosity values and errors for the training set.
156
+ - `{output_dir}/chemberta_residuals_training.log`: A log file containing the training process.
157
+ - `{output_dir}/chemberta_residuals_indexes.json`: A JSON file containing the indexes of the training and test sets used for training and validation.
158
+
159
+ #### Inference
160
+
161
+ Inference can be ran using a pre-trained model included in the package or a custom model trained using the `train` command.
162
+ The input file must be a CSV file containing the following columns:
163
+ - `compound_name`: The name of the compound.
164
+ - `SMILES`: The SMILES representation of the compound.
165
+ - `visc_mPas_exp`: OPTIONAL. The reference viscosity of the compound in mPa·s. If provided, the error between predicted and reference values will be calculated and included in the output.
166
+
167
+ ```bash
168
+ aitw-viscosity-surrogate chemberta-residuals inference --help
169
+ aitw-viscosity-surrogate chemberta-residuals inference \
170
+ --input-file examples/dataset.csv \
171
+ --output-file chemberta_residuals_predictions.csv
172
+ ```
173
+
174
+ It will produce a CSV file containing the predicted log-viscosity values for the input compounds.
175
+
176
+ The `chemberta_model_residuals.pt` file included in the package was trained on the dataset provided in `examples/dataset.csv`.
177
+
178
+ ![ChemBERTa residual training performance](./images/chemberta_res_train.png)
179
+
180
+
181
+ ### 4. Ensemble: Combining XGBoost and Two ChemBERTa Variants
182
+
183
+ The final approach trains three separate models, an XGBoost regressor, a simple fine-tuned ChemBERTa model, and the residual ChemBERTa-descriptor hybrid, and then searches over weighted combinations of their predictions to find the optimal blend. By systematically testing different weightings across all three models on a held-out test set, this ensemble approach demonstrates that combining diverse modeling strategies, each capturing different aspects of molecular structure, yields more robust and accurate viscosity predictions than any single model in isolation.
184
+
185
+ #### Training
186
+
187
+ The training of all the three previous approaches will be performed in a single command.
188
+ The only difference is that the XGBoost model will be trained using only the 10 descriptors used in the ChemBERTa residuals approach, instead of the full set of descriptors.
189
+
190
+ ```bash
191
+ aitw-viscosity-surrogate chemberta train --help
192
+ aitw-viscosity-surrogate chemberta train \
193
+ --input-file examples/dataset.csv
194
+ ```
195
+
196
+ The output of the training command contains the files from the three previous approaches, plus the following files:
197
+
198
+ - `{output_dir}/ensemble_config.json`: A JSON file containing the optimal weights for each model in the ensemble.
199
+ - `{output_dir}/ensemble_performance.png`: A plot of predicted vs. experimental log-viscosity values for the ensemble on the training set.
200
+ - `{output_dir}/ensemble_predictions.csv`: A CSV file containing the predicted log-viscosity values and errors for the ensemble on the training set.
201
+
202
+ #### Inference
203
+
204
+ Inference can be ran using pre-trained models included in the package or custom models trained using the `train` command.
205
+ The input file must be a CSV file containing the following columns:
206
+ - `compound_name`: The name of the compound.
207
+ - `SMILES`: The SMILES representation of the compound.
208
+
209
+ ```bash
210
+ aitw-viscosity-surrogate ensemble inference --help
211
+ aitw-viscosity-surrogate ensemble inference \
212
+ --input-file examples/dataset.csv \
213
+ --output-file ensemble_predictions.csv
214
+ ```
215
+
216
+ The output CSV will contain the predicted log-viscosity and viscosity values with all models plus the ensemble prediction denoted with a suffix of `_ensemble`.
217
+
218
+ ![Ensemble training performance](./images/ensemble_train.png)
@@ -0,0 +1,3 @@
1
+ """AITW Monomer Viscosity Surrogate"""
2
+
3
+ __version__ = '1.0.0'