uma-finetune-wrapper 0.1.0.dev2__py3-none-any.whl

uma_finetune/__init__.py

@@ -0,0 +1,6 @@
+'''Simple UMA dataset preparation and fine-tuning tools.'''
+
+from .dataset import create_dataset
+from .training import finetune
+
+__all__ = ['create_dataset', 'finetune']

uma_finetune/dataset.py

@@ -0,0 +1,62 @@
+'''Prepare datasets with FAIRChem's existing UMA conversion script.'''
+
+import subprocess
+import sys
+from importlib.resources import as_file, files
+from pathlib import Path
+
+
+def create_dataset(
+    train_path,
+    val_path,
+    output_dir,
+    *,
+    task='oc25',
+    targets='ef',
+    base_model=None,
+    num_workers=1,
+):
+    '''Convert directories of labelled ASE-readable files and return their output path.
+
+    targets: 'e' (energies), 'ef' (energies and forces), or 'efs' (also stress).
+    base_model=None uses FAIRChem's default for the generated configuration.
+    output_dir must not already exist. Templates are included in the package.
+    Conversion failures raise subprocess.CalledProcessError; logs remain visible.
+    '''
+    train_path = Path(train_path).expanduser().resolve()
+    val_path = Path(val_path).expanduser().resolve()
+    output_dir = Path(output_dir).expanduser().resolve()
+
+    for path in (train_path, val_path):
+        if not path.is_dir():
+            raise NotADirectoryError(f'Input directory does not exist: {path}')
+
+    if output_dir.exists():
+        raise FileExistsError(f'Output directory already exists: {output_dir}')
+    if task not in ('omol', 'omat', 'odac', 'oc20', 'oc25', 'omc'):
+        raise ValueError('task must be omol, omat, odac, oc20, oc25, or omc.')
+    if targets not in ('e', 'ef', 'efs'):
+        raise ValueError('targets must be e, ef, or efs.')
+    if type(num_workers) is not int or num_workers < 1:
+        raise ValueError('num_workers must be a positive integer.')
+
+    command = [
+        sys.executable,
+        '-m', 'fairchem.core.scripts.create_uma_finetune_dataset',
+        '--train-dir', str(train_path),
+        '--val-dir', str(val_path),
+        '--output-dir', str(output_dir),
+        '--uma-task', task,
+        '--regression-tasks', targets,
+        '--num-workers', str(num_workers),
+    ]
+
+    if base_model is not None:
+        command.extend(['--base-model', base_model])
+
+    # FAIRChem reads configs/uma/finetune relative to its working directory.
+    # Input and output paths are absolute, so only template lookup changes.
+    resources = files('uma_finetune').joinpath('resources')
+    with as_file(resources) as resource_root:
+        subprocess.run(command, check=True, cwd=resource_root)
+    return output_dir

uma_finetune/resources/FAIRChem-LICENSE.md

@@ -0,0 +1,23 @@
+# License
+
+MIT License
+
+Copyright (c) Meta, Inc. and its affiliates.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

uma_finetune/resources/NOTICE.md

@@ -0,0 +1,13 @@
+# FAIRChem templates
+
+The four YAML files under configs/uma/finetune are unmodified copies from
+FAIRChem 2.22.0, supplied from the working source checkout for this project.
+Upstream project: https://github.com/facebookresearch/fairchem
+
+Copyright (c) Meta, Inc. and its affiliates.
+These templates are distributed under the MIT license reproduced in
+FAIRChem-LICENSE.md. This notice does not assign a license to the wrapper itself.
+
+The conversion script is supplied by the installed fairchem-core dependency;
+it is not copied into this package. Keep these resources matched to the pinned
+FAIRChem version when upgrading.

uma_finetune/resources/configs/uma/finetune/data/uma_conserving_data_task_energy.yaml

