diff --git a/README.md b/README.md
index a73941b..1c7cf15 100644
--- a/README.md
+++ b/README.md
@@ -56,6 +56,10 @@ See the [List of Supported Models](#list-of-supported-models) section for all av
| Model | Required Packages |
|-------|-------------------|
| HFPretrainedMolecularEncoder | transformers |
+| HFPretrainedMolecularGenerator | transformers |
+| HFPretrainedMolecularGenerator (MolGen) | transformers, [selfies](https://github.com/aspuru-guzik-group/selfies) |
+| HFPretrainedMolecularGenerator (Molexar) | transformers, [fragment-selfies](https://github.com/fairydance/Fragment-SELFIES), [molexar](https://github.com/fairydance/Molexar) |
+| HFPretrainedMolecularGenerator (SAFE-GPT) | transformers, [safe-mol](https://github.com/datamol-io/safe) |
| BFGNNMolecularPredictor | torch-scatter |
| GRINMolecularPredictor | torch-scatter |
| GRINMolecularPredictor (if enable `repetition_augmentation=True`) | CombineMols |
@@ -66,6 +70,23 @@ See the [List of Supported Models](#list-of-supported-models) section for all av
**For models that require `transformers`:** `pip install transformers`
+**For MolGen (`selfies`):** `pip install "selfies>=2.1"`. Source: [aspuru-guzik-group/selfies](https://github.com/aspuru-guzik-group/selfies).
+
+**For Molexar:** `pip install fragment-selfies loguru` ([Fragment-SELFIES](https://github.com/fairydance/Fragment-SELFIES)) and `pip install git+https://github.com/fairydance/Molexar.git` ([Molexar](https://github.com/fairydance/Molexar)). Molexar itself requires `transformers>=5.8`.
+
+**For SAFE-GPT:** `pip install safe-mol` ([SAFE](https://github.com/datamol-io/safe)).
+
+```python
+from torch_molecule import HFPretrainedMolecularGenerator
+
+model = HFPretrainedMolecularGenerator(
+ repo_id="datamol-io/safe-gpt",
+)
+model.fit()
+print(model.generate(n_samples=5))
+print(model.generate(n_samples=5, scaffold="c1ccccc1"))
+```
+
## Usage
> More examples can be found in the `examples` and `tests` folders.
@@ -107,15 +128,21 @@ assert molecular_data.target is None
### Fit a Model
-After preparing the dataset, we can easily fit a model similar to how we use sklearn (actually, the coding is even simpler than sklearn, as we still need to do feature engineering in sklearn to convert molecule SMILES into vectors):
+After preparing the dataset, split it, then fit a model with an sklearn-style API (no extra SMILES featurization is required):
```python
+from torch_molecule.datasets import load_qm9
from torch_molecule import GREAMolecularPredictor
-split = int(0.8 * len(smiles_list))
+data = load_qm9(local_dir='torchmol_data')
+# "random" | "scaffold" | "butina" | "size"
+# scaffold: unseen Bemis-Murcko scaffolds; butina: Tanimoto clusters; size: heavy-atom count
+# Split the full dataset. subsample() is only for local debugging / CI — do not
+# shrink QM9 (or any benchmark) just to make Butina cheaper.
+train, val = data.train_test_split(test_size=0.2, method="scaffold", seed=42)
grea = GREAMolecularPredictor(
- num_task=num_task,
+ num_task=1,
task_type="regression",
evaluate_higher_better=False,
verbose="progress_bar" #or "print_statement" recommended for jupyter notebooks, or "none"
@@ -123,10 +150,10 @@ grea = GREAMolecularPredictor(
# Fit with automatic hyperparameter tuning with 10 attempts, or implement .fit() with the default/manual hyperparameters
grea.autofit(
- X_train=smiles_list[:split],
- y_train=property_np_array[:split],
- X_val=smiles_list[split:],
- y_val=property_np_array[split:],
+ X_train=train.data,
+ y_train=train.target,
+ X_val=val.data,
+ y_val=val.target,
n_trials=10,
)
```
@@ -196,6 +223,7 @@ new_model.load_from_local("qm9_grea.pt")
| JTVAE | [Junction Tree Variational Autoencoder for Molecular Graph Generation. ICML 2018.](https://proceedings.mlr.press/v80/jin18a) |
| GraphGA | [A Graph-Based Genetic Algorithm and Its Application to the Multiobjective Evolution of Median Molecules. Journal of Chemical Information and Computer Sciences 2004](https://pubs.acs.org/doi/10.1021/ci034290p) |
| LSTM (SMILES) | [Long short-term memory (Neural Computation 1997)](https://ieeexplore.ieee.org/abstract/document/6795963) based on SMILES strings |
+| Pretrained | [NovoMolGen](https://huggingface.co/chandar-lab/NovoMolGen_32M_SMILES_BPE): Causal LM pretrained on ZINC-22 for de novo SMILES generation.
[MolGen-large](https://huggingface.co/zjunlp/MolGen-large): Seq2Seq SELFIES generator with high chemical validity.
[MolGen-large-opt](https://huggingface.co/zjunlp/MolGen-large-opt): MolGen-large fine-tuned for QED / p-logP optimization.
[Molexar-10M-base](https://huggingface.co/fairydance/molexar-10m-base): Fragment-SELFIES de novo and fragment-constrained generation.
[Molexar-10M-omni](https://huggingface.co/fairydance/molexar-10m-omni): Multi-condition Molexar model for property-guided generation.
[SAFE-GPT](https://huggingface.co/datamol-io/safe-gpt): GPT-2 causal LM pretrained on SAFE strings for de novo generation and scaffold-prefix completion. |
### Representation Models
diff --git a/docs/source/install.rst b/docs/source/install.rst
index e37d2f2..b20ab38 100644
--- a/docs/source/install.rst
+++ b/docs/source/install.rst
@@ -66,12 +66,28 @@ Additional Packages
Some models require extra libraries. Install these packages if you use the corresponding model:
-+------------------------------+-------------------+
-| Model | Required Package |
-+==============================+===================+
-| HFPretrainedMolecularEncoder | transformers |
-+------------------------------+-------------------+
-| BFGNNMolecularPredictor | torch-scatter |
-+------------------------------+-------------------+
-| GRINMolecularPredictor | torch-scatter |
-+------------------------------+-------------------+
++----------------------------------------------+----------------------------------------------+
+| Model | Required Package |
++==============================================+==============================================+
+| HFPretrainedMolecularEncoder | transformers |
++----------------------------------------------+----------------------------------------------+
+| HFPretrainedMolecularGenerator | transformers |
++----------------------------------------------+----------------------------------------------+
+| HFPretrainedMolecularGenerator (MolGen) | transformers, selfies |
++----------------------------------------------+----------------------------------------------+
+| HFPretrainedMolecularGenerator (Molexar) | transformers, fragment-selfies, molexar |
++----------------------------------------------+----------------------------------------------+
+| HFPretrainedMolecularGenerator (SAFE-GPT) | transformers, safe-mol |
++----------------------------------------------+----------------------------------------------+
+| BFGNNMolecularPredictor | torch-scatter |
++----------------------------------------------+----------------------------------------------+
+| GRINMolecularPredictor | torch-scatter |
++----------------------------------------------+----------------------------------------------+
+
+**For models that require** ``transformers``: ``pip install transformers``
+
+**For MolGen** (``selfies``): ``pip install "selfies>=2.1"``. Source: `aspuru-guzik-group/selfies `_.
+
+**For Molexar:** ``pip install fragment-selfies loguru`` (`Fragment-SELFIES `_) and ``pip install git+https://github.com/fairydance/Molexar.git`` (`Molexar `_). Molexar itself requires ``transformers>=5.8``.
+
+**For SAFE-GPT:** ``pip install safe-mol`` (`SAFE `_).
diff --git a/pyproject.toml b/pyproject.toml
index af1055a..b1599a7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -57,6 +57,9 @@ include-package-data = true
[tool.pytest.ini_options]
addopts = "--verbose"
testpaths = ["tests"]
+markers = [
+ "integration: tests that download models or require network access",
+]
[project.optional-dependencies]
dev = [
diff --git a/tests/datasets/test_split.py b/tests/datasets/test_split.py
new file mode 100644
index 0000000..1616c18
--- /dev/null
+++ b/tests/datasets/test_split.py
@@ -0,0 +1,327 @@
+import numpy as np
+import pytest
+from rdkit import Chem
+from rdkit.Chem.Scaffolds import MurckoScaffold
+
+from torch_molecule.datasets import SMILESDataset, subsample, train_test_split
+from torch_molecule.datasets.split import _BUTINA_OOM_MESSAGE
+
+
+def _benzene_family():
+ return [
+ "c1ccccc1",
+ "c1ccccc1O",
+ "c1ccccc1N",
+ "c1ccccc1C",
+ "c1ccccc1Cl",
+ "c1ccccc1F",
+ "Cc1ccccc1O",
+ "Nc1ccccc1O",
+ ]
+
+
+def _other_molecules():
+ return [
+ "CCO",
+ "CCN",
+ "CCC",
+ "C1CCCCC1",
+ "n1ccccc1",
+ "C1CCNCC1",
+ "CC(=O)O",
+ "CC(C)O",
+ ]
+
+
+def _labeled_dataset():
+ smiles = _benzene_family() + _other_molecules()
+ y = np.arange(len(smiles), dtype=np.float32).reshape(-1, 1)
+ return SMILESDataset(data=smiles, target=y)
+
+
+def _scaffold_smiles(smiles: str) -> str:
+ mol = Chem.MolFromSmiles(smiles)
+ scaffold = MurckoScaffold.GetScaffoldForMol(mol)
+ if scaffold is None or scaffold.GetNumAtoms() == 0:
+ return Chem.MolToSmiles(mol)
+ return Chem.MolToSmiles(scaffold)
+
+
+def test_random_split_reproducible():
+ data = _labeled_dataset()
+ train_a, test_a = train_test_split(data, test_size=0.25, method="random", seed=42)
+ train_b, test_b = train_test_split(data, test_size=0.25, method="random", seed=42)
+ assert train_a.data == train_b.data
+ assert test_a.data == test_b.data
+ np.testing.assert_array_equal(train_a.target, train_b.target)
+
+
+def test_random_split_different_seeds():
+ data = _labeled_dataset()
+ train_a, _ = train_test_split(data, test_size=0.25, method="random", seed=0)
+ train_b, _ = train_test_split(data, test_size=0.25, method="random", seed=1)
+ assert train_a.data != train_b.data
+
+
+def test_random_split_ratio_and_coverage():
+ data = _labeled_dataset()
+ train, holdout = train_test_split(data, test_size=0.25, method="random", seed=42)
+ n = len(data.data)
+ assert len(train.data) + len(holdout.data) == n
+ assert abs(len(holdout.data) / n - 0.25) < 1e-9
+ assert set(train.data).isdisjoint(holdout.data)
+ assert set(train.data) | set(holdout.data) == set(data.data)
+
+
+def test_random_preserves_multitask_target():
+ smiles = _benzene_family() + _other_molecules()
+ y = np.column_stack(
+ [
+ np.arange(len(smiles), dtype=np.float32),
+ np.arange(len(smiles), dtype=np.float32) * 2,
+ ]
+ )
+ data = SMILESDataset(data=smiles, target=y)
+ train, holdout = train_test_split(data, test_size=0.2, method="random", seed=7)
+ assert train.target.shape[1] == 2
+ assert holdout.target.shape[1] == 2
+ assert train.target.shape[0] == len(train.data)
+
+
+def test_split_unlabeled_dataset():
+ data = SMILESDataset(data=_benzene_family() + _other_molecules(), target=None)
+ train, holdout = train_test_split(data, test_size=0.2, method="random", seed=1)
+ assert train.target is None
+ assert holdout.target is None
+ assert len(train.data) + len(holdout.data) == len(data.data)
+
+
+def test_scaffold_no_leakage():
+ data = _labeled_dataset()
+ train, holdout = train_test_split(data, test_size=0.3, method="scaffold", seed=42)
+ train_scaffolds = {_scaffold_smiles(s) for s in train.data}
+ holdout_scaffolds = {_scaffold_smiles(s) for s in holdout.data}
+ assert train_scaffolds.isdisjoint(holdout_scaffolds)
+ assert len(train.data) + len(holdout.data) == len(data.data)
+ assert len(set(train.data) & set(holdout.data)) == 0
+
+
+def test_scaffold_keeps_benzene_family_together():
+ data = _labeled_dataset()
+ train, holdout = train_test_split(data, test_size=0.3, method="scaffold")
+ benzene_scaffold = _scaffold_smiles("c1ccccc1")
+ train_has = any(_scaffold_smiles(s) == benzene_scaffold for s in train.data)
+ holdout_has = any(_scaffold_smiles(s) == benzene_scaffold for s in holdout.data)
+ assert train_has ^ holdout_has
+
+
+def test_scaffold_acyclic_molecules_do_not_crash():
+ data = SMILESDataset(
+ data=["CCO", "CCN", "CCC", "CC", "C"],
+ target=np.arange(5).reshape(-1, 1),
+ )
+ train, holdout = train_test_split(data, test_size=0.4, method="scaffold")
+ assert len(train.data) >= 1
+ assert len(holdout.data) >= 1
+
+
+def test_scaffold_invalid_smiles_raises():
+ data = SMILESDataset(data=["CCO", "not_a_smiles"], target=None)
+ with pytest.raises(ValueError, match="Invalid SMILES"):
+ train_test_split(data, method="scaffold")
+
+
+def test_unknown_method_raises():
+ data = _labeled_dataset()
+ with pytest.raises(ValueError, match="Unknown split method"):
+ train_test_split(data, method="kmeans")
+
+
+def test_subsample_reproducible_and_size():
+ data = _labeled_dataset()
+ a = subsample(data, n=5, seed=0)
+ b = subsample(data, n=5, seed=0)
+ c = subsample(data, n=5, seed=1)
+ assert len(a.data) == 5
+ assert a.data == b.data
+ assert a.data != c.data
+ assert a.target.shape == (5, 1)
+
+
+def test_subsample_and_split_methods_on_dataset():
+ data = _labeled_dataset()
+ small = data.subsample(n=10, seed=3)
+ assert len(small.data) == 10
+ train, holdout = small.train_test_split(test_size=0.3, method="random", seed=4)
+ assert len(train.data) + len(holdout.data) == 10
+ butina_train, butina_hold = small.train_test_split(
+ test_size=0.3, method="butina", similarity_cutoff=0.4
+ )
+ assert len(butina_train.data) + len(butina_hold.data) == 10
+ size_train, size_hold = small.train_test_split(test_size=0.3, method="size")
+ assert len(size_train.data) + len(size_hold.data) == 10
+
+
+def test_subsample_rejects_too_large_n():
+ data = _labeled_dataset()
+ with pytest.raises(ValueError, match="larger than the dataset size"):
+ data.subsample(n=len(data.data) + 1)
+
+
+def test_target_row_mismatch_raises():
+ data = SMILESDataset(data=["CCO", "CCC"], target=np.array([[1.0]]))
+ with pytest.raises(ValueError, match="target has"):
+ train_test_split(data, method="random")
+
+
+def _rdkit_butina_clusters(smiles, similarity_cutoff=0.65):
+ from rdkit.ML.Cluster import Butina
+ from rdkit.Chem import rdFingerprintGenerator
+ from rdkit import DataStructs as RDS
+
+ gen = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=2048)
+ fps = [gen.GetFingerprint(Chem.MolFromSmiles(s)) for s in smiles]
+ dists = []
+ for i in range(1, len(fps)):
+ sims = RDS.BulkTanimotoSimilarity(fps[i], fps[:i])
+ dists.extend([1.0 - x for x in sims])
+ return Butina.ClusterData(
+ dists, len(fps), 1.0 - similarity_cutoff, isDistData=True
+ )
+
+
+def _cluster_membership(clusters):
+ return {frozenset(cluster) for cluster in clusters}
+
+
+def test_butina_matches_rdkit_clusterdata_membership():
+ from torch_molecule.datasets.split import _butina_clusters
+
+ smiles = _benzene_family() + _other_molecules()
+ for cutoff in (0.3, 0.4, 0.65):
+ ours = _butina_clusters(smiles, similarity_cutoff=cutoff)
+ rdkit = _rdkit_butina_clusters(smiles, similarity_cutoff=cutoff)
+ assert _cluster_membership(ours) == _cluster_membership(rdkit)
+ for cluster, ref in zip(ours, rdkit):
+ assert cluster[0] == ref[0]
+
+
+def test_butina_no_cluster_leakage():
+ from torch_molecule.datasets.split import _butina_clusters
+
+ data = _labeled_dataset()
+ cutoff = 0.4
+ train, holdout = train_test_split(
+ data, test_size=0.3, method="butina", similarity_cutoff=cutoff
+ )
+ smiles_to_idx = {s: i for i, s in enumerate(data.data)}
+ clusters = _butina_clusters(data.data, similarity_cutoff=cutoff)
+ train_idx = {smiles_to_idx[s] for s in train.data}
+ holdout_idx = {smiles_to_idx[s] for s in holdout.data}
+ assert train_idx.isdisjoint(holdout_idx)
+ assert train_idx | holdout_idx == set(range(len(data.data)))
+ for cluster in clusters:
+ members = set(cluster)
+ assert members <= train_idx or members <= holdout_idx
+
+
+def test_butina_ignores_seed_and_is_reproducible():
+ data = _labeled_dataset()
+ a_train, a_hold = train_test_split(
+ data, test_size=0.3, method="butina", seed=0, similarity_cutoff=0.4
+ )
+ b_train, b_hold = train_test_split(
+ data, test_size=0.3, method="butina", seed=1, similarity_cutoff=0.4
+ )
+ assert a_train.data == b_train.data
+ assert a_hold.data == b_hold.data
+
+
+def test_butina_invalid_smiles_raises():
+ data = SMILESDataset(data=["CCO", "not_a_smiles"], target=None)
+ with pytest.raises(ValueError, match="Invalid SMILES"):
+ train_test_split(data, method="butina")
+
+
+def test_butina_rejects_bad_cutoff():
+ data = _labeled_dataset()
+ with pytest.raises(ValueError, match="similarity_cutoff"):
+ train_test_split(data, method="butina", similarity_cutoff=0.0)
+ with pytest.raises(ValueError, match="similarity_cutoff"):
+ train_test_split(data, method="butina", similarity_cutoff=1.5)
+
+
+def _heavy_atoms(smiles: str) -> int:
+ return Chem.MolFromSmiles(smiles).GetNumHeavyAtoms()
+
+
+def test_size_split_holds_out_larger_molecules():
+ smiles = ["C", "CC", "CCC", "CCCC", "c1ccccc1", "c1ccccc1c1ccccc1"]
+ y = np.arange(len(smiles), dtype=np.float32).reshape(-1, 1)
+ data = SMILESDataset(data=smiles, target=y)
+ train, holdout = train_test_split(
+ data, test_size=1 / 3, method="size", direction="small_to_large"
+ )
+ assert len(holdout.data) == 2
+ assert max(_heavy_atoms(s) for s in train.data) <= min(
+ _heavy_atoms(s) for s in holdout.data
+ )
+ assert set(train.data) | set(holdout.data) == set(smiles)
+
+
+def test_size_split_large_to_small_holds_out_smaller_molecules():
+ smiles = ["C", "CC", "CCC", "CCCC", "c1ccccc1", "c1ccccc1c1ccccc1"]
+ data = SMILESDataset(data=smiles, target=None)
+ train, holdout = train_test_split(
+ data, test_size=1 / 3, method="size", direction="large_to_small"
+ )
+ assert max(_heavy_atoms(s) for s in holdout.data) <= min(
+ _heavy_atoms(s) for s in train.data
+ )
+
+
+def test_size_split_sizeshiftreg_protocol():
+ smiles = [
+ "C",
+ "CC",
+ "CCC",
+ "CCCC",
+ "CCCCC",
+ "CCCCCC",
+ "CCCCCCC",
+ "CCCCCCCC",
+ "CCCCCCCCC",
+ "c1ccccc1",
+ ]
+ data = SMILESDataset(data=smiles, target=None)
+ train, holdout = train_test_split(
+ data, test_size=0.2, method="size", mode="sizeshiftreg"
+ )
+ n = len(smiles)
+ assert len(train.data) == int(round(0.5 * n))
+ assert len(holdout.data) == int(round(0.1 * n))
+ assert len(train.data) + len(holdout.data) < n
+ assert max(_heavy_atoms(s) for s in train.data) <= min(
+ _heavy_atoms(s) for s in holdout.data
+ )
+
+
+def test_size_invalid_smiles_raises():
+ data = SMILESDataset(data=["CCO", "not_a_smiles"], target=None)
+ with pytest.raises(ValueError, match="Invalid SMILES"):
+ train_test_split(data, method="size")
+
+
+def test_butina_oom_does_not_suggest_subsample(monkeypatch):
+ data = _labeled_dataset()
+
+ def _boom(*args, **kwargs):
+ raise MemoryError("Unable to allocate array")
+
+ monkeypatch.setattr(
+ "torch_molecule.datasets.split._butina_groups", _boom
+ )
+ with pytest.raises(MemoryError, match="Do not subsample") as excinfo:
+ train_test_split(data, method="butina")
+ assert "chemfp" in str(excinfo.value)
+ assert "Do not subsample" in _BUTINA_OOM_MESSAGE
diff --git a/tests/generator/pretrained_molexar.py b/tests/generator/pretrained_molexar.py
new file mode 100644
index 0000000..26ee478
--- /dev/null
+++ b/tests/generator/pretrained_molexar.py
@@ -0,0 +1,66 @@
+import os
+import shutil
+
+from torch_molecule import HFPretrainedMolecularGenerator
+
+REPO_ID = "fairydance/molexar-10m-base"
+N_SAMPLES = 2
+START_SMILES = "[*]C1(CC#N)CN(S(=O)(=O)CC)C1"
+
+
+def test_molexar_generator():
+ print("\n=== Testing Molexar initialization ===")
+ model = HFPretrainedMolecularGenerator(
+ repo_id=REPO_ID,
+ verbose="progress_bar",
+ )
+ print("Molexar initialized successfully")
+
+ print("\n=== Testing Molexar loading from Hugging Face ===")
+ model.fit()
+ print("Molexar loaded successfully")
+
+ print("\n=== Testing Molexar de novo generation ===")
+ generated_smiles = model.generate(
+ n_samples=N_SAMPLES,
+ max_new_tokens=64,
+ temperature=0.8,
+ )
+ print(f"Generated {len(generated_smiles)} molecules")
+ print("Example generated SMILES:", generated_smiles[:2])
+
+ print("\n=== Testing Molexar fragment-constrained generation ===")
+ generated_smiles = model.generate(
+ n_samples=N_SAMPLES,
+ start_smiles=START_SMILES,
+ generation_task="motif_extension",
+ max_new_tokens=64,
+ temperature=0.8,
+ )
+ print(f"Generated {len(generated_smiles)} molecules from start_smiles {START_SMILES}")
+ print("Example generated SMILES:", generated_smiles[:2])
+
+ print("\n=== Testing Molexar saving and loading ===")
+ save_path = "pretrained_molexar_test_model"
+ model.save_to_local(save_path)
+ print(f"Molexar saved to {save_path}")
+
+ loaded_model = HFPretrainedMolecularGenerator(repo_id=REPO_ID)
+ loaded_model.load_from_local(save_path)
+ print("Molexar loaded from local directory")
+
+ generated_smiles = loaded_model.generate(
+ n_samples=N_SAMPLES,
+ max_new_tokens=64,
+ temperature=0.8,
+ )
+ print(f"Generated {len(generated_smiles)} molecules with loaded model")
+ print("Example generated SMILES:", generated_smiles[:2])
+
+ if os.path.exists(save_path):
+ shutil.rmtree(save_path)
+ print(f"Cleaned up {save_path}")
+
+
+if __name__ == "__main__":
+ test_molexar_generator()
diff --git a/tests/generator/pretrained_molgen.py b/tests/generator/pretrained_molgen.py
new file mode 100644
index 0000000..1f9835a
--- /dev/null
+++ b/tests/generator/pretrained_molgen.py
@@ -0,0 +1,74 @@
+import os
+import shutil
+
+from torch_molecule import HFPretrainedMolecularGenerator
+
+N_SAMPLES = 2
+SCAFFOLD = "c1ccccc1"
+PREFIX_SELFIES = "[C][=C][C][=C][C][=C][Ring1][=Branch1]"
+
+
+def test_molgen_generator():
+ models_to_test = [
+ {"repo_id": "zjunlp/MolGen-large", "model_name": "MolGen-large"},
+ {"repo_id": "zjunlp/MolGen-large-opt", "model_name": "MolGen-large-opt"},
+ ]
+
+ for model_config in models_to_test:
+ name = model_config["model_name"]
+ repo_id = model_config["repo_id"]
+
+ print(f"\n=== Testing {name} initialization ===")
+ model = HFPretrainedMolecularGenerator(
+ repo_id=repo_id,
+ verbose="progress_bar",
+ )
+ print(f"{name} initialized successfully")
+
+ print(f"\n=== Testing {name} loading from Hugging Face ===")
+ model.fit()
+ print(f"{name} loaded successfully")
+
+ print(f"\n=== Testing {name} de novo generation ===")
+ generated_smiles = model.generate(n_samples=N_SAMPLES, num_beams=5)
+ print(f"Generated {len(generated_smiles)} molecules")
+ print("Example generated SMILES:", generated_smiles[:2])
+
+ print(f"\n=== Testing {name} scaffold generation ===")
+ generated_smiles = model.generate(
+ n_samples=N_SAMPLES,
+ scaffold=SCAFFOLD,
+ num_beams=5,
+ )
+ print(f"Generated {len(generated_smiles)} molecules from scaffold {SCAFFOLD}")
+ print("Example generated SMILES:", generated_smiles[:2])
+
+ print(f"\n=== Testing {name} prefix_selfies generation ===")
+ generated_smiles = model.generate(
+ n_samples=N_SAMPLES,
+ prefix_selfies=PREFIX_SELFIES,
+ num_beams=5,
+ )
+ print(f"Generated {len(generated_smiles)} molecules from prefix_selfies")
+ print("Example generated SMILES:", generated_smiles[:2])
+
+ print(f"\n=== Testing {name} saving and loading ===")
+ save_path = f"pretrained_{name.lower().replace('-', '_')}_test_model"
+ model.save_to_local(save_path)
+ print(f"{name} saved to {save_path}")
+
+ loaded_model = HFPretrainedMolecularGenerator(repo_id=repo_id)
+ loaded_model.load_from_local(save_path)
+ print(f"{name} loaded from local directory")
+
+ generated_smiles = loaded_model.generate(n_samples=N_SAMPLES, num_beams=5)
+ print(f"Generated {len(generated_smiles)} molecules with loaded model")
+ print("Example generated SMILES:", generated_smiles[:2])
+
+ if os.path.exists(save_path):
+ shutil.rmtree(save_path)
+ print(f"Cleaned up {save_path}")
+
+
+if __name__ == "__main__":
+ test_molgen_generator()
diff --git a/tests/generator/pretrained_novomolgen.py b/tests/generator/pretrained_novomolgen.py
new file mode 100644
index 0000000..d04cfbe
--- /dev/null
+++ b/tests/generator/pretrained_novomolgen.py
@@ -0,0 +1,86 @@
+import os
+import shutil
+
+from torch_molecule import HFPretrainedMolecularGenerator
+
+REPO_ID = "chandar-lab/NovoMolGen_32M_SMILES_BPE"
+N_SAMPLES = 5
+TRAIN_SMILES = [
+ "CC(=O)O",
+ "CCO",
+ "CCCC",
+ "c1ccccc1",
+ "CCN",
+]
+
+
+def test_novomolgen_generator():
+ print("\n=== Testing NovoMolGen initialization ===")
+ model = HFPretrainedMolecularGenerator(
+ repo_id=REPO_ID,
+ verbose="progress_bar",
+ )
+ print("NovoMolGen initialized successfully")
+
+ print("\n=== Testing NovoMolGen loading from Hugging Face ===")
+ model.fit()
+ print("NovoMolGen loaded successfully")
+
+ print("\n=== Testing NovoMolGen de novo generation ===")
+ generated_smiles = model.generate(n_samples=N_SAMPLES)
+ print(f"Generated {len(generated_smiles)} molecules")
+ print("Example generated SMILES:", generated_smiles[:2])
+
+ print("\n=== Testing NovoMolGen saving and loading ===")
+ save_path = "pretrained_novomolgen_test_model"
+ model.save_to_local(save_path)
+ print(f"NovoMolGen saved to {save_path}")
+
+ loaded_model = HFPretrainedMolecularGenerator(repo_id=REPO_ID)
+ loaded_model.load_from_local(save_path)
+ print("NovoMolGen loaded from local directory")
+
+ generated_smiles = loaded_model.generate(n_samples=2)
+ print(f"Generated {len(generated_smiles)} molecules with loaded model")
+ print("Example generated SMILES:", generated_smiles[:2])
+
+ if os.path.exists(save_path):
+ shutil.rmtree(save_path)
+ print(f"Cleaned up {save_path}")
+
+ print("\n=== Testing NovoMolGen fine-tuning ===")
+ finetune_model = HFPretrainedMolecularGenerator(
+ repo_id=REPO_ID,
+ batch_size=2,
+ epochs=1,
+ verbose="progress_bar",
+ )
+ finetune_model.fit(TRAIN_SMILES)
+ print("Fine-tuning completed")
+ print(f"Fitting epochs: {finetune_model.fitting_epoch + 1}")
+ print(f"Fitting loss: {finetune_model.fitting_loss}")
+
+ generated_smiles = finetune_model.generate(n_samples=2)
+ print(f"Generated {len(generated_smiles)} molecules after fine-tuning")
+ print("Example generated SMILES:", generated_smiles[:2])
+
+ print("\n=== Testing fine-tuned NovoMolGen saving and loading ===")
+ save_path = "pretrained_novomolgen_finetune_test_model"
+ finetune_model.save_to_local(save_path)
+ print(f"Fine-tuned NovoMolGen saved to {save_path}")
+
+ loaded_finetune = HFPretrainedMolecularGenerator(repo_id=REPO_ID)
+ loaded_finetune.load_from_local(save_path)
+ print("Fine-tuned NovoMolGen loaded from local directory")
+
+ generated_smiles = loaded_finetune.generate(n_samples=2)
+ print(f"Generated {len(generated_smiles)} molecules with loaded fine-tuned model")
+ print("Example generated SMILES:", generated_smiles[:2])
+
+ if os.path.exists(save_path):
+ shutil.rmtree(save_path)
+ print(f"Cleaned up {save_path}")
+
+
+if __name__ == "__main__":
+ test_novomolgen_generator()
diff --git a/tests/generator/pretrained_safe_gpt.py b/tests/generator/pretrained_safe_gpt.py
new file mode 100644
index 0000000..7d179bd
--- /dev/null
+++ b/tests/generator/pretrained_safe_gpt.py
@@ -0,0 +1,58 @@
+import os
+import shutil
+
+from torch_molecule import HFPretrainedMolecularGenerator
+
+REPO_ID = "datamol-io/safe-gpt"
+N_SAMPLES = 5
+SHORT_SCAFFOLD = "c1ccccc1"
+LONG_SCAFFOLD = "CC1=CC=C(C=C1)C2=CC(=NN2C3=CC=C(C=C3)S(=O)(=O)N)C(F)(F)F"
+
+
+def test_safe_gpt_generator():
+ print("\n=== Testing SAFE-GPT initialization ===")
+ model = HFPretrainedMolecularGenerator(
+ repo_id=REPO_ID,
+ verbose="progress_bar",
+ )
+ print("SAFE-GPT initialized successfully")
+
+ print("\n=== Testing SAFE-GPT loading from Hugging Face ===")
+ model.fit()
+ print("SAFE-GPT loaded successfully")
+
+ print("\n=== Testing SAFE-GPT de novo generation ===")
+ generated_smiles = model.generate(n_samples=N_SAMPLES)
+ print(f"Generated {len(generated_smiles)} molecules")
+ print("Example generated SMILES:", generated_smiles[:2])
+
+ print("\n=== Testing SAFE-GPT short scaffold generation ===")
+ generated_smiles = model.generate(n_samples=N_SAMPLES, scaffold=SHORT_SCAFFOLD)
+ print(f"Generated {len(generated_smiles)} molecules from scaffold {SHORT_SCAFFOLD}")
+ print("Example generated SMILES:", generated_smiles[:2])
+
+ print("\n=== Testing SAFE-GPT long scaffold generation ===")
+ generated_smiles = model.generate(n_samples=2, scaffold=LONG_SCAFFOLD)
+ print(f"Generated {len(generated_smiles)} molecules from long scaffold")
+ print("Example generated SMILES:", generated_smiles[:2])
+
+ print("\n=== Testing SAFE-GPT saving and loading ===")
+ save_path = "pretrained_safe_gpt_test_model"
+ model.save_to_local(save_path)
+ print(f"SAFE-GPT saved to {save_path}")
+
+ loaded_model = HFPretrainedMolecularGenerator(repo_id=REPO_ID)
+ loaded_model.load_from_local(save_path)
+ print("SAFE-GPT loaded from local directory")
+
+ generated_smiles = loaded_model.generate(n_samples=2, scaffold=SHORT_SCAFFOLD)
+ print(f"Generated {len(generated_smiles)} molecules with loaded model")
+ print("Example generated SMILES:", generated_smiles[:2])
+
+ if os.path.exists(save_path):
+ shutil.rmtree(save_path)
+ print(f"Cleaned up {save_path}")
+
+
+if __name__ == "__main__":
+ test_safe_gpt_generator()
diff --git a/torch_molecule/__init__.py b/torch_molecule/__init__.py
index 5b6ebff..11a1e4a 100644
--- a/torch_molecule/__init__.py
+++ b/torch_molecule/__init__.py
@@ -36,6 +36,7 @@
from .generator.lstm import LSTMMolecularGenerator
from .generator.molgpt import MolGPTMolecularGenerator
from .generator.defog import DeFoGMolecularGenerator
+from .generator.pretrained import HFPretrainedMolecularGenerator
__all__ = [
# 'BaseMolecularPredictor',
@@ -69,4 +70,5 @@
'MolGPTMolecularGenerator',
'LSTMMolecularGenerator',
'DeFoGMolecularGenerator',
+ 'HFPretrainedMolecularGenerator',
]
\ No newline at end of file
diff --git a/torch_molecule/datasets/__init__.py b/torch_molecule/datasets/__init__.py
index 407ce15..9fcb111 100644
--- a/torch_molecule/datasets/__init__.py
+++ b/torch_molecule/datasets/__init__.py
@@ -1,7 +1,12 @@
+from .constant import SMILESDataset
from .load_hf_dataset import load_qm9, load_chembl2k, load_broad6k, load_toxcast, load_admet, load_zinc250k
from .load_local_csv import load_gasperm
+from .split import subsample, train_test_split
__all__ = [
+ "SMILESDataset",
+ "subsample",
+ "train_test_split",
"load_qm9",
"load_chembl2k",
"load_broad6k",
diff --git a/torch_molecule/datasets/constant.py b/torch_molecule/datasets/constant.py
index 859b02c..f717fee 100644
--- a/torch_molecule/datasets/constant.py
+++ b/torch_molecule/datasets/constant.py
@@ -1,5 +1,5 @@
from dataclasses import dataclass
-from typing import List
+from typing import List, Tuple
import numpy as np
@dataclass
@@ -14,6 +14,59 @@ class SMILESDataset:
data: List[str]
target: np.ndarray | None
+ def subsample(self, n: int, seed: int = 0) -> "SMILESDataset":
+ """Draw a random subset without replacement.
+
+ Intended for local debugging and CI, not as a way to make structure-
+ aware splits cheaper. Do not subsample a benchmark (for example QM9)
+ because Butina is slow or memory-heavy; that changes what the split
+ measures.
+
+ Parameters
+ ----------
+ n : int
+ Number of molecules to keep.
+ seed : int, default=0
+ Random seed.
+ """
+ from .split import subsample
+
+ return subsample(self, n=n, seed=seed)
+
+ def train_test_split(
+ self,
+ test_size: float = 0.2,
+ method: str = "random",
+ seed: int = 42,
+ **kwargs,
+ ) -> Tuple["SMILESDataset", "SMILESDataset"]:
+ """Split into train and holdout ``SMILESDataset`` objects.
+
+ Parameters
+ ----------
+ test_size : float, default=0.2
+ Requested holdout fraction.
+ method : {"random", "scaffold", "butina", "size"}, default="random"
+ Split protocol. ``random`` is an i.i.d. baseline; ``scaffold``
+ holds out unseen Bemis-Murcko scaffolds; ``butina`` holds out
+ unseen Taylor-Butina clusters; ``size`` splits by heavy-atom count.
+ seed : int, default=42
+ Random seed (used by ``random``).
+ **kwargs
+ Extra options forwarded to the splitter (``use_csk`` for scaffold,
+ ``similarity_cutoff`` for butina, ``direction`` / ``mode`` for size).
+
+ Returns
+ -------
+ train, holdout : SMILESDataset
+ The second dataset is intended as validation data for ``fit``.
+ """
+ from .split import train_test_split
+
+ return train_test_split(
+ self, test_size=test_size, method=method, seed=seed, **kwargs
+ )
+
TOXCAST_TASKS = [
'ACEA_T47D_80hr_Negative', 'ACEA_T47D_80hr_Positive',
diff --git a/torch_molecule/datasets/split.py b/torch_molecule/datasets/split.py
new file mode 100644
index 0000000..f62c105
--- /dev/null
+++ b/torch_molecule/datasets/split.py
@@ -0,0 +1,421 @@
+"""Train/test splitting utilities for molecular SMILES datasets.
+
+Random splitting is an i.i.d. baseline. Scaffold splitting groups molecules by
+Bemis-Murcko frameworks so that the same scaffold does not appear in both
+splits. Butina splitting groups by Taylor-Butina clusters on Morgan fingerprints
+(sparse Tanimoto neighbor graph). Size splitting holds out larger (or smaller)
+molecules by heavy-atom count.
+
+The second split returned by ``train_test_split`` is a holdout set. Pass it to
+``fit(..., X_val, y_val)`` as validation data. A disjoint final test set requires
+a later three-way split API.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import Dict, List, Optional, Sequence, Tuple
+
+import numpy as np
+from rdkit import Chem, DataStructs
+from rdkit.Chem import rdFingerprintGenerator
+from rdkit.Chem.Scaffolds import MurckoScaffold
+
+from .constant import SMILESDataset
+
+_SUPPORTED_METHODS = ("random", "scaffold", "butina", "size")
+_SIZE_DIRECTIONS = ("small_to_large", "large_to_small")
+_SIZE_MODES = ("standard", "sizeshiftreg")
+_BUTINA_RADIUS = 2
+_BUTINA_FP_SIZE = 2048
+_BUTINA_DEFAULT_CUTOFF = 0.65
+_SIZESHIFTREG_TRAIN_FRACTION = 0.5
+_SIZESHIFTREG_TEST_FRACTION = 0.1
+_BUTINA_OOM_MESSAGE = (
+ "Butina split ran out of memory on the full dataset. "
+ "Do not subsample to work around a split failure; that changes what "
+ "the split measures. QM9-scale data is expected to fit; for much "
+ "larger libraries an optional chemfp backend may be added later."
+)
+
+_MORGAN_FP_GEN = rdFingerprintGenerator.GetMorganGenerator(
+ radius=_BUTINA_RADIUS, fpSize=_BUTINA_FP_SIZE
+)
+
+
+def subsample(
+ dataset: SMILESDataset,
+ n: int,
+ seed: int = 0,
+) -> SMILESDataset:
+ """Draw a random subset without replacement.
+
+ Intended for local debugging and CI, not as a way to make structure-aware
+ splits cheaper. Do not subsample a benchmark because Butina is slow or
+ memory-heavy; that changes what the split measures.
+
+ Parameters
+ ----------
+ dataset : SMILESDataset
+ Input dataset.
+ n : int
+ Number of molecules to keep. Must be at least 1 and at most the dataset
+ size.
+ seed : int, default=0
+ Random seed.
+
+ Returns
+ -------
+ SMILESDataset
+ Subsampled dataset. If ``n`` equals the dataset size, a copy with the
+ original order is returned.
+ """
+ _check_dataset(dataset)
+ n_total = len(dataset.data)
+ if n < 1:
+ raise ValueError(f"n must be >= 1, got {n}.")
+ if n > n_total:
+ raise ValueError(
+ f"n={n} is larger than the dataset size ({n_total})."
+ )
+ if n == n_total:
+ return _subset(dataset, list(range(n_total)))
+
+ rng = np.random.RandomState(seed)
+ indices = rng.choice(n_total, size=n, replace=False)
+ return _subset(dataset, indices.tolist())
+
+
+def train_test_split(
+ dataset: SMILESDataset,
+ test_size: float = 0.2,
+ method: str = "random",
+ seed: int = 42,
+ *,
+ use_csk: bool = False,
+ similarity_cutoff: float = _BUTINA_DEFAULT_CUTOFF,
+ direction: str = "small_to_large",
+ mode: str = "standard",
+) -> Tuple[SMILESDataset, SMILESDataset]:
+ """Split a SMILES dataset into train and holdout subsets.
+
+ Parameters
+ ----------
+ dataset : SMILESDataset
+ Input dataset.
+ test_size : float, default=0.2
+ Fraction of molecules requested for the holdout set. For scaffold and
+ Butina splits the realized fraction can differ because whole groups are
+ assigned together. Ignored when ``method="size"`` and
+ ``mode="sizeshiftreg"``.
+ method : {"random", "scaffold", "butina", "size"}, default="random"
+ ``"random"`` is an i.i.d. baseline (often optimistic for molecules).
+ ``"scaffold"`` holds out unseen Bemis-Murcko scaffolds.
+ ``"butina"`` holds out unseen Taylor-Butina clusters (Morgan / Tanimoto).
+ ``"size"`` holds out molecules by heavy-atom count.
+ seed : int, default=42
+ Random seed. Used by ``random``. Scaffold, Butina clustering, and size
+ assignment are deterministic; ``seed`` is accepted for API stability
+ and ignored.
+ use_csk : bool, default=False
+ Scaffold only. If True, generic cyclic skeletons are used (all atoms
+ as carbon). If False, atom types are kept (RDKit default).
+ similarity_cutoff : float, default=0.65
+ Butina only. Tanimoto **similarity** threshold in ``(0, 1]``. Molecules
+ with similarity at least this value are neighbors. This is not
+ DeepChem's distance cutoff.
+ direction : {"small_to_large", "large_to_small"}, default="small_to_large"
+ Size only (``mode="standard"``). ``small_to_large`` puts smaller
+ molecules in train and larger ones in holdout.
+ mode : {"standard", "sizeshiftreg"}, default="standard"
+ Size only. ``sizeshiftreg`` uses the SizeShiftReg protocol: smallest
+ 50% train, largest 10% holdout; the middle 40% is unused.
+
+ Returns
+ -------
+ train, holdout : SMILESDataset
+ The second dataset is a holdout split intended as validation data for
+ ``fit`` / ``autofit``.
+ """
+ _check_dataset(dataset)
+ if method not in _SUPPORTED_METHODS:
+ raise ValueError(
+ f"Unknown split method {method!r}. "
+ f"Supported methods: {list(_SUPPORTED_METHODS)}."
+ )
+ if not 0.0 < test_size < 1.0:
+ raise ValueError(f"test_size must be in (0, 1), got {test_size}.")
+
+ n = len(dataset.data)
+ if n < 2:
+ raise ValueError("Need at least 2 molecules to split a dataset.")
+
+ if method == "random":
+ idx_train, idx_test = _random_split(n, test_size, seed)
+ elif method == "scaffold":
+ groups = _scaffold_groups(dataset.data, use_csk=use_csk)
+ idx_train, idx_test = _group_split(groups, test_size)
+ elif method == "butina":
+ if not 0.0 < similarity_cutoff <= 1.0:
+ raise ValueError(
+ f"similarity_cutoff must be in (0, 1], got {similarity_cutoff}."
+ )
+ groups = _butina_groups_or_oom(
+ dataset.data, similarity_cutoff=similarity_cutoff
+ )
+ idx_train, idx_test = _group_split(groups, test_size)
+ else:
+ idx_train, idx_test = _size_split(
+ dataset.data,
+ test_size=test_size,
+ direction=direction,
+ mode=mode,
+ )
+
+ return _subset(dataset, idx_train), _subset(dataset, idx_test)
+
+
+def _check_dataset(dataset: SMILESDataset) -> None:
+ if not isinstance(dataset, SMILESDataset):
+ raise TypeError(
+ f"dataset must be a SMILESDataset, got {type(dataset).__name__}."
+ )
+ if not isinstance(dataset.data, list):
+ raise TypeError("dataset.data must be a list of SMILES strings.")
+ if dataset.target is not None:
+ target = np.asarray(dataset.target)
+ if target.shape[0] != len(dataset.data):
+ raise ValueError(
+ f"target has {target.shape[0]} rows but data has "
+ f"{len(dataset.data)} molecules."
+ )
+
+
+def _random_split(
+ n: int, test_size: float, seed: int
+) -> Tuple[List[int], List[int]]:
+ rng = np.random.RandomState(seed)
+ perm = rng.permutation(n)
+ n_test = int(round(n * test_size))
+ n_test = min(max(n_test, 1), n - 1)
+ idx_test = np.sort(perm[:n_test]).tolist()
+ idx_train = np.sort(perm[n_test:]).tolist()
+ return idx_train, idx_test
+
+
+def _mols_from_smiles(smiles_list: Sequence[str]) -> List[Chem.Mol]:
+ from ..utils.checker import MolecularInputChecker
+
+ invalid = []
+ mols: List[Optional[Chem.Mol]] = [None] * len(smiles_list)
+ for i, smiles in enumerate(smiles_list):
+ if not isinstance(smiles, str):
+ invalid.append(f"Non-string SMILES at index {i}: {smiles!r}")
+ continue
+ is_valid, error_msg, mol = MolecularInputChecker.validate_smiles(
+ smiles, i
+ )
+ if not is_valid:
+ invalid.append(error_msg)
+ continue
+ mols[i] = mol
+
+ if invalid:
+ raise ValueError("Invalid SMILES found:\n" + "\n".join(invalid))
+ return mols # type: ignore[return-value]
+
+
+def _scaffold_groups(
+ smiles_list: Sequence[str], use_csk: bool = False
+) -> Dict[str, List[int]]:
+ groups: Dict[str, List[int]] = defaultdict(list)
+ for i, mol in enumerate(_mols_from_smiles(smiles_list)):
+ groups[_scaffold_key(mol, use_csk=use_csk)].append(i)
+ return dict(groups)
+
+
+def _scaffold_key(mol: Chem.Mol, use_csk: bool = False) -> str:
+ scaffold = MurckoScaffold.GetScaffoldForMol(mol)
+ if scaffold is None or scaffold.GetNumAtoms() == 0:
+ return Chem.MolToSmiles(mol)
+
+ if use_csk:
+ scaffold = MurckoScaffold.MakeScaffoldGeneric(scaffold)
+ if scaffold is None or scaffold.GetNumAtoms() == 0:
+ return Chem.MolToSmiles(mol)
+ return Chem.MolToSmiles(scaffold)
+
+
+def _butina_groups(
+ smiles_list: Sequence[str],
+ similarity_cutoff: float = _BUTINA_DEFAULT_CUTOFF,
+) -> Dict[str, List[int]]:
+ clusters = _butina_clusters(smiles_list, similarity_cutoff=similarity_cutoff)
+ return {f"cluster_{i}": members for i, members in enumerate(clusters)}
+
+
+def _butina_groups_or_oom(
+ smiles_list: Sequence[str],
+ similarity_cutoff: float = _BUTINA_DEFAULT_CUTOFF,
+) -> Dict[str, List[int]]:
+ try:
+ return _butina_groups(
+ smiles_list, similarity_cutoff=similarity_cutoff
+ )
+ except MemoryError as exc:
+ raise MemoryError(_BUTINA_OOM_MESSAGE) from exc
+
+
+def _butina_clusters(
+ smiles_list: Sequence[str],
+ similarity_cutoff: float = _BUTINA_DEFAULT_CUTOFF,
+) -> List[List[int]]:
+ """Exact Taylor-Butina clusters via a sparse Tanimoto neighbor graph.
+
+ Matches RDKit ``Butina.ClusterData`` membership (``reordering=False``)
+ without storing the condensed distance matrix. Neighbor edges are pairs
+ with Tanimoto similarity >= ``similarity_cutoff``. Degree includes self,
+ matching RDKit's zero self-distance. Ties break like RDKit: higher index
+ first among equal degrees.
+ """
+ mols = _mols_from_smiles(smiles_list)
+ fps = [_MORGAN_FP_GEN.GetFingerprint(mol) for mol in mols]
+ neighbor_lists = _butina_neighbor_lists(fps, similarity_cutoff)
+ return _butina_exclusion_spheres(neighbor_lists)
+
+
+def _butina_neighbor_lists(
+ fps: Sequence[DataStructs.ExplicitBitVect],
+ similarity_cutoff: float,
+) -> List[List[int]]:
+ n = len(fps)
+ neighbors: List[List[int]] = [[] for _ in range(n)]
+ for i in range(n):
+ neighbors[i].append(i)
+ if i == 0:
+ continue
+ sims = DataStructs.BulkTanimotoSimilarity(fps[i], fps[:i])
+ if sims:
+ hits = np.flatnonzero(np.asarray(sims, dtype=np.float64) >= similarity_cutoff)
+ for j in hits.tolist():
+ neighbors[i].append(int(j))
+ neighbors[j].append(i)
+ for i in range(n):
+ neighbors[i].sort()
+ return neighbors
+
+
+def _butina_exclusion_spheres(
+ neighbor_lists: Sequence[Sequence[int]],
+) -> List[List[int]]:
+ n = len(neighbor_lists)
+ sorted_indices = [
+ (len(nbrs), idx) for idx, nbrs in enumerate(neighbor_lists)
+ ]
+ sorted_indices.sort(reverse=True)
+
+ clusters: List[List[int]] = []
+ seen = np.zeros(n, dtype=bool)
+
+ while sorted_indices and sorted_indices[0][0] > 1:
+ _, idx = sorted_indices.pop(0)
+ if seen[idx]:
+ continue
+ cluster = [idx]
+ seen[idx] = True
+ for neighbor in neighbor_lists[idx]:
+ if not seen[neighbor]:
+ cluster.append(neighbor)
+ seen[neighbor] = True
+ clusters.append(cluster)
+
+ while sorted_indices:
+ _, idx = sorted_indices.pop(0)
+ if seen[idx]:
+ continue
+ clusters.append([idx])
+ return clusters
+
+
+def _size_split(
+ smiles_list: Sequence[str],
+ test_size: float,
+ direction: str = "small_to_large",
+ mode: str = "standard",
+) -> Tuple[List[int], List[int]]:
+ if direction not in _SIZE_DIRECTIONS:
+ raise ValueError(
+ f"Unknown size direction {direction!r}. "
+ f"Supported directions: {list(_SIZE_DIRECTIONS)}."
+ )
+ if mode not in _SIZE_MODES:
+ raise ValueError(
+ f"Unknown size mode {mode!r}. "
+ f"Supported modes: {list(_SIZE_MODES)}."
+ )
+
+ mols = _mols_from_smiles(smiles_list)
+ n_atoms = np.array([mol.GetNumHeavyAtoms() for mol in mols], dtype=np.int64)
+ order = np.argsort(n_atoms, kind="stable")
+ n = len(order)
+
+ if mode == "sizeshiftreg":
+ n_train = int(round(_SIZESHIFTREG_TRAIN_FRACTION * n))
+ n_test = int(round(_SIZESHIFTREG_TEST_FRACTION * n))
+ n_train = min(max(n_train, 1), n - 1)
+ n_test = min(max(n_test, 1), n - n_train)
+ idx_train = order[:n_train].tolist()
+ idx_test = order[-n_test:].tolist()
+ return idx_train, idx_test
+
+ n_test = int(round(n * test_size))
+ n_test = min(max(n_test, 1), n - 1)
+ if direction == "small_to_large":
+ idx_train = order[:-n_test].tolist()
+ idx_test = order[-n_test:].tolist()
+ else:
+ idx_train = order[n_test:].tolist()
+ idx_test = order[:n_test].tolist()
+ return idx_train, idx_test
+
+
+def _group_split(
+ groups: Dict[str, List[int]], test_size: float
+) -> Tuple[List[int], List[int]]:
+ """Assign whole groups with DeepChem-style greedy filling.
+
+ Groups are sorted by decreasing size (group id as a tie-break). A group
+ goes to train if it still fits under the train cutoff; otherwise it goes to
+ the holdout set.
+ """
+ n = sum(len(idx) for idx in groups.values())
+ train_cutoff = (1.0 - test_size) * n
+ ordered = sorted(groups.items(), key=lambda item: (-len(item[1]), item[0]))
+
+ idx_train: List[int] = []
+ idx_test: List[int] = []
+ for _, members in ordered:
+ members_sorted = sorted(members)
+ if len(idx_train) + len(members_sorted) > train_cutoff:
+ idx_test.extend(members_sorted)
+ else:
+ idx_train.extend(members_sorted)
+
+ if not idx_train or not idx_test:
+ raise ValueError(
+ "Group split produced an empty train or holdout set. "
+ "Try a different test_size or a dataset with more diverse groups."
+ )
+ return idx_train, idx_test
+
+
+def _subset(dataset: SMILESDataset, indices: Sequence[int]) -> SMILESDataset:
+ indices = list(indices)
+ data = [dataset.data[i] for i in indices]
+ if dataset.target is None:
+ target: Optional[np.ndarray] = None
+ else:
+ target = np.asarray(dataset.target)[indices]
+ if target.ndim == 1:
+ target = target.reshape(-1, 1)
+ return SMILESDataset(data=data, target=target)
diff --git a/torch_molecule/generator/pretrained/__init__.py b/torch_molecule/generator/pretrained/__init__.py
new file mode 100644
index 0000000..212000f
--- /dev/null
+++ b/torch_molecule/generator/pretrained/__init__.py
@@ -0,0 +1,3 @@
+from .modeling_pretrained import HFPretrainedMolecularGenerator
+
+__all__ = ["HFPretrainedMolecularGenerator"]
diff --git a/torch_molecule/generator/pretrained/checkpoint.py b/torch_molecule/generator/pretrained/checkpoint.py
new file mode 100644
index 0000000..1d46604
--- /dev/null
+++ b/torch_molecule/generator/pretrained/checkpoint.py
@@ -0,0 +1,46 @@
+"""Local save/load helpers for HF pretrained generators."""
+
+import json
+import os
+from typing import Any, Dict, Optional
+
+METADATA_FILENAME = "hf_generator_metadata.json"
+
+
+def build_metadata(
+ repo_id: str,
+ family: str,
+ max_length: int,
+ revision: Optional[str],
+ trust_remote_code: bool,
+ tokenizer_repo_id: Optional[str],
+ generate_max_length: int,
+ model_name: str,
+) -> Dict[str, Any]:
+ return {
+ "repo_id": repo_id,
+ "family": family,
+ "max_length": max_length,
+ "revision": revision,
+ "trust_remote_code": trust_remote_code,
+ "tokenizer_repo_id": tokenizer_repo_id,
+ "generate_max_length": generate_max_length,
+ "model_name": model_name,
+ }
+
+
+def save_metadata(path: str, metadata: Dict[str, Any]) -> None:
+ os.makedirs(path, exist_ok=True)
+ with open(os.path.join(path, METADATA_FILENAME), "w", encoding="utf-8") as handle:
+ json.dump(metadata, handle, indent=2)
+
+
+def load_metadata(path: str) -> Dict[str, Any]:
+ metadata_path = os.path.join(path, METADATA_FILENAME)
+ if not os.path.exists(metadata_path):
+ raise FileNotFoundError(
+ f"Missing {METADATA_FILENAME} in '{path}'. Expected a directory saved by "
+ "HFPretrainedMolecularGenerator.save_to_local()."
+ )
+ with open(metadata_path, "r", encoding="utf-8") as handle:
+ return json.load(handle)
diff --git a/torch_molecule/generator/pretrained/compat.py b/torch_molecule/generator/pretrained/compat.py
new file mode 100644
index 0000000..172572d
--- /dev/null
+++ b/torch_molecule/generator/pretrained/compat.py
@@ -0,0 +1,33 @@
+"""Compatibility helpers for Hugging Face pretrained generators."""
+
+from typing import Any
+
+
+def ensure_safe_transformers_compat() -> None:
+ """Restore generation constraint symbols removed in transformers 5.
+
+ PyPI ``safe-mol`` still imports ``DisjunctiveConstraint`` and
+ ``PhrasalConstraint`` from ``transformers.generation`` when the package
+ is imported. Those classes were removed in transformers 5. Dummy
+ stand-ins are enough to import ``safe.converter`` and run standard GPT-2
+ ``generate()``. Official ``SAFEDesign`` constrained beam search is not used.
+ """
+ try:
+ import transformers.generation as generation
+ except ImportError:
+ return
+
+ if hasattr(generation, "DisjunctiveConstraint") and hasattr(generation, "PhrasalConstraint"):
+ return
+
+ class _RemovedConstraint:
+ def __init__(self, *args: Any, **kwargs: Any) -> None:
+ raise NotImplementedError(
+ "Constrained beam search was removed in transformers 5. "
+ "SAFE-GPT in torch-molecule uses standard GPT-2 generate()."
+ )
+
+ if not hasattr(generation, "DisjunctiveConstraint"):
+ generation.DisjunctiveConstraint = _RemovedConstraint
+ if not hasattr(generation, "PhrasalConstraint"):
+ generation.PhrasalConstraint = _RemovedConstraint
diff --git a/torch_molecule/generator/pretrained/families/__init__.py b/torch_molecule/generator/pretrained/families/__init__.py
new file mode 100644
index 0000000..476efaa
--- /dev/null
+++ b/torch_molecule/generator/pretrained/families/__init__.py
@@ -0,0 +1,10 @@
+from .causal_lm import generate_causal_lm
+from .molexar import generate_molexar
+from .seq2seq import DEFAULT_MOLGEN_PREFIX_SELFIES, generate_seq2seq
+
+__all__ = [
+ "generate_causal_lm",
+ "generate_molexar",
+ "generate_seq2seq",
+ "DEFAULT_MOLGEN_PREFIX_SELFIES",
+]
diff --git a/torch_molecule/generator/pretrained/families/causal_lm.py b/torch_molecule/generator/pretrained/families/causal_lm.py
new file mode 100644
index 0000000..4cddf23
--- /dev/null
+++ b/torch_molecule/generator/pretrained/families/causal_lm.py
@@ -0,0 +1,96 @@
+"""Causal language model generation for SMILES-based HF generators."""
+
+from typing import Any, List, Optional
+
+import torch
+
+
+def generate_causal_lm(
+ model: torch.nn.Module,
+ tokenizer: Any,
+ device: torch.device,
+ n_samples: int,
+ *,
+ family: Optional[str] = None,
+ max_new_tokens: Optional[int] = None,
+ max_length: Optional[int] = None,
+ temperature: float = 1.0,
+ do_sample: bool = True,
+ scaffold: Optional[str] = None,
+ **kwargs: Any,
+) -> List[str]:
+ """Generate SMILES strings with a causal language model.
+
+ Parameters
+ ----------
+ model : torch.nn.Module
+ A Hugging Face causal LM.
+ tokenizer : transformers.PreTrainedTokenizer
+ Tokenizer paired with the model.
+ device : torch.device
+ Device used for generation.
+ n_samples : int
+ Number of molecules to generate.
+ family : Optional[str], default=None
+ Generator family name. Unused by the standard ``generate()`` path;
+ kept for call-site compatibility.
+ max_new_tokens : Optional[int], default=None
+ Maximum number of newly generated tokens. Independent of prefix length,
+ so a long ``scaffold=`` does not consume the generation budget.
+ Used when ``max_length`` is omitted; defaults to 64.
+ max_length : Optional[int], default=None
+ Optional Hugging Face total sequence length (prefix + new tokens).
+ When set without ``max_new_tokens``, this is passed through instead.
+ temperature : float, default=1.0
+ Sampling temperature.
+ do_sample : bool, default=True
+ Whether to use sampling during generation.
+ scaffold : Optional[str], default=None
+ Optional tokenized prefix. Callers that need a SMILES-to-SAFE
+ conversion (SAFE-GPT) should pass the already-encoded prefix.
+
+ Returns
+ -------
+ List[str]
+ Raw decoded strings from the tokenizer (may contain spaces).
+ """
+ del family # dispatch is done by HFPretrainedMolecularGenerator
+ pad_token_id = tokenizer.pad_token_id
+ if pad_token_id is None:
+ pad_token_id = tokenizer.eos_token_id
+
+ generate_kwargs = {
+ "do_sample": do_sample,
+ "pad_token_id": pad_token_id,
+ }
+ # Hugging Face rejects passing both; prefer max_new_tokens unless the caller
+ # explicitly asks for total-length max_length.
+ if max_length is not None and max_new_tokens is None:
+ generate_kwargs["max_length"] = max_length
+ else:
+ generate_kwargs["max_new_tokens"] = (
+ max_new_tokens if max_new_tokens is not None else 64
+ )
+ if do_sample:
+ generate_kwargs["temperature"] = temperature
+ generate_kwargs.update(kwargs)
+
+ if scaffold:
+ encoded = tokenizer(scaffold, return_tensors="pt", add_special_tokens=False)
+ input_ids = encoded["input_ids"]
+ input_ids = input_ids.to(device).expand(n_samples, -1).contiguous()
+ generate_kwargs["input_ids"] = input_ids
+ generate_kwargs["attention_mask"] = torch.ones_like(input_ids)
+ else:
+ if tokenizer.bos_token_id is None:
+ raise ValueError(
+ "Tokenizer has no BOS token. Provide `scaffold=` or use a model "
+ "with a defined bos_token_id."
+ )
+ input_ids = torch.tensor([[tokenizer.bos_token_id]], device=device)
+ generate_kwargs["input_ids"] = input_ids.expand(n_samples, -1).contiguous()
+
+ with torch.no_grad():
+ outputs = model.generate(**generate_kwargs)
+
+ return tokenizer.batch_decode(outputs, skip_special_tokens=True)
diff --git a/torch_molecule/generator/pretrained/families/molexar.py b/torch_molecule/generator/pretrained/families/molexar.py
new file mode 100644
index 0000000..7d3fdc8
--- /dev/null
+++ b/torch_molecule/generator/pretrained/families/molexar.py
@@ -0,0 +1,129 @@
+"""Molexar Fragment-SELFIES generation."""
+
+from typing import Any, Dict, List, Optional
+
+SINGLE_FRAGMENT_TASKS = frozenset({"motif_extension", "scaffold_decoration"})
+TWO_FRAGMENT_TASKS = frozenset({"linker_design", "scaffold_morphing"})
+FRAGMENT_CONSTRAINED_TASKS = SINGLE_FRAGMENT_TASKS | TWO_FRAGMENT_TASKS | frozenset({"superstructure"})
+
+PROPERTY_KEYS = (
+ "mol_hac",
+ "mol_hbdc",
+ "mol_hbac",
+ "mol_rotbc",
+ "mol_wt",
+ "mol_logp",
+ "mol_tpsa",
+ "mol_qed",
+ "mol_sas",
+)
+
+
+def _require_molexar():
+ try:
+ from molexar.inference import MolexarInference # noqa: F401
+ except ImportError as exc:
+ raise ImportError(
+ "The 'molexar' package is required for Molexar generation. "
+ "Install it with `pip install git+https://github.com/fairydance/Molexar.git`."
+ ) from exc
+
+
+def resolve_start_string(
+ *,
+ start_string: Optional[str] = None,
+ start_smiles: Optional[str] = None,
+ start_fragment_selfies: Optional[str] = None,
+ generation_task: Optional[str] = None,
+) -> Optional[str]:
+ """Resolve the Fragment-SELFIES prefix placed after ````."""
+ if start_string is not None:
+ return start_string
+ if start_fragment_selfies is not None:
+ return start_fragment_selfies
+ if start_smiles is None:
+ return None
+
+ task = generation_task or "motif_extension"
+ if task == "de_novo":
+ raise ValueError("de_novo generation does not accept start_smiles")
+
+ from molexar.data.converter import smiles_fragment_to_fragment_selfies
+
+ if task in SINGLE_FRAGMENT_TASKS | {"superstructure"}:
+ encoded = smiles_fragment_to_fragment_selfies(start_smiles, randomized=True)
+ return f"{encoded}[Attach:0]"
+
+ if task in TWO_FRAGMENT_TASKS:
+ fragments = [fragment.strip() for fragment in start_smiles.split(".") if fragment.strip()]
+ if len(fragments) != 2:
+ raise ValueError(
+ "linker_design and scaffold_morphing require exactly two "
+ "dot-separated SMILES fragments"
+ )
+ encoded_fragments = [
+ smiles_fragment_to_fragment_selfies(fragment, randomized=True) for fragment in fragments
+ ]
+ return "".join(encoded_fragments) + "[Attach:0]"
+
+ raise ValueError(
+ f"Unsupported generation_task '{task}'. Supported tasks: de_novo, "
+ "motif_extension, scaffold_decoration, linker_design, scaffold_morphing, superstructure"
+ )
+
+
+def extract_conditions(kwargs: Dict[str, Any]) -> Dict[str, Any]:
+ """Extract Molexar condition kwargs into a conditions dictionary."""
+ conditions = dict(kwargs.pop("conditions", {}) or {})
+ for key in PROPERTY_KEYS:
+ if key in kwargs:
+ conditions[key] = kwargs.pop(key)
+ for key in ("mol_pharma_fp", "prot_seq_esm_emb", "prot_poc_gvp_emb"):
+ if key in kwargs:
+ conditions[key] = kwargs.pop(key)
+ return conditions
+
+
+def generate_molexar(
+ engine: Any,
+ n_samples: int,
+ *,
+ start_string: Optional[str] = None,
+ start_smiles: Optional[str] = None,
+ start_fragment_selfies: Optional[str] = None,
+ generation_task: Optional[str] = None,
+ conditions: Optional[Dict[str, Any]] = None,
+ max_new_tokens: Optional[int] = None,
+ temperature: float = 0.8,
+ top_p: float = 0.95,
+ top_k: int = 50,
+ do_sample: bool = True,
+ repetition_penalty: float = 1.0,
+ batch_size: int = 100,
+ **kwargs: Any,
+) -> List[str]:
+ """Generate Fragment-SELFIES strings with a Molexar inference engine."""
+ _require_molexar()
+
+ merged_conditions = dict(conditions or {})
+
+ resolved_start = resolve_start_string(
+ start_string=start_string,
+ start_smiles=start_smiles,
+ start_fragment_selfies=start_fragment_selfies,
+ generation_task=generation_task,
+ )
+
+ return engine.generate(
+ conditions=merged_conditions,
+ start_string=resolved_start,
+ max_new_tokens=max_new_tokens,
+ num_samples=n_samples,
+ temperature=temperature,
+ top_p=top_p,
+ top_k=top_k,
+ do_sample=do_sample,
+ repetition_penalty=repetition_penalty,
+ batch_size=batch_size,
+ **kwargs,
+ )
diff --git a/torch_molecule/generator/pretrained/families/seq2seq.py b/torch_molecule/generator/pretrained/families/seq2seq.py
new file mode 100644
index 0000000..99d11f1
--- /dev/null
+++ b/torch_molecule/generator/pretrained/families/seq2seq.py
@@ -0,0 +1,77 @@
+"""Seq2Seq generation for SELFIES-based HF generators such as MolGen."""
+
+from typing import Any, List, Optional
+
+import torch
+
+DEFAULT_MOLGEN_PREFIX_SELFIES = "[C][=C][C][=C][C][=C][Ring1][=Branch1]"
+
+
+def generate_seq2seq(
+ model: torch.nn.Module,
+ tokenizer: Any,
+ device: torch.device,
+ n_samples: int,
+ *,
+ prefix_selfies: Optional[str] = None,
+ max_length: int = 15,
+ min_length: int = 5,
+ num_beams: int = 5,
+ **kwargs: Any,
+) -> List[str]:
+ """Generate SELFIES strings with a seq2seq language model.
+
+ MolGen uses a corrupted SELFIES prefix as input and generates a completed
+ SELFIES sequence via beam search.
+
+ Parameters
+ ----------
+ model : torch.nn.Module
+ A Hugging Face seq2seq model.
+ tokenizer : transformers.PreTrainedTokenizer
+ Tokenizer paired with the model.
+ device : torch.device
+ Device used for generation.
+ n_samples : int
+ Number of molecules to generate.
+ prefix_selfies : Optional[str], default=None
+ SELFIES prefix used as model input. Defaults to a benzene ring fragment.
+ max_length : int, default=15
+ Maximum generated sequence length.
+ min_length : int, default=5
+ Minimum generated sequence length.
+ num_beams : int, default=5
+ Beam width for beam search.
+
+ Returns
+ -------
+ List[str]
+ Raw decoded SELFIES strings from the tokenizer.
+ """
+ prefix = prefix_selfies or DEFAULT_MOLGEN_PREFIX_SELFIES
+ encoded = tokenizer(prefix, return_tensors="pt")
+ input_ids = encoded["input_ids"].to(device)
+ attention_mask = encoded["attention_mask"].to(device)
+
+ beam_width = max(num_beams, n_samples)
+ generate_kwargs = {
+ "input_ids": input_ids,
+ "attention_mask": attention_mask,
+ "max_length": max_length,
+ "min_length": min_length,
+ "num_return_sequences": n_samples,
+ "num_beams": beam_width,
+ }
+ generate_kwargs.update(kwargs)
+
+ with torch.no_grad():
+ outputs = model.generate(**generate_kwargs)
+
+ return [
+ tokenizer.decode(
+ sequence,
+ skip_special_tokens=True,
+ clean_up_tokenization_spaces=True,
+ )
+ for sequence in outputs
+ ]
diff --git a/torch_molecule/generator/pretrained/finetune.py b/torch_molecule/generator/pretrained/finetune.py
new file mode 100644
index 0000000..e9199fd
--- /dev/null
+++ b/torch_molecule/generator/pretrained/finetune.py
@@ -0,0 +1,288 @@
+"""Fine-tuning utilities for Hugging Face pretrained generators."""
+
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+
+from .registry import MOLEXAR_FAMILIES, SEQ2SEQ_FAMILIES
+
+
+def corrupt_token_ids(
+ input_ids: torch.Tensor,
+ tokenizer: Any,
+ *,
+ mask_prob: float = 0.15,
+ generator: Optional[torch.Generator] = None,
+) -> torch.Tensor:
+ """Replace a random subset of non-special tokens with the tokenizer mask id.
+
+ MolGen is trained as a denoising seq2seq model: corrupted SELFIES in,
+ clean SELFIES as labels. Special tokens (BOS/EOS/PAD/mask) are left intact.
+ """
+ if tokenizer.mask_token_id is None:
+ raise ValueError(
+ "MolGen denoising fine-tuning requires tokenizer.mask_token_id. "
+ "MolGen tokenizers provide a token."
+ )
+
+ corrupted = input_ids.clone()
+ special = torch.zeros_like(corrupted, dtype=torch.bool)
+ for special_id in tokenizer.all_special_ids:
+ special |= corrupted == special_id
+ if tokenizer.pad_token_id is not None:
+ special |= corrupted == tokenizer.pad_token_id
+
+ probs = torch.rand(corrupted.shape, generator=generator, device=corrupted.device)
+ to_mask = (probs < mask_prob) & ~special
+ corrupted[to_mask] = tokenizer.mask_token_id
+ return corrupted
+
+
+class _TokenizedDataset(Dataset):
+ def __init__(self, input_ids: List[List[int]], attention_mask: List[List[int]]):
+ self.input_ids = input_ids
+ self.attention_mask = attention_mask
+
+ def __len__(self) -> int:
+ return len(self.input_ids)
+
+ def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+ return {
+ "input_ids": torch.tensor(self.input_ids[idx]),
+ "attention_mask": torch.tensor(self.attention_mask[idx]),
+ }
+
+
+def _batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]:
+ return {key: value.to(device) for key, value in batch.items()}
+
+
+def _run_training_loop(
+ model: torch.nn.Module,
+ train_loader: DataLoader,
+ optimizer: torch.optim.Optimizer,
+ device: torch.device,
+ epochs: int,
+ grad_norm_clip: Optional[float],
+ verbose: str,
+) -> Tuple[List[float], int]:
+ model.train()
+ epoch_losses: List[float] = []
+ last_epoch = 0
+
+ for epoch in range(epochs):
+ last_epoch = epoch
+ batch_losses: List[float] = []
+ iterator = train_loader
+ if verbose in {"progress_bar", "print_statement"}:
+ iterator = tqdm(train_loader, desc=f"Fine-tuning epoch {epoch + 1}/{epochs}")
+
+ for batch in iterator:
+ batch = _batch_to_device(batch, device)
+ optimizer.zero_grad()
+ outputs = model(**batch)
+ loss = outputs.loss
+ loss.backward()
+
+ if grad_norm_clip is not None:
+ torch.nn.utils.clip_grad_norm_(model.parameters(), grad_norm_clip)
+
+ optimizer.step()
+ batch_losses.append(float(loss.detach().cpu()))
+
+ epoch_losses.append(float(np.mean(batch_losses)) if batch_losses else 0.0)
+ if verbose == "print_statement":
+ print(f"Epoch {epoch + 1}/{epochs} loss: {epoch_losses[-1]:.4f}")
+
+ model.eval()
+ return epoch_losses, last_epoch
+
+
+def finetune_causal_lm(
+ model: torch.nn.Module,
+ tokenizer: Any,
+ texts: List[str],
+ device: torch.device,
+ *,
+ max_length: int,
+ batch_size: int,
+ epochs: int,
+ learning_rate: float,
+ weight_decay: float,
+ grad_norm_clip: Optional[float],
+ verbose: str,
+) -> Tuple[List[float], int]:
+ """Fine-tune a causal language model with next-token prediction."""
+ import transformers
+
+ tokenized = tokenizer(
+ texts,
+ truncation=True,
+ max_length=max_length,
+ padding=False,
+ )
+ dataset = _TokenizedDataset(tokenized["input_ids"], tokenized["attention_mask"])
+ collator = transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False)
+ train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collator)
+
+ optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
+ return _run_training_loop(model, train_loader, optimizer, device, epochs, grad_norm_clip, verbose)
+
+
+def finetune_seq2seq(
+ model: torch.nn.Module,
+ tokenizer: Any,
+ texts: List[str],
+ device: torch.device,
+ *,
+ max_length: int,
+ batch_size: int,
+ epochs: int,
+ learning_rate: float,
+ weight_decay: float,
+ grad_norm_clip: Optional[float],
+ verbose: str,
+ mask_prob: float = 0.15,
+) -> Tuple[List[float], int]:
+ """Fine-tune a seq2seq model with MolGen-style denoising.
+
+ Encoder inputs are token-masked SELFIES; labels remain the original
+ clean sequence. Causal LM and Molexar fine-tuning are unchanged.
+ """
+ import transformers
+
+ class _Seq2SeqDataset(Dataset):
+ def __init__(self, items: List[str]):
+ self.items = items
+
+ def __len__(self):
+ return len(self.items)
+
+ def __getitem__(self, idx):
+ encoded = tokenizer(
+ self.items[idx],
+ truncation=True,
+ max_length=max_length,
+ padding=False,
+ )
+ item = {key: torch.tensor(value) for key, value in encoded.items()}
+ item["labels"] = item["input_ids"].clone()
+ item["input_ids"] = corrupt_token_ids(
+ item["input_ids"],
+ tokenizer,
+ mask_prob=mask_prob,
+ )
+ return item
+
+ collator = transformers.DataCollatorForSeq2Seq(tokenizer, model=model, padding=True)
+ train_loader = DataLoader(
+ _Seq2SeqDataset(texts),
+ batch_size=batch_size,
+ shuffle=True,
+ collate_fn=collator,
+ )
+ optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
+ return _run_training_loop(model, train_loader, optimizer, device, epochs, grad_norm_clip, verbose)
+
+
+def finetune_molexar(
+ model: torch.nn.Module,
+ tokenizer: Any,
+ config: Any,
+ texts: List[str],
+ device: torch.device,
+ *,
+ max_length: int,
+ batch_size: int,
+ epochs: int,
+ learning_rate: float,
+ weight_decay: float,
+ grad_norm_clip: Optional[float],
+ verbose: str,
+) -> Tuple[List[float], int]:
+ """Fine-tune a Molexar model on Fragment-SELFIES training templates."""
+ from molexar.templates import build_condition_template, build_training_text
+
+ condition_block, _ = build_condition_template(config)
+ training_texts = [build_training_text(config, condition_block, text) for text in texts]
+ return finetune_causal_lm(
+ model,
+ tokenizer,
+ training_texts,
+ device,
+ max_length=max_length,
+ batch_size=batch_size,
+ epochs=epochs,
+ learning_rate=learning_rate,
+ weight_decay=weight_decay,
+ grad_norm_clip=grad_norm_clip,
+ verbose=verbose,
+ )
+
+
+def finetune_generator(
+ family: str,
+ model: torch.nn.Module,
+ tokenizer: Any,
+ texts: List[str],
+ device: torch.device,
+ *,
+ config: Optional[Any] = None,
+ max_length: int,
+ batch_size: int,
+ epochs: int,
+ learning_rate: float,
+ weight_decay: float,
+ grad_norm_clip: Optional[float],
+ verbose: str,
+) -> Tuple[List[float], int]:
+ """Dispatch fine-tuning to the family-specific routine."""
+ if family in SEQ2SEQ_FAMILIES:
+ return finetune_seq2seq(
+ model,
+ tokenizer,
+ texts,
+ device,
+ max_length=max_length,
+ batch_size=batch_size,
+ epochs=epochs,
+ learning_rate=learning_rate,
+ weight_decay=weight_decay,
+ grad_norm_clip=grad_norm_clip,
+ verbose=verbose,
+ )
+
+ if family in MOLEXAR_FAMILIES:
+ if config is None:
+ raise ValueError("Molexar fine-tuning requires a model config.")
+ return finetune_molexar(
+ model,
+ tokenizer,
+ config,
+ texts,
+ device,
+ max_length=max_length,
+ batch_size=batch_size,
+ epochs=epochs,
+ learning_rate=learning_rate,
+ weight_decay=weight_decay,
+ grad_norm_clip=grad_norm_clip,
+ verbose=verbose,
+ )
+
+ return finetune_causal_lm(
+ model,
+ tokenizer,
+ texts,
+ device,
+ max_length=max_length,
+ batch_size=batch_size,
+ epochs=epochs,
+ learning_rate=learning_rate,
+ weight_decay=weight_decay,
+ grad_norm_clip=grad_norm_clip,
+ verbose=verbose,
+ )
diff --git a/torch_molecule/generator/pretrained/modeling_pretrained.py b/torch_molecule/generator/pretrained/modeling_pretrained.py
index f87f5c1..010364d 100644
--- a/torch_molecule/generator/pretrained/modeling_pretrained.py
+++ b/torch_molecule/generator/pretrained/modeling_pretrained.py
@@ -1 +1,639 @@
-# TODO
\ No newline at end of file
+import warnings
+import os
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from ...base import BaseMolecularGenerator
+from .checkpoint import build_metadata, load_metadata, save_metadata
+from .families.causal_lm import generate_causal_lm
+from .families.molexar import extract_conditions, generate_molexar
+from .families.seq2seq import generate_seq2seq
+from .finetune import finetune_generator
+from .registry import (
+ CAUSAL_LM_FAMILIES,
+ MOLEXAR_FAMILIES,
+ SAFE_GPT_FAMILIES,
+ SEQ2SEQ_FAMILIES,
+ resolve_family,
+)
+
+
+class HFPretrainedMolecularGenerator(BaseMolecularGenerator):
+ """Hugging Face pretrained models as molecular generators.
+
+ This class loads pretrained generative models from Hugging Face and exposes
+ a sklearn-style ``fit`` / ``generate`` interface consistent with other
+ generators in torch-molecule.
+
+ Supported generation modes depend on the model family:
+
+ - NovoMolGen: de novo SMILES generation from BOS.
+ - MolGen-large / MolGen-large-opt: SELFIES seq2seq generation via ``prefix_selfies=``
+ or ``scaffold=`` (SMILES converted internally).
+ - Molexar: Fragment-SELFIES de novo and fragment-constrained generation via
+ ``start_smiles`` / ``start_string`` / ``conditions`` (omni).
+ - SAFE-GPT: GPT-2 causal LM on SAFE strings; de novo and ``scaffold=`` prefix.
+
+ Other registered families can be loaded but may raise ``NotImplementedError``
+ until later phases are implemented.
+
+ Tested models include:
+
+ - NovoMolGen: Causal LM pretrained on ZINC-22 for de novo SMILES generation.
+
+ repo_id: ``"chandar-lab/NovoMolGen_32M_SMILES_BPE"``
+ (https://huggingface.co/chandar-lab/NovoMolGen_32M_SMILES_BPE)
+
+ - MolGen-large: Seq2Seq SELFIES generator with high chemical validity.
+
+ repo_id: ``"zjunlp/MolGen-large"``
+ (https://huggingface.co/zjunlp/MolGen-large)
+
+ - MolGen-large-opt: MolGen-large fine-tuned for QED / p-logP optimization.
+
+ repo_id: ``"zjunlp/MolGen-large-opt"``
+ (https://huggingface.co/zjunlp/MolGen-large-opt)
+
+ - Molexar-10M-base: Fragment-SELFIES de novo and fragment-constrained generation.
+
+ repo_id: ``"fairydance/molexar-10m-base"``
+ (https://huggingface.co/fairydance/molexar-10m-base)
+
+ - Molexar-10M-omni: Multi-condition Molexar model for property-guided generation.
+
+ repo_id: ``"fairydance/molexar-10m-omni"``
+ (https://huggingface.co/fairydance/molexar-10m-omni)
+
+ - SAFE-GPT: GPT-2 causal LM pretrained on 1.1B SAFE strings for de novo
+ generation and scaffold-prefix completion.
+
+ repo_id: ``"datamol-io/safe-gpt"``
+ (https://huggingface.co/datamol-io/safe-gpt)
+
+ Parameters
+ ----------
+ repo_id : str
+ Hugging Face repository id of the pretrained generator.
+ max_length : int, default=128
+ Maximum sequence length used when loading the tokenizer.
+ revision : Optional[str], default=None
+ Model revision on the Hugging Face Hub. NovoMolGen defaults to
+ ``"hf-checkpoint"`` so standard ``model.generate`` works out of the box.
+ trust_remote_code : bool, default=False
+ Whether to trust remote code when loading from Hugging Face.
+ Automatically enabled for Molexar.
+ tokenizer_repo_id : Optional[str], default=None
+ Optional Hugging Face repo for the tokenizer.
+ generate_max_length : int, default=64
+ Default generation length. For causal LMs (NovoMolGen, SAFE-GPT) this is
+ ``max_new_tokens`` so a ``scaffold=`` prefix does not consume the budget.
+ For MolGen this is still decoder ``max_length``.
+ batch_size : int, default=8
+ Batch size used when fine-tuning on SMILES data.
+ epochs : int, default=1
+ Number of fine-tuning epochs when ``fit(X)`` is called.
+ learning_rate : float, default=5e-5
+ Learning rate for fine-tuning.
+ weight_decay : float, default=0.01
+ Weight decay for fine-tuning.
+ grad_norm_clip : Optional[float], default=1.0
+ Maximum gradient norm during fine-tuning. Set to ``None`` to disable clipping.
+ device : Optional[Union[torch.device, str]], default=None
+ Device to run the model on.
+ model_name : str, default="HFPretrainedMolecularGenerator"
+ Name identifier for the model instance.
+ verbose : str, default="none"
+ Progress display mode: ``"none"``, ``"progress_bar"``, or
+ ``"print_statement"``.
+ """
+
+ def __init__(
+ self,
+ repo_id: str,
+ max_length: int = 128,
+ revision: Optional[str] = None,
+ trust_remote_code: bool = False,
+ tokenizer_repo_id: Optional[str] = None,
+ generate_max_length: int = 64,
+ batch_size: int = 8,
+ epochs: int = 1,
+ learning_rate: float = 5e-5,
+ weight_decay: float = 0.01,
+ grad_norm_clip: Optional[float] = 1.0,
+ *,
+ device: Optional[Union[torch.device, str]] = None,
+ model_name: str = "HFPretrainedMolecularGenerator",
+ verbose: str = "none",
+ ):
+ super().__init__(device=device, model_name=model_name, verbose=verbose)
+
+ self.repo_id = repo_id
+ self.max_length = max_length
+ self.revision = revision
+ self.trust_remote_code = trust_remote_code
+ self.tokenizer_repo_id = tokenizer_repo_id
+ self.generate_max_length = generate_max_length
+ self.batch_size = batch_size
+ self.epochs = epochs
+ self.learning_rate = learning_rate
+ self.weight_decay = weight_decay
+ self.grad_norm_clip = grad_norm_clip
+ self.fitting_loss: List[float] = []
+
+ self._family: Optional[str] = None
+ self.tokenizer = None
+ self._safe_tokenizer = None
+ self._molexar_engine = None
+ self._model_local_path: Optional[str] = None
+ self.fitting_epoch = -1
+
+ self._require_transformers()
+
+ if resolve_family(self.repo_id) == "causal_lm":
+ warnings.warn(
+ f"Unknown repo_id: {self.repo_id}. The class will try to load the "
+ "model from Hugging Face as a causal LM, but generation may fail "
+ "if the architecture is not supported.",
+ stacklevel=2,
+ )
+
+ @staticmethod
+ def _get_param_names() -> List[str]:
+ return [
+ "repo_id",
+ "max_length",
+ "revision",
+ "trust_remote_code",
+ "tokenizer_repo_id",
+ "generate_max_length",
+ "batch_size",
+ "epochs",
+ "learning_rate",
+ "weight_decay",
+ "grad_norm_clip",
+ "model_name",
+ ]
+
+ def _get_model_params(self) -> Dict[str, Any]:
+ return {
+ "repo_id": self.repo_id,
+ "max_length": self.max_length,
+ "generate_max_length": self.generate_max_length,
+ }
+
+ def _setup_optimizers(self) -> Tuple[torch.optim.Optimizer, Optional[Any]]:
+ optimizer = torch.optim.AdamW(
+ self.model.parameters(),
+ lr=self.learning_rate,
+ weight_decay=self.weight_decay,
+ )
+ return optimizer, None
+
+ def _train_epoch(self, train_loader, optimizer) -> Dict[str, float]:
+ raise NotImplementedError(
+ "Use fit(X) for fine-tuning HFPretrainedMolecularGenerator."
+ )
+
+ def save_to_local(self, path: str) -> None:
+ """Save the model and tokenizer to a local directory."""
+ self._check_is_fitted()
+ os.makedirs(path, exist_ok=True)
+
+ self.model.save_pretrained(path)
+ self._save_generator_tokenizer(path)
+
+ save_metadata(
+ path,
+ build_metadata(
+ repo_id=self.repo_id,
+ family=self._family,
+ max_length=self.max_length,
+ revision=self.revision,
+ trust_remote_code=self.trust_remote_code,
+ tokenizer_repo_id=self.tokenizer_repo_id,
+ generate_max_length=self.generate_max_length,
+ model_name=self.model_name,
+ ),
+ )
+ self._model_local_path = path
+
+ def load_from_local(self, path: str) -> None:
+ """Load a model and tokenizer saved by :meth:`save_to_local`."""
+ metadata = load_metadata(path)
+ self.repo_id = metadata["repo_id"]
+ self._family = metadata["family"]
+ self.max_length = metadata.get("max_length", self.max_length)
+ self.revision = metadata.get("revision")
+ self.trust_remote_code = metadata.get("trust_remote_code", self.trust_remote_code)
+ self.tokenizer_repo_id = metadata.get("tokenizer_repo_id")
+ self.generate_max_length = metadata.get("generate_max_length", self.generate_max_length)
+ self.model_name = metadata.get("model_name", self.model_name)
+ self._model_local_path = path
+
+ if self._family in MOLEXAR_FAMILIES:
+ self._load_molexar_pretrained(local_path=path)
+ else:
+ self._load_pretrained(local_path=path)
+
+ self.is_fitted_ = True
+
+ def save_to_hf(self, repo_id: str, **kwargs) -> None:
+ raise NotImplementedError(
+ "HFPretrainedMolecularGenerator does not support saving to Hugging Face."
+ )
+
+ def load_from_hf(self, repo_id: Optional[str] = None, **kwargs) -> None:
+ """Load the pretrained model from Hugging Face (same as ``fit()``)."""
+ if repo_id is not None:
+ self.repo_id = repo_id
+ self.fit()
+
+ def load(self, path: Optional[str] = None, repo_id: Optional[str] = None, **kwargs) -> None:
+ """Load the model from a local directory or Hugging Face."""
+ if path is not None:
+ self.load_from_local(path)
+ return
+ if repo_id is not None:
+ self.repo_id = repo_id
+ self.fit()
+
+ def fit(
+ self,
+ X: Optional[List[str]] = None,
+ y: Optional[np.ndarray] = None,
+ ) -> "HFPretrainedMolecularGenerator":
+ """Load the pretrained model from Hugging Face.
+
+ Parameters
+ ----------
+ X : Optional[List[str]], default=None
+ Optional SMILES strings for fine-tuning. When provided, the pretrained
+ weights are adapted on the encoded family-specific representation.
+ y : Optional[np.ndarray], default=None
+ Reserved for future conditional fine-tuning. Currently ignored with a warning.
+
+ Returns
+ -------
+ HFPretrainedMolecularGenerator
+ The fitted generator instance.
+ """
+ assert self.repo_id is not None, "repo_id is not set"
+ self._require_transformers()
+
+ self._family = resolve_family(self.repo_id)
+ self._load_pretrained()
+
+ if X is not None:
+ if y is not None:
+ warnings.warn(
+ "Conditional fine-tuning with y is not implemented yet; continuing with "
+ "unconditional language-model fine-tuning.",
+ stacklevel=2,
+ )
+ y = None
+ X, y = self._validate_inputs(X, y, return_rdkit_mol=False)
+ X = self._encode_inputs(X)
+ self._finetune(X, y)
+
+ self.is_fitted_ = True
+ return self
+
+ def generate(self, n_samples: int = 10, **kwargs) -> List[str]:
+ """Generate molecules as SMILES strings.
+
+ Parameters
+ ----------
+ n_samples : int, default=10
+ Number of molecules to generate.
+ **kwargs
+ Additional arguments forwarded to the family-specific generator.
+ For causal LMs, common options include ``max_new_tokens``,
+ ``temperature``, ``do_sample``, and ``scaffold``. ``max_length`` is
+ still accepted as a Hugging Face total-length override. For MolGen,
+ use ``prefix_selfies`` or ``scaffold`` plus optional ``num_beams``,
+ ``min_length``, and ``max_length``. For Molexar, use ``start_smiles``,
+ ``start_string``, ``generation_task``, or ``conditions`` for omni
+ models. For SAFE-GPT, use ``scaffold=`` with a SMILES prefix
+ (converted to SAFE internally).
+
+ Returns
+ -------
+ List[str]
+ Generated SMILES strings. For MolGen, Molexar, and SAFE-GPT, invalid
+ decodes are dropped, so the list may be shorter than ``n_samples``.
+ """
+ self._check_is_fitted()
+
+ if self._family in CAUSAL_LM_FAMILIES:
+ scaffold = kwargs.pop("scaffold", None)
+ if scaffold is not None and self._family in SAFE_GPT_FAMILIES:
+ from .utils import smiles_to_safe
+
+ scaffold = smiles_to_safe([scaffold])[0]
+ max_new_tokens = kwargs.pop("max_new_tokens", None)
+ max_length = kwargs.pop("max_length", None)
+ if max_new_tokens is not None and max_length is not None:
+ warnings.warn(
+ "Both max_new_tokens and max_length were passed; using max_new_tokens.",
+ stacklevel=2,
+ )
+ max_length = None
+ elif max_new_tokens is None and max_length is None:
+ max_new_tokens = self.generate_max_length
+ raw = generate_causal_lm(
+ self.model,
+ self.tokenizer,
+ self.device,
+ n_samples,
+ family=self._family,
+ max_new_tokens=max_new_tokens,
+ max_length=max_length,
+ temperature=kwargs.pop("temperature", 1.0),
+ do_sample=kwargs.pop("do_sample", True),
+ scaffold=scaffold,
+ **kwargs,
+ )
+ return self._decode_outputs(raw)
+
+ if self._family in SEQ2SEQ_FAMILIES:
+ raw = generate_seq2seq(
+ self.model,
+ self.tokenizer,
+ self.device,
+ n_samples,
+ prefix_selfies=self._resolve_prefix_selfies(kwargs),
+ max_length=kwargs.pop("max_length", self.generate_max_length),
+ min_length=kwargs.pop("min_length", 5),
+ num_beams=kwargs.pop("num_beams", 5),
+ **kwargs,
+ )
+ return self._decode_outputs(raw)
+
+ if self._family in MOLEXAR_FAMILIES:
+ conditions = extract_conditions(kwargs)
+ raw = generate_molexar(
+ self._molexar_engine,
+ n_samples,
+ conditions=conditions or None,
+ max_new_tokens=kwargs.pop("max_new_tokens", None),
+ temperature=kwargs.pop("temperature", 0.8),
+ top_p=kwargs.pop("top_p", 0.95),
+ top_k=kwargs.pop("top_k", 50),
+ do_sample=kwargs.pop("do_sample", True),
+ repetition_penalty=kwargs.pop("repetition_penalty", 1.0),
+ batch_size=kwargs.pop("batch_size", 100),
+ start_string=kwargs.pop("start_string", None),
+ start_smiles=kwargs.pop("start_smiles", None),
+ start_fragment_selfies=kwargs.pop("start_fragment_selfies", None),
+ generation_task=kwargs.pop("generation_task", None),
+ **kwargs,
+ )
+ return self._decode_outputs(raw)
+
+ raise NotImplementedError(f"Generation is not implemented for family '{self._family}'.")
+
+ def _encode_inputs(self, smiles: List[str]) -> List[str]:
+ """Convert SMILES inputs to the representation expected by the model family."""
+ if self._family in SEQ2SEQ_FAMILIES:
+ from .utils import smiles_to_selfies
+
+ return smiles_to_selfies(smiles)
+ if self._family in MOLEXAR_FAMILIES:
+ from .utils import smiles_to_fragment_selfies
+
+ return smiles_to_fragment_selfies(smiles)
+ if self._family in SAFE_GPT_FAMILIES:
+ from .utils import smiles_to_safe
+
+ return smiles_to_safe(smiles)
+ return smiles
+
+ def _resolve_prefix_selfies(self, kwargs: Dict[str, Any]) -> Optional[str]:
+ """Resolve a MolGen SELFIES prefix from kwargs."""
+ prefix_selfies = kwargs.pop("prefix_selfies", None)
+ scaffold = kwargs.pop("scaffold", None)
+
+ if prefix_selfies is not None:
+ return prefix_selfies
+ if scaffold is not None:
+ from .utils import smiles_to_selfies
+
+ return smiles_to_selfies([scaffold])[0]
+ return None
+
+ def _decode_outputs(self, outputs: List[str]) -> List[str]:
+ """Normalize raw model strings to SMILES.
+
+ MolGen, Molexar, and SAFE-GPT drop strings that cannot be decoded. The
+ returned list length is the number of successful SMILES, which may be
+ smaller than ``n_samples``.
+ """
+ n_attempted = len(outputs)
+
+ if self._family in SEQ2SEQ_FAMILIES:
+ from .utils import selfies_to_smiles
+
+ cleaned = [output.replace(" ", "") for output in outputs]
+ smiles = selfies_to_smiles(cleaned)
+ elif self._family in MOLEXAR_FAMILIES:
+ from .utils import fragment_selfies_to_smiles
+
+ smiles = fragment_selfies_to_smiles(outputs)
+ elif self._family in SAFE_GPT_FAMILIES:
+ from .utils import safe_to_smiles
+
+ smiles = safe_to_smiles(outputs)
+ else:
+ return [output.replace(" ", "") for output in outputs]
+
+ if len(smiles) < n_attempted:
+ warnings.warn(
+ f"got {len(smiles)}/{n_attempted} valid SMILES",
+ stacklevel=2,
+ )
+ return smiles
+
+ def _load_pretrained(self, local_path: Optional[str] = None) -> None:
+ import transformers
+
+ if self._family in MOLEXAR_FAMILIES:
+ self._load_molexar_pretrained(local_path=local_path)
+ return
+
+ load_kwargs = self._get_load_kwargs()
+ model_cls = self._get_model_class()
+ model_source = local_path or self.repo_id
+ tokenizer_repo = local_path or self.tokenizer_repo_id or self.repo_id
+
+ if self._family in SAFE_GPT_FAMILIES:
+ self.tokenizer = self._load_safe_gpt_tokenizer(tokenizer_repo)
+ else:
+ self.tokenizer = transformers.AutoTokenizer.from_pretrained(
+ tokenizer_repo,
+ model_max_length=self.max_length,
+ **load_kwargs,
+ )
+ self.model = model_cls.from_pretrained(model_source, **load_kwargs)
+ self._setup_tokenizer()
+ self.model.to(self.device)
+ self.model.eval()
+
+ def _uses_safe_tokenizer(self) -> bool:
+ if self._safe_tokenizer is not None:
+ return True
+ if self._family in SAFE_GPT_FAMILIES:
+ return True
+ if self.repo_id is not None and resolve_family(self.repo_id) in SAFE_GPT_FAMILIES:
+ return True
+ return False
+
+ def _save_generator_tokenizer(self, path: str) -> None:
+ """Save the tokenizer. SAFE's custom pre-tokenizer cannot use HF serialization."""
+ if self._uses_safe_tokenizer():
+ self._save_safe_gpt_tokenizer(path)
+ return
+ try:
+ self.tokenizer.save_pretrained(path)
+ except Exception as exc:
+ if "cannot be serialized" not in str(exc):
+ raise
+ self._save_safe_gpt_tokenizer(path)
+
+ def _save_safe_gpt_tokenizer(self, path: str) -> None:
+ """Save the SAFE tokenizer JSON. The HF fast wrapper cannot be serialized."""
+ from .utils import _require_safe
+
+ _require_safe()
+ from safe.tokenizer import SAFETokenizer
+
+ if self._safe_tokenizer is None:
+ tokenizer_kwargs = {}
+ if self.revision is not None:
+ tokenizer_kwargs["revision"] = self.revision
+ self._safe_tokenizer = SAFETokenizer.from_pretrained(
+ self.tokenizer_repo_id or self.repo_id,
+ **tokenizer_kwargs,
+ )
+ self._safe_tokenizer.save_pretrained(path)
+
+ def _load_safe_gpt_tokenizer(self, tokenizer_repo: str):
+ """Load the custom SAFE tokenizer as a Hugging Face fast tokenizer."""
+ from .utils import _require_safe
+
+ _require_safe()
+ from safe.tokenizer import SAFETokenizer
+
+ tokenizer_kwargs = {}
+ if self.revision is not None:
+ tokenizer_kwargs["revision"] = self.revision
+ self._safe_tokenizer = SAFETokenizer.from_pretrained(tokenizer_repo, **tokenizer_kwargs)
+ tokenizer = self._safe_tokenizer.get_pretrained()
+ tokenizer.model_max_length = self.max_length
+ return tokenizer
+
+ def _load_molexar_pretrained(self, local_path: Optional[str] = None) -> None:
+ from huggingface_hub import snapshot_download
+
+ from .families.molexar import _require_molexar
+
+ _require_molexar()
+ from molexar.inference import MolexarInference
+
+ if local_path is None:
+ self._model_local_path = snapshot_download(self.repo_id)
+ else:
+ self._model_local_path = local_path
+
+ device = str(self.device)
+ self._molexar_engine = MolexarInference(
+ self._model_local_path,
+ device=device,
+ tokenizer_path=self.tokenizer_repo_id,
+ )
+ self.model = self._molexar_engine.model
+ self.tokenizer = self._molexar_engine.tokenizer
+
+ def _get_molexar_config(self):
+ if self._molexar_engine is not None:
+ return self._molexar_engine.config
+ return getattr(self.model, "config", None)
+
+ def _finetune(self, X: List[str], y: Optional[np.ndarray]) -> None:
+ if len(X) == 0:
+ raise ValueError("Fine-tuning requires at least one training example.")
+
+ config = self._get_molexar_config() if self._family in MOLEXAR_FAMILIES else None
+ losses, last_epoch = finetune_generator(
+ self._family,
+ self.model,
+ self.tokenizer,
+ X,
+ self.device,
+ config=config,
+ max_length=self.max_length,
+ batch_size=self.batch_size,
+ epochs=self.epochs,
+ learning_rate=self.learning_rate,
+ weight_decay=self.weight_decay,
+ grad_norm_clip=self.grad_norm_clip,
+ verbose=self.verbose,
+ )
+ self.fitting_loss = losses
+ self.fitting_epoch = last_epoch
+ self.model.eval()
+
+ def _get_model_class(self):
+ import transformers
+
+ if self._family in SEQ2SEQ_FAMILIES:
+ return transformers.AutoModelForSeq2SeqLM
+ if self._family in SAFE_GPT_FAMILIES:
+ # Hub config lists SAFEDoubleHeadsModel; the LM head is standard GPT-2.
+ return transformers.GPT2LMHeadModel
+ return transformers.AutoModelForCausalLM
+
+ def _get_load_kwargs(self) -> Dict[str, Any]:
+ load_kwargs: Dict[str, Any] = {}
+
+ if self._family == "novomolgen" and self.revision is None:
+ load_kwargs["revision"] = "hf-checkpoint"
+ elif self.revision is not None:
+ load_kwargs["revision"] = self.revision
+
+ if self._family in MOLEXAR_FAMILIES or self.trust_remote_code:
+ load_kwargs["trust_remote_code"] = True
+
+ if self._family in SAFE_GPT_FAMILIES:
+ # Extra property-prediction head in the checkpoint is unused.
+ load_kwargs["ignore_mismatched_sizes"] = True
+
+ return load_kwargs
+
+ def _setup_tokenizer(self) -> None:
+ if self.tokenizer.pad_token is None:
+ if self.tokenizer.eos_token is not None:
+ self.tokenizer.pad_token = self.tokenizer.eos_token
+ else:
+ self.tokenizer.add_special_tokens({"pad_token": ""})
+ self.model.resize_token_embeddings(len(self.tokenizer))
+
+ if self._family in SAFE_GPT_FAMILIES:
+ config = getattr(self.model, "config", None)
+ if self.tokenizer.bos_token_id is None and getattr(config, "bos_token_id", None) is not None:
+ self.tokenizer.bos_token_id = config.bos_token_id
+ if self.tokenizer.eos_token_id is None and getattr(config, "eos_token_id", None) is not None:
+ self.tokenizer.eos_token_id = config.eos_token_id
+ if self.tokenizer.pad_token_id is None and getattr(config, "pad_token_id", None) is not None:
+ self.tokenizer.pad_token_id = config.pad_token_id
+
+ @staticmethod
+ def _require_transformers() -> None:
+ try:
+ import transformers # noqa: F401
+ except ImportError as exc:
+ raise ImportError(
+ "The 'transformers' package is required for HFPretrainedMolecularGenerator. "
+ "Please install it using `pip install transformers`."
+ ) from exc
diff --git a/torch_molecule/generator/pretrained/registry.py b/torch_molecule/generator/pretrained/registry.py
new file mode 100644
index 0000000..c66919f
--- /dev/null
+++ b/torch_molecule/generator/pretrained/registry.py
@@ -0,0 +1,42 @@
+"""Family registry for Hugging Face pretrained molecular generators."""
+
+from typing import Dict
+
+KNOWN_REPOS: Dict[str, str] = {
+ "chandar-lab/NovoMolGen_32M_SMILES_BPE": "novomolgen",
+ "zjunlp/MolGen-large": "molgen",
+ "zjunlp/MolGen-large-opt": "molgen",
+ "fairydance/molexar-10m-base": "molexar",
+ "fairydance/molexar-10m-omni": "molexar",
+ "datamol-io/safe-gpt": "safe_gpt",
+}
+
+FAMILY_PREFIXES: Dict[str, str] = {
+ "chandar-lab/NovoMolGen": "novomolgen",
+ "zjunlp/MolGen": "molgen",
+ "fairydance/molexar": "molexar",
+ "datamol-io/safe": "safe_gpt",
+}
+
+MOLEXAR_OMNI_REPOS = frozenset(
+ {
+ "fairydance/molexar-10m-omni",
+ }
+)
+
+CAUSAL_LM_FAMILIES = frozenset({"novomolgen", "safe_gpt", "causal_lm"})
+SEQ2SEQ_FAMILIES = frozenset({"molgen"})
+MOLEXAR_FAMILIES = frozenset({"molexar"})
+SAFE_GPT_FAMILIES = frozenset({"safe_gpt"})
+
+
+def resolve_family(repo_id: str) -> str:
+ """Map a Hugging Face repo id to a generator family."""
+ if repo_id in KNOWN_REPOS:
+ return KNOWN_REPOS[repo_id]
+
+ for prefix, family in FAMILY_PREFIXES.items():
+ if repo_id.startswith(prefix):
+ return family
+
+ return "causal_lm"
diff --git a/torch_molecule/generator/pretrained/utils.py b/torch_molecule/generator/pretrained/utils.py
new file mode 100644
index 0000000..c480dc4
--- /dev/null
+++ b/torch_molecule/generator/pretrained/utils.py
@@ -0,0 +1,268 @@
+"""Representation conversion utilities for HF pretrained generators."""
+
+import warnings
+from typing import List
+
+from rdkit import Chem
+
+
+def _require_selfies():
+ try:
+ import selfies # noqa: F401
+ except ImportError as exc:
+ raise ImportError(
+ "The 'selfies' package is required for SELFIES conversion. "
+ "Install it with `pip install selfies`."
+ ) from exc
+
+
+def smiles_to_selfies(smiles: List[str]) -> List[str]:
+ """Convert SMILES strings to SELFIES representations.
+
+ Parameters
+ ----------
+ smiles : List[str]
+ Input SMILES strings.
+
+ Returns
+ -------
+ List[str]
+ SELFIES strings in the same order as the input.
+
+ Raises
+ ------
+ ValueError
+ If a SMILES string is invalid or cannot be encoded as SELFIES.
+ """
+ _require_selfies()
+ import selfies as sf
+
+ selfies_list: List[str] = []
+ for idx, smiles_string in enumerate(smiles):
+ mol = Chem.MolFromSmiles(smiles_string)
+ if mol is None:
+ raise ValueError(f"Invalid SMILES at index {idx}: {smiles_string}")
+ canonical = Chem.MolToSmiles(mol)
+ try:
+ selfies_list.append(sf.encoder(canonical))
+ except sf.EncoderError as exc:
+ raise ValueError(
+ f"SMILES at index {idx} is RDKit-valid but not SELFIES-encodable: {smiles_string}"
+ ) from exc
+ return selfies_list
+
+
+def selfies_to_smiles(selfies_list: List[str]) -> List[str]:
+ """Convert SELFIES strings to canonical SMILES.
+
+ Entries that cannot be decoded are dropped and a warning is emitted.
+ Unexpected errors such as ``ImportError`` are not swallowed.
+
+ Parameters
+ ----------
+ selfies_list : List[str]
+ Input SELFIES strings.
+
+ Returns
+ -------
+ List[str]
+ Canonical SMILES strings for entries that decoded successfully.
+ """
+ _require_selfies()
+ import selfies as sf
+
+ smiles_list: List[str] = []
+ n_dropped = 0
+ for selfies_string in selfies_list:
+ if not selfies_string or not str(selfies_string).strip():
+ n_dropped += 1
+ continue
+ try:
+ decoded = sf.decoder(selfies_string)
+ except sf.DecoderError:
+ n_dropped += 1
+ continue
+ mol = Chem.MolFromSmiles(decoded) if decoded else None
+ if mol is None:
+ n_dropped += 1
+ continue
+ smiles_list.append(Chem.MolToSmiles(mol))
+
+ if n_dropped:
+ warnings.warn(f"dropped {n_dropped} invalid SELFIES", stacklevel=2)
+ return smiles_list
+
+
+def _require_fragment_selfies():
+ try:
+ from fragment_selfies import FragmentSelfiesCodec # noqa: F401
+ except ImportError as exc:
+ raise ImportError(
+ "The 'fragment-selfies' package is required for Fragment-SELFIES conversion. "
+ "Install it with `pip install fragment-selfies`."
+ ) from exc
+
+
+def smiles_to_fragment_selfies(smiles: List[str]) -> List[str]:
+ """Convert SMILES strings to Fragment-SELFIES representations."""
+ _require_fragment_selfies()
+
+ try:
+ from molexar.data.converter import smiles_to_fragment_selfies as encode_one
+ except ImportError as exc:
+ raise ImportError(
+ "Molexar conversion utilities require the 'molexar' package. "
+ "Install it with `pip install git+https://github.com/fairydance/Molexar.git`."
+ ) from exc
+
+ return [encode_one(smiles_string, canonical=True) for smiles_string in smiles]
+
+
+def fragment_selfies_to_smiles(fragment_selfies_list: List[str]) -> List[str]:
+ """Convert Fragment-SELFIES strings to canonical SMILES.
+
+ Entries that cannot be decoded are dropped and a warning is emitted.
+ Unexpected errors such as ``ImportError`` are not swallowed.
+ """
+ _require_fragment_selfies()
+
+ try:
+ from molexar.data.converter import fragment_selfies_to_smiles as decode_one
+ except ImportError as exc:
+ raise ImportError(
+ "Molexar conversion utilities require the 'molexar' package. "
+ "Install it with `pip install git+https://github.com/fairydance/Molexar.git`."
+ ) from exc
+
+ smiles_list: List[str] = []
+ n_dropped = 0
+ for fragment_selfies in fragment_selfies_list:
+ if not fragment_selfies or not str(fragment_selfies).strip():
+ n_dropped += 1
+ continue
+ try:
+ decoded = decode_one(fragment_selfies, canonical=True, ignore_errors=False)
+ except (ValueError, TypeError):
+ n_dropped += 1
+ continue
+ if not decoded:
+ n_dropped += 1
+ continue
+ smiles_list.append(decoded)
+
+ if n_dropped:
+ warnings.warn(f"dropped {n_dropped} invalid Fragment-SELFIES", stacklevel=2)
+ return smiles_list
+
+
+def _require_safe():
+ from .compat import ensure_safe_transformers_compat
+
+ ensure_safe_transformers_compat()
+ try:
+ from safe.converter import encode # noqa: F401
+ except ImportError as exc:
+ raise ImportError(
+ "The 'safe-mol' package is required for SAFE-GPT conversion. "
+ "Install it with `pip install safe-mol`."
+ ) from exc
+
+
+def smiles_to_safe(smiles: List[str]) -> List[str]:
+ """Convert SMILES strings to SAFE representations.
+
+ Parameters
+ ----------
+ smiles : List[str]
+ Input SMILES strings.
+
+ Returns
+ -------
+ List[str]
+ SAFE strings in the same order as the input.
+
+ Raises
+ ------
+ ValueError
+ If a SMILES string is invalid or cannot be encoded as SAFE.
+ """
+ _require_safe()
+ try:
+ from safe._exception import SAFEEncodeError, SAFEFragmentationError
+ except ImportError:
+ from safe.converter import SAFEEncodeError, SAFEFragmentationError # type: ignore
+ from safe.converter import encode
+
+ safe_list: List[str] = []
+ for idx, smiles_string in enumerate(smiles):
+ mol = Chem.MolFromSmiles(smiles_string)
+ if mol is None:
+ raise ValueError(f"Invalid SMILES at index {idx}: {smiles_string}")
+ canonical = Chem.MolToSmiles(mol)
+ try:
+ try:
+ encoded = encode(canonical, canonical=True, allow_empty=True)
+ except TypeError:
+ encoded = encode(canonical, canonical=True)
+ except (SAFEEncodeError, SAFEFragmentationError) as exc:
+ raise ValueError(
+ f"SMILES at index {idx} is RDKit-valid but not SAFE-encodable: {smiles_string}"
+ ) from exc
+ if not encoded:
+ raise ValueError(
+ f"SMILES at index {idx} is RDKit-valid but not SAFE-encodable: {smiles_string}"
+ )
+ safe_list.append(encoded)
+ return safe_list
+
+
+def safe_to_smiles(safe_list: List[str]) -> List[str]:
+ """Convert SAFE strings to canonical SMILES.
+
+ Entries that cannot be decoded are dropped and a warning is emitted.
+ Unexpected errors such as ``ImportError`` are not swallowed.
+
+ Parameters
+ ----------
+ safe_list : List[str]
+ Input SAFE strings.
+
+ Returns
+ -------
+ List[str]
+ Canonical SMILES strings for entries that decoded successfully.
+ """
+ _require_safe()
+ try:
+ from safe._exception import SAFEDecodeError
+ except ImportError:
+ from safe.converter import SAFEDecodeError # type: ignore
+ from safe.converter import decode
+
+ smiles_list: List[str] = []
+ n_dropped = 0
+ for safe_string in safe_list:
+ if not safe_string or not str(safe_string).strip():
+ n_dropped += 1
+ continue
+ try:
+ try:
+ decoded = decode(
+ str(safe_string).replace(" ", ""),
+ canonical=True,
+ ignore_errors=False,
+ )
+ except TypeError:
+ decoded = decode(str(safe_string).replace(" ", ""))
+ except SAFEDecodeError:
+ n_dropped += 1
+ continue
+ mol = Chem.MolFromSmiles(decoded) if decoded else None
+ if mol is None:
+ n_dropped += 1
+ continue
+ smiles_list.append(Chem.MolToSmiles(mol))
+
+ if n_dropped:
+ warnings.warn(f"dropped {n_dropped} invalid SAFE", stacklevel=2)
+ return smiles_list