Metadata-Version: 2.4
Name: uma-finetune-wrapper
Version: 0.1.0.dev2
Summary: A simple Python interface for UMA dataset preparation and fine-tuning with FAIRChem.
Author: Timmo-hendrik Pukk
License-Expression: MIT
Keywords: UMA,FAIRChem,fine-tuning,atomistic,machine-learning
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.12
Classifier: Operating System :: POSIX :: Linux
Classifier: Topic :: Scientific/Engineering :: Chemistry
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.12
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: src/uma_finetune/resources/FAIRChem-LICENSE.md
License-File: src/uma_finetune/resources/NOTICE.md
Requires-Dist: fairchem-core==2.22.0
Requires-Dist: PyYAML>=6.0
Dynamic: license-file

# UMA fine-tuning wrapper

Prepare a dataset and fine-tune a UMA model through two Python functions:
`create_dataset()` and `finetune()`.

This package wraps FAIRChem's existing dataset-generation and training tools.
It provides options for backbone freezing, energy and force loss coefficients,
and basic training settings. The training loop and model implementation remain
in FAIRChem.

The distribution name is **uma-finetune-wrapper**; the Python import is
**uma_finetune**. This is an independent wrapper, not an official FAIRChem package.

## Installation

Python 3.12 and Linux are the tested development environment. The wrapper pins
`fairchem-core==2.22.0`; compatibility with other FAIRChem versions has not been
established.

Install a PyTorch build suitable for your hardware first. The tested Tesla V100
configuration uses:

```bash
python -m pip install 'torch==2.13.0' \
    --index-url https://download.pytorch.org/whl/cu126
```