@@ -0,0 +1,89 @@
+dataset_name: ??
+elem_refs: ??
+normalizer_rmsd: ??
+
+train_dataset:
+  splits:
+    train:
+      src: ??
+  format: ase_db
+  transforms:
+    common_transform:
+      dataset_name: ${data.dataset_name}
+    stress_reshape_transform:
+      dataset_name: ${data.dataset_name}
+
+val_dataset:
+  splits:
+    val:
+      src: ??
+  format: ase_db
+  transforms:
+    common_transform:
+      dataset_name: ${data.dataset_name}
+    stress_reshape_transform:
+      dataset_name: ${data.dataset_name}
+
+regress_stress: True
+pass_through_head_outputs: True
+
+heads:
+  efs:
+    module: fairchem.core.models.uma.escn_md.MLP_EFS_Head
+
+tasks_list:
+- _target_: fairchem.core.units.mlip_unit.mlip_unit.Task
+  name: energy
+  level: system
+  property: energy
+  loss_fn:
+    _target_: fairchem.core.modules.loss.DDPMTLoss
+    loss_fn:
+      _target_: fairchem.core.modules.loss.PerAtomMAELoss
+    coefficient: 20
+  out_spec:
+    dim: [1]
+    dtype: float32
+  normalizer:
+    _target_: fairchem.core.modules.normalization.normalizer.Normalizer
+    mean: 0.0
+    rmsd: ${data.normalizer_rmsd}
+  element_references:
+    _target_: fairchem.core.modules.normalization.element_references.ElementReferences
+    element_references:
+      _target_: torch.DoubleTensor
+      _args_:
+        - ${data.elem_refs}
+  datasets:
+    - ${data.dataset_name}
+  metrics:
+    - mae
+    - per_atom_mae
+- _target_: fairchem.core.units.mlip_unit.mlip_unit.Task
+  name: forces
+  level: atom
+  property: forces
+  out_spec:
+    dim: [3]
+    dtype: float32
+  normalizer:
+    _target_: fairchem.core.modules.normalization.normalizer.Normalizer
+    mean: 0.0
+    rmsd: ${data.normalizer_rmsd}
+  datasets:
+    - ${data.dataset_name}
+  inference_only: True
+- _target_: fairchem.core.units.mlip_unit.mlip_unit.Task
+  name: stress
+  level: system
+  property: stress
+  out_spec:
+    dim: [1, 9]
+    dtype: float32
+  normalizer:
+    _target_: fairchem.core.modules.normalization.normalizer.Normalizer
+    mean: 0.0
+    rmsd: ${data.normalizer_rmsd}
+  datasets:
+    - ${data.dataset_name}
+  inference_only: True

uma_finetune/resources/configs/uma/finetune/data/uma_conserving_data_task_energy_force.yaml

@@ -0,0 +1,100 @@
+dataset_name: ??
+elem_refs: ??
+normalizer_rmsd: ??
+
+train_dataset:
+  splits:
+    train:
+      src: ??
+  format: ase_db
+  transforms:
+    common_transform:
+      dataset_name: ${data.dataset_name}
+    stress_reshape_transform:
+      dataset_name: ${data.dataset_name}
+
+val_dataset:
+  splits:
+    val:
+      src: ??
+  format: ase_db
+  transforms:
+    common_transform:
+      dataset_name: ${data.dataset_name}
+    stress_reshape_transform:
+      dataset_name: ${data.dataset_name}
+
+regress_stress: True
+pass_through_head_outputs: True
+
+heads:
+  efs:
+    module: fairchem.core.models.uma.escn_md.MLP_EFS_Head
+
+tasks_list:
+- _target_: fairchem.core.units.mlip_unit.mlip_unit.Task
+  name: energy
+  level: system
+  property: energy
+  loss_fn:
+    _target_: fairchem.core.modules.loss.DDPMTLoss
+    loss_fn:
+      _target_: fairchem.core.modules.loss.PerAtomMAELoss
+    coefficient: 20
+  out_spec:
+    dim: [1]
+    dtype: float32
+  normalizer:
+    _target_: fairchem.core.modules.normalization.normalizer.Normalizer
+    mean: 0.0
+    rmsd: ${data.normalizer_rmsd}
+  element_references:
+    _target_: fairchem.core.modules.normalization.element_references.ElementReferences
+    element_references:
+      _target_: torch.DoubleTensor
+      _args_:
+        - ${data.elem_refs}
+  datasets:
+    - ${data.dataset_name}
+  metrics:
+    - mae
+    - per_atom_mae
+- _target_: fairchem.core.units.mlip_unit.mlip_unit.Task
+  name: forces
+  level: atom
+  property: forces
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  loss_fn:
+    _target_: fairchem.core.modules.loss.DDPMTLoss
+    loss_fn:
+      _target_: fairchem.core.modules.loss.L2NormLoss
+    reduction: mean
+    coefficient: 2
+  out_spec:
+    dim: [3]
+    dtype: float32
+  normalizer:
+    _target_: fairchem.core.modules.normalization.normalizer.Normalizer
+    mean: 0.0
+    rmsd: ${data.normalizer_rmsd}
+  datasets:
+    - ${data.dataset_name}
+  metrics:
+    - mae
+    - cosine_similarity
+    - magnitude_error
+- _target_: fairchem.core.units.mlip_unit.mlip_unit.Task
+  name: stress
+  level: system
+  property: stress
+  out_spec:
+    dim: [1, 9]
+    dtype: float32
+  normalizer:
+    _target_: fairchem.core.modules.normalization.normalizer.Normalizer
+    mean: 0.0
+    rmsd: ${data.normalizer_rmsd}
+  datasets:
+    - ${data.dataset_name}
+  inference_only: True

