FoMoNMR model releases

This directory stages the three FoMoNMR checkpoints selected for public release. All variants encode canonical, structured 1H + 13C resonance lists and return a 512-dimensional spectrum embedding. They operate on peak lists rather than raw FIDs or dense traces and do not generate molecular structures or SMILES.

FoMoNMR uses one token per resonance or multiplet, separate proton and carbon Fourier shift adapters, optional gated proton annotations, and a shared permutation-equivariant Transformer without positional encoding. Valid final peak states are mean-pooled to obtain the global embedding. The model accepts up to 60 proton and 60 carbon resonances per record.

Directory	MLflow run	Training stage	Recommended input
pretraining/	8462a988a9264f40943d37429826d3d5	shift-only-heavy pretraining	chemical shifts
posttraining/	52a53277e18a471b97caf84992c50b39	standard continued pretraining	rich proton annotations + shifts
posttraining_unimol2/	f89fff6cef6848c29425ab3b4ff72fe0	UniMol2-relational continued pretraining	rich proton annotations + shifts

Training provenance

Each release directory contains the exact final_*.yaml configuration used by its run. These files record the architecture, objectives, augmentation, optimizer and scheduler settings, early stopping, validation seed, and maximum step budget. Their hashes and the selected checkpoint steps are recorded in the corresponding release_manifest.json.

The remaining launcher settings were recovered from the three MLflow runs. All used a per-process batch size of 2,048, accumulation of 1, 10 data workers, seed 42, bf16-mixed precision, complete logging, and the MACCS probe at every validation. The configured maximum was 1,000,000 optimizer steps.

Variant	Training code revision	Validation/checkpoint interval	Selected step
Pretraining	0bcd896	10,000	720,000
Standard posttraining	c0d2c95	10,000	92,000
UniMol2-relational posttraining	615b835	1,000	123,000

Training used DDP on two H100 nodes, with one process and one GPU per node. With the per-process batch size of 2,048 and no gradient accumulation, the effective global batch size was 4,096. The two posttraining runs continued from the released pretraining lineage.

From a canonical record to an embedding

Inference follows the same model-independent boundary used throughout the project:

schema-v2 Parquet or CanonicalRecord
    -> CanonicalParquetDataset
    -> FoundationNMRProcessor
    -> padded H/C tensors and availability masks
    -> FoMoNMR
    -> 512-dimensional embedding

Run the examples from a checkout of the nmr_fomo repository with PYTHONPATH=scripts. The public held-out test set can be downloaded from the Hugging Face dataset repository and streamed without loading the whole Parquet into memory:

import torch
from huggingface_hub import hf_hub_download
from torch.utils.data import DataLoader

from data import CanonicalParquetDataset
from model import FoundationNMRProcessor
from model.FoMoNMR import FoMoNMR

test_path = hf_hub_download(
    repo_id="niccogreek/nmr-canonical-cleaned",
    repo_type="dataset",
    filename="test_benchmark.parquet",
)

dataset = CanonicalParquetDataset(test_path)
loader = DataLoader(
    dataset,
    batch_size=8,
    num_workers=0,
    collate_fn=FoundationNMRProcessor(),
)
batch = next(iter(loader))

# For shift-only inference with the pretraining release, hide optional fields:
# batch["h"]["availability"].zero_()
# batch["h"]["j_mask"].zero_()

# Use "pretraining/fomonmr-pretrain.ckpt" for the shift-only release, or
# "posttraining_unimol2/fomonmr-posttrain-unimol2.ckpt" for the relational one.
model = FoMoNMR.load_from_checkpoint(
    "posttraining/fomonmr-posttrain.ckpt",
    map_location="cpu",
    weights_only=True,
)
model.eval()

with torch.inference_mode():
    embeddings, peak_states, peak_mask = model(batch)

print(embeddings.shape)  # torch.Size([8, 512])

The processor pads each modality to 60 resonances and constructs masks that keep padding, unavailable annotations and real values distinct. The two commented lines reproduce the shift-only input policy used by the project's FoMoNMR benchmark embedder.

A small local input

For a small local dataset, records can be constructed directly with the public dataclasses. Optional annotations use None when the measurement is unavailable; an empty J list instead means that the source explicitly reports no numerical coupling values.

from data import CanonicalRecord, CarbonPeak, ProtonPeak
from model import FoundationNMRProcessor

record = CanonicalRecord(
    record_id="example-001",
    source="local",
    h_nmr_peaks=(
        ProtonPeak(
            shift=1.20,
            integration=3,
            multiplicity_raw="t",
            multiplicity="t",
            j_values=(7.2,),
            range_min=1.18,
            range_max=1.22,
            range_half_span=0.02,
        ),
        ProtonPeak(
            shift=3.65,
            integration=2,
            multiplicity_raw="q",
            multiplicity="q",
            j_values=(7.2,),
        ),
    ),
    c_nmr_peaks=(
        CarbonPeak(shift=18.3),
        CarbonPeak(shift=58.1),
    ),
)

batch = FoundationNMRProcessor()([record])

The same processor accepts a list of dataclasses or validated JSON-like mappings. Larger local collections should be converted once to canonical schema-v2 Parquet and opened with CanonicalParquetDataset; the reader does not parse or guess fields from arbitrary raw-source formats. Source-specific conversion examples are documented in the scripts/data/canonicalize directory.

Each release contains:

• fomonmr-*.ckpt: the primary Lightning weights-only checkpoint;
• model/: the equivalent MLflow export, retained for provenance but not used by the documented loading path;
• config.yaml: the exact maintained training configuration;
• release_manifest.json: run, code, DVC and artifact provenance;
• README.md: the variant-specific model card.

The release checkpoint contains tensor-for-tensor identical model weights to the selected MLflow export, without optimizer, scheduler, loop or callback state. The common source revision is db8cdda, and the canonical training collection is published as niccogreek/nmr-canonical-cleaned. The complete implementation and scientific documentation are available in the nmr_fomo GitHub repository and its FoMoNMR model documentation.