MolecularDiffusion.modules.models.chefnmr.embedders

NMR spectrum tokenizers + encoder for ChefNMR (MIT, (c) 2025 Ziyu Xiong).

Upstream: src/model/modules/embedders.py (conv tokenizer modified from MarklandGroup/NMR2Struct).

Ported verbatim except that upstream’s three einops.rearrange calls are inlined as plain reshape/permute – einops is not installed in this environment and these three sites are the only thing that wanted it:

rearrange(t, "b n (h d) -> b h n d")  ->  t.reshape(b, n, h, d).permute(0, 2, 1, 3)
rearrange(o, "b h n d -> b n (h d)")  ->  o.permute(0, 2, 1, 3).reshape(b, n, h * d)

The ``embed`` tokenizer only works on a binary vector. _embed_spectrum does x.long() * arange(1, L+1) into an nn.Embedding(L+1, D, padding_idx=0): a bin holding 0 maps to the pad row, a bin holding 1 maps to that bin’s own row, and any other value indexes a different bin’s row. The 80-bin 13C grid is stored as strictly {0.0, 1.0}, and scripts/convert_dataset.py asserts that rather than assuming it.

Classes

Attention	Unmasked multi-head self-attention over spectrum tokens.
AttnPoolToken	CLS-token attention pooling to a fixed-size vector.
FeedForward
NMRSpectraEmbedder	Concatenated 1H/13C spectra -> one (B, output_dim) vector.
PreNorm
SpectraTokenizerConv1D	Conv1d/ReLU/MaxPool stack -> (B, T, D) tokens.
SpectraTokenizerPatch1D	Split a 1-D spectrum into fixed patches and project them.
TransformerEncoder

Module Contents

class MolecularDiffusion.modules.models.chefnmr.embedders.Attention(dim: int, heads: int, dim_head: int | None = None)

Bases: torch.nn.Module

Unmasked multi-head self-attention over spectrum tokens.

forward(x: torch.Tensor) → torch.Tensor

attn

dim_head = None

heads

scale

to_out

to_qkv

class MolecularDiffusion.modules.models.chefnmr.embedders.AttnPoolToken(dim: int, out_dim: int, heads: int = 8, dim_head: int = 64, dropout: float = 0.1)

Bases: torch.nn.Module

CLS-token attention pooling to a fixed-size vector.

forward(x: torch.Tensor) → torch.Tensor

attn

cls

proj

class MolecularDiffusion.modules.models.chefnmr.embedders.FeedForward(dim: int, hidden_dim: int, dropout: float)

Bases: torch.nn.Module

forward(x: torch.Tensor) → torch.Tensor

net

class MolecularDiffusion.modules.models.chefnmr.embedders.NMRSpectraEmbedder(*, use_hnmr: bool = True, use_cnmr: bool = False, hnmr_dim: int = 10000, cnmr_dim: int = 10000, hidden_dim: int = 256, output_dim: int = 768, dropout: float = 0.1, pooling: str = 'flatten', tokenizer_args: dict = None, transformer_args: dict = None)

Bases: torch.nn.Module

Concatenated 1H/13C spectra -> one (B, output_dim) vector.

forward(x: torch.Tensor) → torch.Tensor

pos_encode(x: torch.Tensor) → torch.Tensor

c_tokenizer

cnmr_dim = 10000

dropout

h_tokenizer

hidden_dim = 256

hnmr_dim = 10000

pooling = 'flatten'

pos_enc_mode

use_c_mask_token

use_cnmr = False

use_h_mask_token

use_hnmr = True

class MolecularDiffusion.modules.models.chefnmr.embedders.PreNorm(dim: int, fn: torch.nn.Module)

Bases: torch.nn.Module

forward(x: torch.Tensor) → torch.Tensor

fn

norm

class MolecularDiffusion.modules.models.chefnmr.embedders.SpectraTokenizerConv1D(input_size: int, hidden_dim: int, pool_sizes: list, kernel_sizes: list, out_channels: list)

Bases: torch.nn.Module

Conv1d/ReLU/MaxPool stack -> (B, T, D) tokens.

forward(x: torch.Tensor) → torch.Tensor

static num_tokens(input_size: int, kernel_sizes: list, pool_sizes: list) → int

conv_blocks

linear_after_conv

class MolecularDiffusion.modules.models.chefnmr.embedders.SpectraTokenizerPatch1D(patch_size: int, stride: int, hidden_dim: int)

Bases: torch.nn.Module

Split a 1-D spectrum into fixed patches and project them.

forward(x: torch.Tensor) → torch.Tensor

static num_tokens(input_size: int, patch_size: int, stride: int) → int

patch_size

proj

stride

class MolecularDiffusion.modules.models.chefnmr.embedders.TransformerEncoder(dim: int, depth: int, heads: int, dim_head: int, mlp_ratio: int, dropout: float)

Bases: torch.nn.Module

forward(x: torch.Tensor) → torch.Tensor

final_norm

layers