uma_finetune/resources/configs/uma/finetune/data/uma_conserving_data_task_energy_force_stress.yaml

@@ -0,0 +1,107 @@
+dataset_name: ??
+elem_refs: ??
+normalizer_rmsd: ??
+
+train_dataset:
+  splits:
+    train:
+      src: ??
+  format: ase_db
+  transforms:
+    common_transform:
+      dataset_name: ${data.dataset_name}
+    stress_reshape_transform:
+      dataset_name: ${data.dataset_name}
+
+val_dataset:
+  splits:
+    val:
+      src: ??
+  format: ase_db
+  transforms:
+    common_transform:
+      dataset_name: ${data.dataset_name}
+    stress_reshape_transform:
+      dataset_name: ${data.dataset_name}
+
+regress_stress: True
+pass_through_head_outputs: True
+
+heads:
+  efs:
+    module: fairchem.core.models.uma.escn_md.MLP_EFS_Head
+
+tasks_list:
+- _target_: fairchem.core.units.mlip_unit.mlip_unit.Task
+  name: energy
+  level: system
+  property: energy
+  loss_fn:
+    _target_: fairchem.core.modules.loss.DDPMTLoss
+    loss_fn:
+      _target_: fairchem.core.modules.loss.PerAtomMAELoss
+    coefficient: 20
+  out_spec:
+    dim: [1]
+    dtype: float32
+  normalizer:
+    _target_: fairchem.core.modules.normalization.normalizer.Normalizer
+    mean: 0.0
+    rmsd: ${data.normalizer_rmsd}
+  element_references:
+    _target_: fairchem.core.modules.normalization.element_references.ElementReferences
+    element_references:
+      _target_: torch.DoubleTensor
+      _args_:
+        - ${data.elem_refs}
+  datasets:
+    - ${data.dataset_name}
+  metrics:
+    - mae
+    - per_atom_mae
+- _target_: fairchem.core.units.mlip_unit.mlip_unit.Task
+  name: forces
+  level: atom
+  property: forces
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  loss_fn:
+    _target_: fairchem.core.modules.loss.DDPMTLoss
+    loss_fn:
+      _target_: fairchem.core.modules.loss.L2NormLoss
+    reduction: mean
+    coefficient: 2
+  out_spec:
+    dim: [3]
+    dtype: float32
+  normalizer:
+    _target_: fairchem.core.modules.normalization.normalizer.Normalizer
+    mean: 0.0
+    rmsd: ${data.normalizer_rmsd}
+  datasets:
+    - ${data.dataset_name}
+  metrics:
+    - mae
+    - cosine_similarity
+    - magnitude_error
+- _target_: fairchem.core.units.mlip_unit.mlip_unit.Task
+  name: stress
+  level: system
+  property: stress
+  loss_fn:
+    _target_: fairchem.core.modules.loss.DDPMTLoss
+    loss_fn:
+      _target_: fairchem.core.modules.loss.MAELoss
+    reduction: mean
+    coefficient: 1
+  out_spec:
+    dim: [1, 9]
+    dtype: float32
+  normalizer:
+    _target_: fairchem.core.modules.normalization.normalizer.Normalizer
+    mean: 0.0
+    rmsd: ${data.normalizer_rmsd}
+  datasets:
+    - ${data.dataset_name}
+  metrics:
+    - mae

