mint258-met 0.1.0

MET molecular embeddings and fine-tuning library.

mint258-met

mint258-met packages the fine-tuning-facing part of MET as an installable Python library.

The package focuses on:

• loading the released MET pretrained encoder from Hugging Face
• producing atom-level embeddings with shape [N, D]
• fine-tuning on user CSV or XYZ data with built-in dataloaders
• exposing both a default prediction head and custom-head injection
• controlling which encoder blocks remain trainable during fine-tuning

Install

pip install mint258-met

For local development inside this repository:

pip install -e ./PyPI

Core APIs

from met import METEncoder, METFineTuner

Atom Embeddings

encoder = METEncoder.from_pretrained()
atom_embeddings = encoder.embed_smiles("CCO", num_conformers=4)
print(atom_embeddings.shape)  # [N, D]

METEncoder returns atom-level embeddings with shape [N, D], where N is the number of atoms and D is the MET embedding dimension from the pretrained checkpoint.

Supported encoder convenience methods:

• embed_smiles(smiles, ...)
• embed_xyz(path, ...)
• embed_smiles_batch(smiles_list, ...)
• embed_data(batch_or_graph)

Fine-Tuning

model = METFineTuner.from_pretrained(
    output_dim=1,
    hidden_dim=256,
    layers=2,
    dropout=0.1,
    pooling_time="after_last_layer",
    pooling_method="attention",
    normalization="layernorm",
)

model.set_trainable_encoder_blocks(trainable_blocks=1)

history = model.fit_csv(
    csv_path="my_dataset.csv",
    target_columns=["target"],
    smiles_column="smiles",
    num_conformers=8,
    epochs=5,
    batch_size=8,
)

Default Prediction Head

The default head is controlled by the following arguments in METFineTuner.from_pretrained(...) or build_default_head(...):

• output_dim: number of downstream targets
• hidden_dim: hidden width used by the default head
• layers: number of feed-forward blocks
• dropout: dropout used inside the default head
• pooling_time: "before_head" or "after_last_layer"
• pooling_method: "attention", "mean", "sum", or "max"
• normalization: None, "layernorm", or "batchnorm"

pooling_time="after_last_layer" means node embeddings are transformed first and pooled at the end.
pooling_time="before_head" means raw node embeddings are pooled first and the graph representation is processed afterward.

Choosing Trainable Encoder Layers

METFineTuner exposes three layer-control styles:

• freeze_encoder(): train only the downstream head
• unfreeze_encoder(): fine-tune the full backbone
• set_trainable_encoder_blocks(trainable_blocks=...): unfreeze the last k interaction blocks plus optional projection layers
• set_trainable_modules([...]): explicitly unfreeze named encoder modules

You can inspect available module names with:

print(model.list_encoder_modules())

Custom Prediction Heads

If you want to replace the default head, pass a custom torch.nn.Module to METFineTuner.from_pretrained(head=...) or later call replace_head(...).

The custom head should accept:

• node_embeddings: shape [batch, max_nodes, dim]
• mask: shape [batch, max_nodes]

and return:

• predictions of shape [batch, output_dim]

Built-In Fine-Tuning Helpers

For downstream work the package provides:

• fit_csv(...)
• fit_xyz(...)
• fit_dataset(...)
• predict_csv(...)
• predict_xyz(...)
• predict_smiles(...)
• evaluate_csv(...)
• evaluate_xyz(...)
• evaluate_dataset(...)

CSV Example

from met import METFineTuner

model = METFineTuner.from_pretrained(
    output_dim=2,
    hidden_dim=256,
    layers=3,
    dropout=0.1,
    pooling_time="after_last_layer",
    pooling_method="attention",
    normalization="layernorm",
)

model.set_trainable_encoder_blocks(trainable_blocks=2)

history = model.fit_csv(
    csv_path="downstream.csv",
    smiles_column="smiles",
    target_columns=["property_a", "property_b"],
    num_conformers=8,
    conformer_seed=2024,
    use_uff_optimization=True,
    epochs=10,
    batch_size=16,
    learning_rate=1e-4,
)

XYZ Example

from met import METFineTuner

model = METFineTuner.from_pretrained(output_dim=1)
model.set_trainable_encoder_blocks(trainable_blocks=1)

history = model.fit_xyz(
    root="data/QM7/train_database_300",
    target_columns=["atomization_energy"],
    epochs=5,
    batch_size=8,
)

Hugging Face Assets

• Dataset: https://huggingface.co/datasets/Mint258/MET-dateset
• Models: https://huggingface.co/Mint258/MET-models

Project Links

• GitHub: https://github.com/mint258/MET
• Paper DOI: https://doi.org/10.1039/D5ME00173K