See the [PyTorch installation instructions](https://pytorch.org/get-started/locally/)
for other hardware. This CUDA build is not a recommendation for every GPU.

Before publication, install from the package source directory:

```bash
python -m pip install .
```

To build and install a wheel instead:

```bash
python -m pip wheel --no-deps --wheel-dir dist .
python -m pip install dist/uma_finetune_wrapper-0.1.0.dev2-py3-none-any.whl
python -m pip check
```

The installation command installs the declared dependencies. A separate FAIRChem
source checkout is not needed: this package includes the required YAML templates.
You also need access to the chosen UMA model and any authentication required by
its [model repository](https://huggingface.co/facebook/UMA). Model weights are not
included in this package.

## Quick start

Prepare separate training and validation directories containing labelled
ASE-readable files, such as `.db`, `.traj`, or `.extxyz` files. The labels must be
readable as ASE calculator results. Split your data before calling the wrapper;
for related structures, keep the appropriate groups together to avoid leakage.

Save this as a Python script and run it in your compute allocation:

```python
from uma_finetune import create_dataset, finetune

if __name__ == '__main__':
    dataset = create_dataset(
        train_path='data/train',
        val_path='data/validation',
        output_dir='data/prepared',
        task='oc25',
        targets='ef',
        base_model='uma-s-1p2',
        num_workers=1,
    )

    output = finetune(
        dataset=dataset,
        output_dir='runs/example',
        mode='head_only',
        energy_coef=20,
        force_coef=2,
        epochs=1,
        batch_size=4,
        lr=1e-4,
        device='cuda',
    )

    print('Run output root:', output)
```

`data/prepared` must not already exist. To train again using prepared data, call
only `finetune(dataset='data/prepared', ...)`. Each call starts a new training run.
On an HPC cluster, submit training inside a GPU allocation; the wrapper does not
submit scheduler jobs.

## Dataset options

`create_dataset(train_path, val_path, output_dir, ...)` accepts directories for
both inputs and returns the absolute prepared-directory path.

| Option | Default | Meaning |
| --- | --- | --- |
| `task` | `'oc25'` | UMA task: `omol`, `omat`, `odac`, `oc20`, `oc25`, or `omc` |
| `targets` | `'ef'` | `e`: energy; `ef`: energy and forces; `efs`: also stress |
| `base_model` | `None` | Model name for the generated configuration; `None` uses FAIRChem's default |
| `num_workers` | `1` | Positive number of conversion workers |

Use the task corresponding to your data and pretrained model. Dataset conversion,
normalization, and element-reference fitting are performed by FAIRChem. Its input
requirements apply; real-data testing of this wrapper has used `targets='ef'`.

## Training options

`finetune(dataset, output_dir, ...)` uses the generated configuration and returns
the absolute output root. FAIRChem creates a timestamped run directory inside it.

| Option | Default | Meaning |
| --- | --- | --- |
| `model` | `None` | Pretrained model name; `None` keeps the dataset configuration's model |
| `mode` | `'full'` | `full`, `head_only`, or `partial` |
| `freeze_blocks` | `0` | Number of initial backbone blocks frozen in partial mode |
| `energy_coef` | `20` | Energy loss coefficient |
| `force_coef` | `2` | Force loss coefficient |
| `epochs` | `30` | Total training epochs |
| `batch_size` | `4` | Training and validation batch size |
| `lr` | `1e-4` | Learning rate |
| `device` | `'cuda'` | `cuda` or `cpu` |
| `max_neighbors` | `300` | Model neighbor limit |
| `max_saved_checkpoints` | `2` | Retention limit passed to FAIRChem's checkpoint callback |

Loss coefficients must be finite and nonnegative. Set either coefficient to
`None` to keep its YAML setting. For energy-only training, use `force_coef=None`.
Stress training retains its generated loss coefficient.

### Freezing

- **full:** all model parameters are trainable.
- **head_only:** the backbone is frozen; parameters outside it remain trainable.
- **partial:** the first `freeze_blocks` blocks are frozen; subsequent blocks
  and heads are trainable. Other backbone parameters remain frozen.

For example, freeze the first block:

```python
output = finetune(
    dataset='data/prepared',
    output_dir='runs/partial',
    mode='partial',
    freeze_blocks=1,
    epochs=1,
    batch_size=1,
)
```

`freeze_blocks` must be between zero and the model's block count. A nonzero value
is only accepted in partial mode. Partial mode with zero frozen blocks still
freezes backbone parameters outside the blocks, so it differs from full mode.
The initializer prints the number of trainable parameters.

## Use a trained model

After a completed run, find:

```text
<output_dir>/<run_id>/checkpoints/final/inference_ckpt.pt
```

Load it through FAIRChem using the same task used for fine-tuning:

```python
from ase.io import read
from fairchem.core import FAIRChemCalculator
from fairchem.core.units.mlip_unit import load_predict_unit

predictor = load_predict_unit(
    'runs/example/<run_id>/checkpoints/final/inference_ckpt.pt',
    device='cuda',
)
atoms = read('data/validation/example.db', index=0)
atoms.calc = FAIRChemCalculator(predictor, task_name='oc25')

energy = atoms.get_potential_energy()  # eV
forces = atoms.get_forces(apply_constraint=False)  # eV/angstrom
```

Replace `<run_id>` and the input filename with actual paths. The wrapper does
not currently implement checkpoint resume or automatic best-checkpoint selection.

## Validation status

The following checks were completed on 17 October 2026:

- Dataset conversion with the bundled templates: 2,130 training structures.
- Head-only GPU training, final checkpoint loading, and finite energy/force
  predictions on a 54-atom validation structure.
- Partial training with `freeze_blocks=1`, producing a final checkpoint.
- Fresh-environment wheel installation followed by dataset conversion,
  head-only GPU training, and checkpoint inference without an editable FAIRChem
  checkout.

The tested combination was Python 3.12, FAIRChem 2.22.0, PyTorch 2.13.0+cu126,
and a Tesla V100-SXM2-32GB. Full-mode GPU training, CPU training, other tasks,
and real-data energy-only/stress training have not been tested for this wrapper.
These are functionality checks, not evidence of model accuracy or generalization.

Lightweight wrapper tests can be run from the source directory:

```bash
python -m unittest discover -s tests -v
```

For development, install with `python -m pip install -e .`. Keep a snapshot of
an environment that passed your tests using `python -m pip freeze > requirements-tested.txt`.
This snapshot records an environment; it does not replace the package's dependency
metadata or the CUDA-specific installation instructions.

## License and attribution

The wrapper is distributed under the MIT license; see `LICENSE`.
The bundled FAIRChem 2.22.0 YAML templates retain their original copyright and
MIT license in `src/uma_finetune/resources/FAIRChem-LICENSE.md`, with provenance
in `src/uma_finetune/resources/NOTICE.md`.

The wrapper's license does not change the terms for UMA model weights, datasets,
or other dependencies. Follow the upstream model's terms and citation guidance
when using it in research.

Upstream: [FAIRChem](https://github.com/facebookresearch/fairchem).
