Metadata-Version: 2.4
Name: mint258-met
Version: 0.1.0
Summary: MET molecular embeddings and fine-tuning library.
Author: Mint258
License-Expression: MIT
Project-URL: Homepage, https://github.com/mint258/MET
Project-URL: Repository, https://github.com/mint258/MET
Project-URL: Documentation, https://github.com/mint258/MET
Project-URL: Dataset, https://huggingface.co/datasets/Mint258/MET-dateset
Project-URL: Models, https://huggingface.co/Mint258/MET-models
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: torch
Requires-Dist: torch-geometric
Requires-Dist: numpy
Requires-Dist: scipy
Requires-Dist: pandas
Requires-Dist: scikit-learn
Requires-Dist: sympy
Requires-Dist: rdkit
Requires-Dist: periodictable
Requires-Dist: tqdm
Requires-Dist: huggingface_hub
Provides-Extra: dev
Requires-Dist: build; extra == "dev"
Requires-Dist: twine; extra == "dev"

# 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

```bash
pip install mint258-met
```

For local development inside this repository:

```bash
pip install -e ./PyPI
```

## Core APIs

```python
from met import METEncoder, METFineTuner
```

### Atom Embeddings

```python
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

```python
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:

```python
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

```python
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

```python
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`