uma_finetune/resources/configs/uma/finetune/uma_sm_finetune_template.yaml

@@ -0,0 +1,118 @@
+defaults:
+  - data: ??
+  - _self_
+
+job:
+  device_type: CUDA
+  scheduler:
+    mode: LOCAL
+    ranks_per_node: 1
+    num_nodes: 1
+  debug: True
+  run_dir: /tmp/uma_finetune_runs/
+  run_name: uma_finetune
+  logger:
+    _target_: fairchem.core.common.logger.WandBSingletonLogger.init_wandb
+    _partial_: true
+    entity: example
+    project: uma_finetune
+
+# refers to a model name that can be retrieved from (huggingface)[https://huggingface.co/facebook/UMA]
+base_model_name: uma-s-1p1
+# 300 is effectively infinite neighbors, but if you are memory strained, using 100 is usually just as good
+max_neighbors: 300
+# choose to either run for integer number of epochs or steps, only 1 can be step, the other must be null
+epochs: 1
+steps: null
+# customize batch size depending on your system
+batch_size: 2
+lr: 4e-4
+weight_decay: 1e-3
+evaluate_every_n_steps: 100
+checkpoint_every_n_steps: 1000
+
+
+train_dataset:
+  _target_: fairchem.core.datasets.mt_concat_dataset.create_concat_dataset
+  dataset_configs:
+    DATASET_NAME: ${data.train_dataset}
+  combined_dataset_config: { sampling: {type: temperature, temperature: 1.0} }
+
+train_dataloader:
+  _target_: fairchem.core.components.common.dataloader_builder.get_dataloader
+  dataset: ${train_dataset}
+  batch_sampler_fn:
+    _target_: fairchem.core.common.data_parallel.BalancedBatchSampler
+    _partial_: True
+    batch_size: ${batch_size}
+    shuffle: True
+    seed: 0
+  num_workers: 0
+  collate_fn:
+    _target_: fairchem.core.units.mlip_unit.mlip_unit.mt_collater_adapter
+    tasks: ${data.tasks_list}
+
+val_dataset:
+  _target_: fairchem.core.datasets.mt_concat_dataset.create_concat_dataset
+  dataset_configs:
+    DATASET_NAME: ${data.val_dataset}
+  combined_dataset_config: { sampling: {type: temperature, temperature: 1.0} }
+
+eval_dataloader:
+  _target_: fairchem.core.components.common.dataloader_builder.get_dataloader
+  dataset: ${val_dataset}
+  batch_sampler_fn:
+    _target_: fairchem.core.common.data_parallel.BalancedBatchSampler
+    _partial_: True
+    batch_size: ${batch_size}
+    shuffle: False
+    seed: 0
+  num_workers: 0
+  collate_fn:
+    _target_: fairchem.core.units.mlip_unit.mlip_unit.mt_collater_adapter
+    tasks: ${data.tasks_list}
+
+runner:
+  _target_: fairchem.core.components.train.train_runner.TrainEvalRunner
+  train_dataloader: ${train_dataloader}
+  eval_dataloader: ${eval_dataloader}
+  train_eval_unit:
+    _target_: fairchem.core.units.mlip_unit.mlip_unit.MLIPTrainEvalUnit
+    job_config: ${job}
+    tasks: ${data.tasks_list}
+    model:
+      _target_: fairchem.core.units.mlip_unit.mlip_unit.initialize_finetuning_model
+      checkpoint_location:
+        _target_: fairchem.core.calculate.pretrained_mlip.pretrained_checkpoint_path_from_name
+        model_name: ${base_model_name}
+      overrides:
+        backbone:
+          otf_graph: True
+          max_neighbors: ${max_neighbors}
+          regress_stress: ${data.regress_stress}
+          always_use_pbc: False
+        pass_through_head_outputs: ${data.pass_through_head_outputs}
+      heads: ${data.heads}
+    optimizer_fn:
+      _target_: torch.optim.AdamW
+      _partial_: true
+      lr: ${lr}
+      weight_decay: ${weight_decay}
+    cosine_lr_scheduler_fn:
+      _target_: fairchem.core.units.mlip_unit.mlip_unit._get_consine_lr_scheduler
+      _partial_: true
+      warmup_factor: 0.2
+      warmup_epochs: 0.01
+      lr_min_factor: 0.01
+      epochs: ${epochs}
+      steps: ${steps}
+    print_every: 10
+    clip_grad_norm: 100
+  max_epochs: ${epochs}
+  max_steps: ${steps}
+  evaluate_every_n_steps: ${evaluate_every_n_steps}
+  callbacks:
+    - _target_: fairchem.core.components.callbacks.TrainCheckpointCallback
+      checkpoint_every_n_steps: ${checkpoint_every_n_steps}
+      max_saved_checkpoints: 5
+    - _target_: torchtnt.framework.callbacks.TQDMProgressBar

