uma-finetune-wrapper 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 (21) hide show
  1. uma_finetune_wrapper-0.1.0.dev2/LICENSE +21 -0
  2. uma_finetune_wrapper-0.1.0.dev2/PKG-INFO +249 -0
  3. uma_finetune_wrapper-0.1.0.dev2/README.md +226 -0
  4. uma_finetune_wrapper-0.1.0.dev2/pyproject.toml +43 -0
  5. uma_finetune_wrapper-0.1.0.dev2/setup.cfg +4 -0
  6. uma_finetune_wrapper-0.1.0.dev2/src/uma_finetune/__init__.py +6 -0
  7. uma_finetune_wrapper-0.1.0.dev2/src/uma_finetune/dataset.py +62 -0
  8. uma_finetune_wrapper-0.1.0.dev2/src/uma_finetune/resources/FAIRChem-LICENSE.md +23 -0
  9. uma_finetune_wrapper-0.1.0.dev2/src/uma_finetune/resources/NOTICE.md +13 -0
  10. uma_finetune_wrapper-0.1.0.dev2/src/uma_finetune/resources/configs/uma/finetune/data/uma_conserving_data_task_energy.yaml +89 -0
  11. uma_finetune_wrapper-0.1.0.dev2/src/uma_finetune/resources/configs/uma/finetune/data/uma_conserving_data_task_energy_force.yaml +100 -0
  12. uma_finetune_wrapper-0.1.0.dev2/src/uma_finetune/resources/configs/uma/finetune/data/uma_conserving_data_task_energy_force_stress.yaml +107 -0
  13. uma_finetune_wrapper-0.1.0.dev2/src/uma_finetune/resources/configs/uma/finetune/uma_sm_finetune_template.yaml +118 -0
  14. uma_finetune_wrapper-0.1.0.dev2/src/uma_finetune/training.py +153 -0
  15. uma_finetune_wrapper-0.1.0.dev2/src/uma_finetune_wrapper.egg-info/PKG-INFO +249 -0
  16. uma_finetune_wrapper-0.1.0.dev2/src/uma_finetune_wrapper.egg-info/SOURCES.txt +19 -0
  17. uma_finetune_wrapper-0.1.0.dev2/src/uma_finetune_wrapper.egg-info/dependency_links.txt +1 -0
  18. uma_finetune_wrapper-0.1.0.dev2/src/uma_finetune_wrapper.egg-info/requires.txt +2 -0
  19. uma_finetune_wrapper-0.1.0.dev2/src/uma_finetune_wrapper.egg-info/top_level.txt +1 -0
  20. uma_finetune_wrapper-0.1.0.dev2/tests/test_dataset.py +69 -0
  21. uma_finetune_wrapper-0.1.0.dev2/tests/test_training.py +105 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Timmo-hendrik Pukk
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,249 @@
1
+ Metadata-Version: 2.4
2
+ Name: uma-finetune-wrapper
3
+ Version: 0.1.0.dev2
4
+ Summary: A simple Python interface for UMA dataset preparation and fine-tuning with FAIRChem.
5
+ Author: Timmo-hendrik Pukk
6
+ License-Expression: MIT
7
+ Keywords: UMA,FAIRChem,fine-tuning,atomistic,machine-learning
8
+ Classifier: Development Status :: 3 - Alpha
9
+ Classifier: Intended Audience :: Science/Research
10
+ Classifier: Programming Language :: Python :: 3
11
+ Classifier: Programming Language :: Python :: 3.12
12
+ Classifier: Operating System :: POSIX :: Linux
13
+ Classifier: Topic :: Scientific/Engineering :: Chemistry
14
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
15
+ Requires-Python: >=3.12
16
+ Description-Content-Type: text/markdown
17
+ License-File: LICENSE
18
+ License-File: src/uma_finetune/resources/FAIRChem-LICENSE.md
19
+ License-File: src/uma_finetune/resources/NOTICE.md
20
+ Requires-Dist: fairchem-core==2.22.0
21
+ Requires-Dist: PyYAML>=6.0
22
+ Dynamic: license-file
23
+
24
+ # UMA fine-tuning wrapper
25
+
26
+ Prepare a dataset and fine-tune a UMA model through two Python functions:
27
+ `create_dataset()` and `finetune()`.
28
+
29
+ This package wraps FAIRChem's existing dataset-generation and training tools.
30
+ It provides options for backbone freezing, energy and force loss coefficients,
31
+ and basic training settings. The training loop and model implementation remain
32
+ in FAIRChem.
33
+
34
+ The distribution name is **uma-finetune-wrapper**; the Python import is
35
+ **uma_finetune**. This is an independent wrapper, not an official FAIRChem package.
36
+
37
+ ## Installation
38
+
39
+ Python 3.12 and Linux are the tested development environment. The wrapper pins
40
+ `fairchem-core==2.22.0`; compatibility with other FAIRChem versions has not been
41
+ established.
42
+
43
+ Install a PyTorch build suitable for your hardware first. The tested Tesla V100
44
+ configuration uses:
45
+
46
+ ```bash
47
+ python -m pip install 'torch==2.13.0' \
48
+ --index-url https://download.pytorch.org/whl/cu126
49
+ ```
50
+
51
+ See the [PyTorch installation instructions](https://pytorch.org/get-started/locally/)
52
+ for other hardware. This CUDA build is not a recommendation for every GPU.
53
+
54
+ Before publication, install from the package source directory:
55
+
56
+ ```bash
57
+ python -m pip install .
58
+ ```
59
+
60
+ To build and install a wheel instead:
61
+
62
+ ```bash
63
+ python -m pip wheel --no-deps --wheel-dir dist .
64
+ python -m pip install dist/uma_finetune_wrapper-0.1.0.dev2-py3-none-any.whl
65
+ python -m pip check
66
+ ```
67
+
68
+ The installation command installs the declared dependencies. A separate FAIRChem
69
+ source checkout is not needed: this package includes the required YAML templates.
70
+ You also need access to the chosen UMA model and any authentication required by
71
+ its [model repository](https://huggingface.co/facebook/UMA). Model weights are not
72
+ included in this package.
73
+
74
+ ## Quick start
75
+
76
+ Prepare separate training and validation directories containing labelled
77
+ ASE-readable files, such as `.db`, `.traj`, or `.extxyz` files. The labels must be
78
+ readable as ASE calculator results. Split your data before calling the wrapper;
79
+ for related structures, keep the appropriate groups together to avoid leakage.
80
+
81
+ Save this as a Python script and run it in your compute allocation:
82
+
83
+ ```python
84
+ from uma_finetune import create_dataset, finetune
85
+
86
+ if __name__ == '__main__':
87
+ dataset = create_dataset(
88
+ train_path='data/train',
89
+ val_path='data/validation',
90
+ output_dir='data/prepared',
91
+ task='oc25',
92
+ targets='ef',
93
+ base_model='uma-s-1p2',
94
+ num_workers=1,
95
+ )
96
+
97
+ output = finetune(
98
+ dataset=dataset,
99
+ output_dir='runs/example',
100
+ mode='head_only',
101
+ energy_coef=20,
102
+ force_coef=2,
103
+ epochs=1,
104
+ batch_size=4,
105
+ lr=1e-4,
106
+ device='cuda',
107
+ )
108
+
109
+ print('Run output root:', output)
110
+ ```
111
+
112
+ `data/prepared` must not already exist. To train again using prepared data, call
113
+ only `finetune(dataset='data/prepared', ...)`. Each call starts a new training run.
114
+ On an HPC cluster, submit training inside a GPU allocation; the wrapper does not
115
+ submit scheduler jobs.
116
+
117
+ ## Dataset options
118
+
119
+ `create_dataset(train_path, val_path, output_dir, ...)` accepts directories for
120
+ both inputs and returns the absolute prepared-directory path.
121
+
122
+ | Option | Default | Meaning |
123
+ | --- | --- | --- |
124
+ | `task` | `'oc25'` | UMA task: `omol`, `omat`, `odac`, `oc20`, `oc25`, or `omc` |
125
+ | `targets` | `'ef'` | `e`: energy; `ef`: energy and forces; `efs`: also stress |
126
+ | `base_model` | `None` | Model name for the generated configuration; `None` uses FAIRChem's default |
127
+ | `num_workers` | `1` | Positive number of conversion workers |
128
+
129
+ Use the task corresponding to your data and pretrained model. Dataset conversion,
130
+ normalization, and element-reference fitting are performed by FAIRChem. Its input
131
+ requirements apply; real-data testing of this wrapper has used `targets='ef'`.
132
+
133
+ ## Training options
134
+
135
+ `finetune(dataset, output_dir, ...)` uses the generated configuration and returns
136
+ the absolute output root. FAIRChem creates a timestamped run directory inside it.
137
+
138
+ | Option | Default | Meaning |
139
+ | --- | --- | --- |
140
+ | `model` | `None` | Pretrained model name; `None` keeps the dataset configuration's model |
141
+ | `mode` | `'full'` | `full`, `head_only`, or `partial` |
142
+ | `freeze_blocks` | `0` | Number of initial backbone blocks frozen in partial mode |
143
+ | `energy_coef` | `20` | Energy loss coefficient |
144
+ | `force_coef` | `2` | Force loss coefficient |
145
+ | `epochs` | `30` | Total training epochs |
146
+ | `batch_size` | `4` | Training and validation batch size |
147
+ | `lr` | `1e-4` | Learning rate |
148
+ | `device` | `'cuda'` | `cuda` or `cpu` |
149
+ | `max_neighbors` | `300` | Model neighbor limit |
150
+ | `max_saved_checkpoints` | `2` | Retention limit passed to FAIRChem's checkpoint callback |
151
+
152
+ Loss coefficients must be finite and nonnegative. Set either coefficient to
153
+ `None` to keep its YAML setting. For energy-only training, use `force_coef=None`.
154
+ Stress training retains its generated loss coefficient.
155
+
156
+ ### Freezing
157
+
158
+ - **full:** all model parameters are trainable.
159
+ - **head_only:** the backbone is frozen; parameters outside it remain trainable.
160
+ - **partial:** the first `freeze_blocks` blocks are frozen; subsequent blocks
161
+ and heads are trainable. Other backbone parameters remain frozen.
162
+
163
+ For example, freeze the first block:
164
+
165
+ ```python
166
+ output = finetune(
167
+ dataset='data/prepared',
168
+ output_dir='runs/partial',
169
+ mode='partial',
170
+ freeze_blocks=1,
171
+ epochs=1,
172
+ batch_size=1,
173
+ )
174
+ ```
175
+
176
+ `freeze_blocks` must be between zero and the model's block count. A nonzero value
177
+ is only accepted in partial mode. Partial mode with zero frozen blocks still
178
+ freezes backbone parameters outside the blocks, so it differs from full mode.
179
+ The initializer prints the number of trainable parameters.
180
+
181
+ ## Use a trained model
182
+
183
+ After a completed run, find:
184
+
185
+ ```text
186
+ <output_dir>/<run_id>/checkpoints/final/inference_ckpt.pt
187
+ ```
188
+
189
+ Load it through FAIRChem using the same task used for fine-tuning:
190
+
191
+ ```python
192
+ from ase.io import read
193
+ from fairchem.core import FAIRChemCalculator
194
+ from fairchem.core.units.mlip_unit import load_predict_unit
195
+
196
+ predictor = load_predict_unit(
197
+ 'runs/example/<run_id>/checkpoints/final/inference_ckpt.pt',
198
+ device='cuda',
199
+ )
200
+ atoms = read('data/validation/example.db', index=0)
201
+ atoms.calc = FAIRChemCalculator(predictor, task_name='oc25')
202
+
203
+ energy = atoms.get_potential_energy() # eV
204
+ forces = atoms.get_forces(apply_constraint=False) # eV/angstrom
205
+ ```
206
+
207
+ Replace `<run_id>` and the input filename with actual paths. The wrapper does
208
+ not currently implement checkpoint resume or automatic best-checkpoint selection.
209
+
210
+ ## Validation status
211
+
212
+ The following checks were completed on 17 October 2026:
213
+
214
+ - Dataset conversion with the bundled templates: 2,130 training structures.
215
+ - Head-only GPU training, final checkpoint loading, and finite energy/force
216
+ predictions on a 54-atom validation structure.
217
+ - Partial training with `freeze_blocks=1`, producing a final checkpoint.
218
+ - Fresh-environment wheel installation followed by dataset conversion,
219
+ head-only GPU training, and checkpoint inference without an editable FAIRChem
220
+ checkout.
221
+
222
+ The tested combination was Python 3.12, FAIRChem 2.22.0, PyTorch 2.13.0+cu126,
223
+ and a Tesla V100-SXM2-32GB. Full-mode GPU training, CPU training, other tasks,
224
+ and real-data energy-only/stress training have not been tested for this wrapper.
225
+ These are functionality checks, not evidence of model accuracy or generalization.
226
+
227
+ Lightweight wrapper tests can be run from the source directory:
228
+
229
+ ```bash
230
+ python -m unittest discover -s tests -v
231
+ ```
232
+
233
+ For development, install with `python -m pip install -e .`. Keep a snapshot of
234
+ an environment that passed your tests using `python -m pip freeze > requirements-tested.txt`.
235
+ This snapshot records an environment; it does not replace the package's dependency
236
+ metadata or the CUDA-specific installation instructions.
237
+
238
+ ## License and attribution
239
+
240
+ The wrapper is distributed under the MIT license; see `LICENSE`.
241
+ The bundled FAIRChem 2.22.0 YAML templates retain their original copyright and
242
+ MIT license in `src/uma_finetune/resources/FAIRChem-LICENSE.md`, with provenance
243
+ in `src/uma_finetune/resources/NOTICE.md`.
244
+
245
+ The wrapper's license does not change the terms for UMA model weights, datasets,
246
+ or other dependencies. Follow the upstream model's terms and citation guidance
247
+ when using it in research.
248
+
249
+ Upstream: [FAIRChem](https://github.com/facebookresearch/fairchem).
@@ -0,0 +1,226 @@
1
+ # UMA fine-tuning wrapper
2
+
3
+ Prepare a dataset and fine-tune a UMA model through two Python functions:
4
+ `create_dataset()` and `finetune()`.
5
+
6
+ This package wraps FAIRChem's existing dataset-generation and training tools.
7
+ It provides options for backbone freezing, energy and force loss coefficients,
8
+ and basic training settings. The training loop and model implementation remain
9
+ in FAIRChem.
10
+
11
+ The distribution name is **uma-finetune-wrapper**; the Python import is
12
+ **uma_finetune**. This is an independent wrapper, not an official FAIRChem package.
13
+
14
+ ## Installation
15
+
16
+ Python 3.12 and Linux are the tested development environment. The wrapper pins
17
+ `fairchem-core==2.22.0`; compatibility with other FAIRChem versions has not been
18
+ established.
19
+
20
+ Install a PyTorch build suitable for your hardware first. The tested Tesla V100
21
+ configuration uses:
22
+
23
+ ```bash
24
+ python -m pip install 'torch==2.13.0' \
25
+ --index-url https://download.pytorch.org/whl/cu126
26
+ ```
27
+
28
+ See the [PyTorch installation instructions](https://pytorch.org/get-started/locally/)
29
+ for other hardware. This CUDA build is not a recommendation for every GPU.
30
+
31
+ Before publication, install from the package source directory:
32
+
33
+ ```bash
34
+ python -m pip install .
35
+ ```
36
+
37
+ To build and install a wheel instead:
38
+
39
+ ```bash
40
+ python -m pip wheel --no-deps --wheel-dir dist .
41
+ python -m pip install dist/uma_finetune_wrapper-0.1.0.dev2-py3-none-any.whl
42
+ python -m pip check
43
+ ```
44
+
45
+ The installation command installs the declared dependencies. A separate FAIRChem
46
+ source checkout is not needed: this package includes the required YAML templates.
47
+ You also need access to the chosen UMA model and any authentication required by
48
+ its [model repository](https://huggingface.co/facebook/UMA). Model weights are not
49
+ included in this package.
50
+
51
+ ## Quick start
52
+
53
+ Prepare separate training and validation directories containing labelled
54
+ ASE-readable files, such as `.db`, `.traj`, or `.extxyz` files. The labels must be
55
+ readable as ASE calculator results. Split your data before calling the wrapper;
56
+ for related structures, keep the appropriate groups together to avoid leakage.
57
+
58
+ Save this as a Python script and run it in your compute allocation:
59
+
60
+ ```python
61
+ from uma_finetune import create_dataset, finetune
62
+
63
+ if __name__ == '__main__':
64
+ dataset = create_dataset(
65
+ train_path='data/train',
66
+ val_path='data/validation',
67
+ output_dir='data/prepared',
68
+ task='oc25',
69
+ targets='ef',
70
+ base_model='uma-s-1p2',
71
+ num_workers=1,
72
+ )
73
+
74
+ output = finetune(
75
+ dataset=dataset,
76
+ output_dir='runs/example',
77
+ mode='head_only',
78
+ energy_coef=20,
79
+ force_coef=2,
80
+ epochs=1,
81
+ batch_size=4,
82
+ lr=1e-4,
83
+ device='cuda',
84
+ )
85
+
86
+ print('Run output root:', output)
87
+ ```
88
+
89
+ `data/prepared` must not already exist. To train again using prepared data, call
90
+ only `finetune(dataset='data/prepared', ...)`. Each call starts a new training run.
91
+ On an HPC cluster, submit training inside a GPU allocation; the wrapper does not
92
+ submit scheduler jobs.
93
+
94
+ ## Dataset options
95
+
96
+ `create_dataset(train_path, val_path, output_dir, ...)` accepts directories for
97
+ both inputs and returns the absolute prepared-directory path.
98
+
99
+ | Option | Default | Meaning |
100
+ | --- | --- | --- |
101
+ | `task` | `'oc25'` | UMA task: `omol`, `omat`, `odac`, `oc20`, `oc25`, or `omc` |
102
+ | `targets` | `'ef'` | `e`: energy; `ef`: energy and forces; `efs`: also stress |
103
+ | `base_model` | `None` | Model name for the generated configuration; `None` uses FAIRChem's default |
104
+ | `num_workers` | `1` | Positive number of conversion workers |
105
+
106
+ Use the task corresponding to your data and pretrained model. Dataset conversion,
107
+ normalization, and element-reference fitting are performed by FAIRChem. Its input
108
+ requirements apply; real-data testing of this wrapper has used `targets='ef'`.
109
+
110
+ ## Training options
111
+
112
+ `finetune(dataset, output_dir, ...)` uses the generated configuration and returns
113
+ the absolute output root. FAIRChem creates a timestamped run directory inside it.
114
+
115
+ | Option | Default | Meaning |
116
+ | --- | --- | --- |
117
+ | `model` | `None` | Pretrained model name; `None` keeps the dataset configuration's model |
118
+ | `mode` | `'full'` | `full`, `head_only`, or `partial` |
119
+ | `freeze_blocks` | `0` | Number of initial backbone blocks frozen in partial mode |
120
+ | `energy_coef` | `20` | Energy loss coefficient |
121
+ | `force_coef` | `2` | Force loss coefficient |
122
+ | `epochs` | `30` | Total training epochs |
123
+ | `batch_size` | `4` | Training and validation batch size |
124
+ | `lr` | `1e-4` | Learning rate |
125
+ | `device` | `'cuda'` | `cuda` or `cpu` |
126
+ | `max_neighbors` | `300` | Model neighbor limit |
127
+ | `max_saved_checkpoints` | `2` | Retention limit passed to FAIRChem's checkpoint callback |
128
+
129
+ Loss coefficients must be finite and nonnegative. Set either coefficient to
130
+ `None` to keep its YAML setting. For energy-only training, use `force_coef=None`.
131
+ Stress training retains its generated loss coefficient.
132
+
133
+ ### Freezing
134
+
135
+ - **full:** all model parameters are trainable.
136
+ - **head_only:** the backbone is frozen; parameters outside it remain trainable.
137
+ - **partial:** the first `freeze_blocks` blocks are frozen; subsequent blocks
138
+ and heads are trainable. Other backbone parameters remain frozen.
139
+
140
+ For example, freeze the first block:
141
+
142
+ ```python
143
+ output = finetune(
144
+ dataset='data/prepared',
145
+ output_dir='runs/partial',
146
+ mode='partial',
147
+ freeze_blocks=1,
148
+ epochs=1,
149
+ batch_size=1,
150
+ )
151
+ ```
152
+
153
+ `freeze_blocks` must be between zero and the model's block count. A nonzero value
154
+ is only accepted in partial mode. Partial mode with zero frozen blocks still
155
+ freezes backbone parameters outside the blocks, so it differs from full mode.
156
+ The initializer prints the number of trainable parameters.
157
+
158
+ ## Use a trained model
159
+
160
+ After a completed run, find:
161
+
162
+ ```text
163
+ <output_dir>/<run_id>/checkpoints/final/inference_ckpt.pt
164
+ ```
165
+
166
+ Load it through FAIRChem using the same task used for fine-tuning:
167
+
168
+ ```python
169
+ from ase.io import read
170
+ from fairchem.core import FAIRChemCalculator
171
+ from fairchem.core.units.mlip_unit import load_predict_unit
172
+
173
+ predictor = load_predict_unit(
174
+ 'runs/example/<run_id>/checkpoints/final/inference_ckpt.pt',
175
+ device='cuda',
176
+ )
177
+ atoms = read('data/validation/example.db', index=0)
178
+ atoms.calc = FAIRChemCalculator(predictor, task_name='oc25')
179
+
180
+ energy = atoms.get_potential_energy() # eV
181
+ forces = atoms.get_forces(apply_constraint=False) # eV/angstrom
182
+ ```
183
+
184
+ Replace `<run_id>` and the input filename with actual paths. The wrapper does
185
+ not currently implement checkpoint resume or automatic best-checkpoint selection.
186
+
187
+ ## Validation status
188
+
189
+ The following checks were completed on 17 October 2026:
190
+
191
+ - Dataset conversion with the bundled templates: 2,130 training structures.
192
+ - Head-only GPU training, final checkpoint loading, and finite energy/force
193
+ predictions on a 54-atom validation structure.
194
+ - Partial training with `freeze_blocks=1`, producing a final checkpoint.
195
+ - Fresh-environment wheel installation followed by dataset conversion,
196
+ head-only GPU training, and checkpoint inference without an editable FAIRChem
197
+ checkout.
198
+
199
+ The tested combination was Python 3.12, FAIRChem 2.22.0, PyTorch 2.13.0+cu126,
200
+ and a Tesla V100-SXM2-32GB. Full-mode GPU training, CPU training, other tasks,
201
+ and real-data energy-only/stress training have not been tested for this wrapper.
202
+ These are functionality checks, not evidence of model accuracy or generalization.
203
+
204
+ Lightweight wrapper tests can be run from the source directory:
205
+
206
+ ```bash
207
+ python -m unittest discover -s tests -v
208
+ ```
209
+
210
+ For development, install with `python -m pip install -e .`. Keep a snapshot of
211
+ an environment that passed your tests using `python -m pip freeze > requirements-tested.txt`.
212
+ This snapshot records an environment; it does not replace the package's dependency
213
+ metadata or the CUDA-specific installation instructions.
214
+
215
+ ## License and attribution
216
+
217
+ The wrapper is distributed under the MIT license; see `LICENSE`.
218
+ The bundled FAIRChem 2.22.0 YAML templates retain their original copyright and
219
+ MIT license in `src/uma_finetune/resources/FAIRChem-LICENSE.md`, with provenance
220
+ in `src/uma_finetune/resources/NOTICE.md`.
221
+
222
+ The wrapper's license does not change the terms for UMA model weights, datasets,
223
+ or other dependencies. Follow the upstream model's terms and citation guidance
224
+ when using it in research.
225
+
226
+ Upstream: [FAIRChem](https://github.com/facebookresearch/fairchem).
@@ -0,0 +1,43 @@
1
+ [build-system]
2
+ requires = ['setuptools>=77.0.3']
3
+ build-backend = 'setuptools.build_meta'
4
+
5
+ [project]
6
+ name = 'uma-finetune-wrapper'
7
+ version = '0.1.0.dev2'
8
+ description = 'A simple Python interface for UMA dataset preparation and fine-tuning with FAIRChem.'
9
+ readme = 'README.md'
10
+ requires-python = '>=3.12'
11
+ license = 'MIT'
12
+ license-files = [
13
+ 'LICENSE',
14
+ 'src/uma_finetune/resources/FAIRChem-LICENSE.md',
15
+ 'src/uma_finetune/resources/NOTICE.md',
16
+ ]
17
+ authors = [{name = 'Timmo-hendrik Pukk'}]
18
+ keywords = ['UMA', 'FAIRChem', 'fine-tuning', 'atomistic', 'machine-learning']
19
+ classifiers = [
20
+ 'Development Status :: 3 - Alpha',
21
+ 'Intended Audience :: Science/Research',
22
+ 'Programming Language :: Python :: 3',
23
+ 'Programming Language :: Python :: 3.12',
24
+ 'Operating System :: POSIX :: Linux',
25
+ 'Topic :: Scientific/Engineering :: Chemistry',
26
+ 'Topic :: Scientific/Engineering :: Artificial Intelligence',
27
+ ]
28
+ dependencies = [
29
+ 'fairchem-core==2.22.0',
30
+ 'PyYAML>=6.0',
31
+ ]
32
+
33
+ [tool.setuptools.packages.find]
34
+ where = ['src']
35
+ include = ['uma_finetune*']
36
+ namespaces = true
37
+
38
+ [tool.setuptools.package-data]
39
+ uma_finetune = [
40
+ 'resources/*.md',
41
+ 'resources/configs/uma/finetune/*.yaml',
42
+ 'resources/configs/uma/finetune/data/*.yaml',
43
+ ]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,6 @@
1
+ '''Simple UMA dataset preparation and fine-tuning tools.'''
2
+
3
+ from .dataset import create_dataset
4
+ from .training import finetune
5
+
6
+ __all__ = ['create_dataset', 'finetune']
@@ -0,0 +1,62 @@
1
+ '''Prepare datasets with FAIRChem's existing UMA conversion script.'''
2
+
3
+ import subprocess
4
+ import sys
5
+ from importlib.resources import as_file, files
6
+ from pathlib import Path
7
+
8
+
9
+ def create_dataset(
10
+ train_path,
11
+ val_path,
12
+ output_dir,
13
+ *,
14
+ task='oc25',
15
+ targets='ef',
16
+ base_model=None,
17
+ num_workers=1,
18
+ ):
19
+ '''Convert directories of labelled ASE-readable files and return their output path.
20
+
21
+ targets: 'e' (energies), 'ef' (energies and forces), or 'efs' (also stress).
22
+ base_model=None uses FAIRChem's default for the generated configuration.
23
+ output_dir must not already exist. Templates are included in the package.
24
+ Conversion failures raise subprocess.CalledProcessError; logs remain visible.
25
+ '''
26
+ train_path = Path(train_path).expanduser().resolve()
27
+ val_path = Path(val_path).expanduser().resolve()
28
+ output_dir = Path(output_dir).expanduser().resolve()
29
+
30
+ for path in (train_path, val_path):
31
+ if not path.is_dir():
32
+ raise NotADirectoryError(f'Input directory does not exist: {path}')
33
+
34
+ if output_dir.exists():
35
+ raise FileExistsError(f'Output directory already exists: {output_dir}')
36
+ if task not in ('omol', 'omat', 'odac', 'oc20', 'oc25', 'omc'):
37
+ raise ValueError('task must be omol, omat, odac, oc20, oc25, or omc.')
38
+ if targets not in ('e', 'ef', 'efs'):
39
+ raise ValueError('targets must be e, ef, or efs.')
40
+ if type(num_workers) is not int or num_workers < 1:
41
+ raise ValueError('num_workers must be a positive integer.')
42
+
43
+ command = [
44
+ sys.executable,
45
+ '-m', 'fairchem.core.scripts.create_uma_finetune_dataset',
46
+ '--train-dir', str(train_path),
47
+ '--val-dir', str(val_path),
48
+ '--output-dir', str(output_dir),
49
+ '--uma-task', task,
50
+ '--regression-tasks', targets,
51
+ '--num-workers', str(num_workers),
52
+ ]
53
+
54
+ if base_model is not None:
55
+ command.extend(['--base-model', base_model])
56
+
57
+ # FAIRChem reads configs/uma/finetune relative to its working directory.
58
+ # Input and output paths are absolute, so only template lookup changes.
59
+ resources = files('uma_finetune').joinpath('resources')
60
+ with as_file(resources) as resource_root:
61
+ subprocess.run(command, check=True, cwd=resource_root)
62
+ return output_dir
@@ -0,0 +1,23 @@
1
+ # License
2
+
3
+ MIT License
4
+
5
+ Copyright (c) Meta, Inc. and its affiliates.
6
+
7
+ Permission is hereby granted, free of charge, to any person obtaining a copy
8
+ of this software and associated documentation files (the "Software"), to deal
9
+ in the Software without restriction, including without limitation the rights
10
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11
+ copies of the Software, and to permit persons to whom the Software is
12
+ furnished to do so, subject to the following conditions:
13
+
14
+ The above copyright notice and this permission notice shall be included in all
15
+ copies or substantial portions of the Software.
16
+
17
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
23
+ SOFTWARE.
@@ -0,0 +1,13 @@
1
+ # FAIRChem templates
2
+
3
+ The four YAML files under configs/uma/finetune are unmodified copies from
4
+ FAIRChem 2.22.0, supplied from the working source checkout for this project.
5
+ Upstream project: https://github.com/facebookresearch/fairchem
6
+
7
+ Copyright (c) Meta, Inc. and its affiliates.
8
+ These templates are distributed under the MIT license reproduced in
9
+ FAIRChem-LICENSE.md. This notice does not assign a license to the wrapper itself.
10
+
11
+ The conversion script is supplied by the installed fairchem-core dependency;
12
+ it is not copied into this package. Keep these resources matched to the pinned
13
+ FAIRChem version when upgrading.