uma_finetune/training.py

@@ -0,0 +1,153 @@
+'''Run UMA fine-tuning with configurable backbone freezing.'''
+
+import json
+import math
+import subprocess
+import sys
+from pathlib import Path
+
+import yaml
+
+
+def initialize_model(mode='full', freeze_blocks=0, **kwargs):
+    '''Initialize UMA and apply the requested parameter-freezing policy.'''
+    from fairchem.core.units.mlip_unit.mlip_unit import initialize_finetuning_model
+
+    model = initialize_finetuning_model(**kwargs)
+
+    if mode not in ('full', 'head_only', 'partial'):
+        raise ValueError('mode must be full, head_only, or partial.')
+    if type(freeze_blocks) is not int or freeze_blocks < 0:
+        raise ValueError('freeze_blocks must be a nonnegative integer.')
+    if mode != 'partial' and freeze_blocks != 0:
+        raise ValueError('freeze_blocks is only used with mode=partial.')
+
+    if mode == 'partial':
+        blocks = model.backbone.blocks
+        if freeze_blocks > len(blocks):
+            raise ValueError(f'freeze_blocks exceeds the {len(blocks)} backbone blocks.')
+
+    # Keep the heads trainable; control backbone parameters separately.
+    for parameter in model.parameters():
+        parameter.requires_grad = True
+
+    if mode in ('head_only', 'partial'):
+        for parameter in model.backbone.parameters():
+            parameter.requires_grad = False
+
+    if mode == 'partial':
+        for block in blocks[freeze_blocks:]:
+            for parameter in block.parameters():
+                parameter.requires_grad = True
+        print(f'Trainable backbone blocks: {len(blocks) - freeze_blocks}/{len(blocks)}')
+
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total = sum(p.numel() for p in model.parameters())
+    print(f'Trainable parameters: {trainable:,}/{total:,}')
+    return model
+
+
+def finetune(
+    dataset,
+    output_dir,
+    *,
+    model=None,
+    mode='full',
+    freeze_blocks=0,
+    energy_coef=20,
+    force_coef=2,
+    epochs=30,
+    batch_size=4,
+    lr=1e-4,
+    device='cuda',
+    max_neighbors=300,
+    max_saved_checkpoints=2,
+):
+    '''
+    Fine-tune UMA using a directory returned by create_dataset().
+
+    Args:
+        dataset: Directory containing the generated training YAML and datasets.
+        output_dir: Root directory for FAIRChem run outputs.
+        model: Pretrained UMA model name; None keeps the generated YAML setting.
+        mode: 'full', 'head_only', or 'partial'.
+        freeze_blocks: Number of initial backbone blocks frozen in partial mode.
+            Other backbone parameters remain frozen; heads stay trainable.
+        energy_coef: Energy loss coefficient; None keeps the YAML setting.
+        force_coef: Force loss coefficient; None keeps the YAML setting.
+            Use None for an energy-only dataset.
+        epochs: Total training epochs.
+        batch_size: Training and validation batch size.
+        lr: Learning rate.
+        device: 'cuda' or 'cpu'.
+        max_neighbors: Maximum neighbors used by the model.
+        max_saved_checkpoints: Checkpoint retention limit.
+
+    Returns:
+        Absolute output root as a Path. FAIRChem creates its run folder inside.
+        This function starts a new run; it does not implement checkpoint resume.
+    '''
+    dataset = Path(dataset).expanduser().resolve()
+    output_dir = Path(output_dir).expanduser().resolve()
+    config = dataset / 'uma_sm_finetune_template.yaml'
+    if not config.is_file():
+        raise FileNotFoundError(config)
+    if mode not in ('full', 'head_only', 'partial'):
+        raise ValueError('mode must be full, head_only, or partial.')
+    if type(freeze_blocks) is not int or freeze_blocks < 0:
+        raise ValueError('freeze_blocks must be a nonnegative integer.')
+    if mode != 'partial' and freeze_blocks != 0:
+        raise ValueError('freeze_blocks is only used with mode=partial.')
+    device = device.upper()
+    if device not in ('CPU', 'CUDA'):
+        raise ValueError('device must be cpu or cuda.')
+    for name, value in [('epochs', epochs), ('batch_size', batch_size),
+                        ('max_neighbors', max_neighbors),
+                        ('max_saved_checkpoints', max_saved_checkpoints)]:
+        if type(value) is not int or value < 1:
+            raise ValueError(f'{name} must be a positive integer.')
+    if not math.isfinite(lr) or lr <= 0:
+        raise ValueError('lr must be finite and positive.')
+
+    # Invoke the installed FAIRChem entry point with this Python interpreter.
+    entry_point = '''from importlib.metadata import distribution
+next(ep for ep in distribution('fairchem-core').entry_points
+     if ep.group == 'console_scripts' and ep.name == 'fairchem').load()()
+'''
+    command = [
+        sys.executable, '-c', entry_point,
+        '-c', str(config),
+        f'job.device_type={device}',
+        '+job.scheduler.distributed_init_method=FILE',
+        f'job.run_dir={json.dumps(str(output_dir))}',
+        f'epochs={epochs}',
+        f'batch_size={batch_size}',
+        f'lr={lr}',
+        f'max_neighbors={max_neighbors}',
+        'runner.train_eval_unit.model._target_=uma_finetune.training.initialize_model',
+        f'+runner.train_eval_unit.model.mode={mode}',
+        f'+runner.train_eval_unit.model.freeze_blocks={freeze_blocks}',
+        f'runner.callbacks.0.max_saved_checkpoints={max_saved_checkpoints}',
+    ]
+    if model is not None:
+        command.append(f'base_model_name={json.dumps(model)}')
+
+    # Find loss positions by property rather than assuming their list order.
+    template = yaml.safe_load(config.read_text())
+    data_name = next(item['data'] for item in template['defaults']
+                     if isinstance(item, dict) and 'data' in item)
+    data = yaml.safe_load((dataset / 'data' / f'{data_name}.yaml').read_text())
+    for property_name, coefficient in [('energy', energy_coef), ('forces', force_coef)]:
+        if coefficient is None:
+            continue
+        if not math.isfinite(coefficient) or coefficient < 0:
+            raise ValueError(f'{property_name} coefficient must be finite and nonnegative.')
+        index = next((i for i, task in enumerate(data['tasks_list'])
+                      if task['property'] == property_name
+                      and not task.get('inference_only', False)), None)
+        if index is None:
+            raise ValueError(f'No training task for {property_name}; set its coefficient to None.')
+        command.append(f'data.tasks_list.{index}.loss_fn.coefficient={coefficient}')
+
+    subprocess.run(command, check=True, cwd=dataset)
+    return output_dir

uma_finetune_wrapper-0.1.0.dev2.dist-info/METADATA

@@ -0,0 +1,249 @@
+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).

uma_finetune_wrapper-0.1.0.dev2.dist-info/RECORD

@@ -0,0 +1,16 @@
+uma_finetune/__init__.py,sha256=42u-cwQRprTXU_pRLNQPKYrk9YDF9wqzXlhUzrsHAPs,170
+uma_finetune/dataset.py,sha256=s8EoSvie0pA7s-sma2x0474qWfSvgf_ObQ7EW7WF0Gw,2311
+uma_finetune/training.py,sha256=Ei0CkT2TY1ytSvD0zy1BILLjLDmP9DEvY_Bxha7z_tk,6505
+uma_finetune/resources/FAIRChem-LICENSE.md,sha256=PhCCeTiq6eWRZHciyDAGWpWzrych8NfwLzc4abQP6B0,1093
+uma_finetune/resources/NOTICE.md,sha256=BiU7jlt9w5979X4783bfwVzaKPH5pPEsZfEKH0tzME8,622
+uma_finetune/resources/configs/uma/finetune/uma_sm_finetune_template.yaml,sha256=OGQgSaSqvExd6rmpg01IjLx6b6o--ZW4gGDknzGjcq0,3800
+uma_finetune/resources/configs/uma/finetune/data/uma_conserving_data_task_energy.yaml,sha256=BR_7QQvBo_d06MANmvYXrWNAZ-HJKDMOBVOtZUmM_Jo,2041
+uma_finetune/resources/configs/uma/finetune/data/uma_conserving_data_task_energy_force.yaml,sha256=wa6zpC3i5KHOV8swE5dBK4QpdORmDXcaSZkKsqzZ8Q8,2308
+uma_finetune/resources/configs/uma/finetune/data/uma_conserving_data_task_energy_force_stress.yaml,sha256=5IpUVwGBDVR567U3cQZFDlB4UFXTPVETY7DXdWdZPdE,2471
+uma_finetune_wrapper-0.1.0.dev2.dist-info/licenses/LICENSE,sha256=V8T-HymNKX2rCymqfxOzsNgiAkMZ1kVRmG4w8W1Uk5A,1075
+uma_finetune_wrapper-0.1.0.dev2.dist-info/licenses/src/uma_finetune/resources/FAIRChem-LICENSE.md,sha256=PhCCeTiq6eWRZHciyDAGWpWzrych8NfwLzc4abQP6B0,1093
+uma_finetune_wrapper-0.1.0.dev2.dist-info/licenses/src/uma_finetune/resources/NOTICE.md,sha256=BiU7jlt9w5979X4783bfwVzaKPH5pPEsZfEKH0tzME8,622
+uma_finetune_wrapper-0.1.0.dev2.dist-info/METADATA,sha256=7p6Wja0uOaCUulQRiauccXGliC18RhoYwjsOOfZtG28,9332
+uma_finetune_wrapper-0.1.0.dev2.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+uma_finetune_wrapper-0.1.0.dev2.dist-info/top_level.txt,sha256=gKWWMdhx1KrESvpMdYjVKvpvoUJ5HNm52Ijc6AbdJzg,13
+uma_finetune_wrapper-0.1.0.dev2.dist-info/RECORD,,

uma_finetune_wrapper-0.1.0.dev2.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

uma_finetune_wrapper-0.1.0.dev2.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Timmo-hendrik Pukk
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

uma_finetune_wrapper-0.1.0.dev2.dist-info/licenses/src/uma_finetune/resources/FAIRChem-LICENSE.md

@@ -0,0 +1,23 @@
+# License
+
+MIT License
+
+Copyright (c) Meta, Inc. and its affiliates.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

uma_finetune_wrapper-0.1.0.dev2.dist-info/licenses/src/uma_finetune/resources/NOTICE.md

@@ -0,0 +1,13 @@
+# FAIRChem templates
+
+The four YAML files under configs/uma/finetune are unmodified copies from
+FAIRChem 2.22.0, supplied from the working source checkout for this project.
+Upstream project: https://github.com/facebookresearch/fairchem
+
+Copyright (c) Meta, Inc. and its affiliates.
+These templates are distributed under the MIT license reproduced in
+FAIRChem-LICENSE.md. This notice does not assign a license to the wrapper itself.
+
+The conversion script is supplied by the installed fairchem-core dependency;
+it is not copied into this package. Keep these resources matched to the pinned
+FAIRChem version when upgrading.

uma_finetune_wrapper-0.1.0.dev2.dist-info/top_level.txt

@@ -0,0 +1 @@
+uma_finetune