ChemMRL
Collection
SMILES Matryoshka Representation Learning Embedding Transformer
•
5 items
•
Updated
•
1
This is a Chem-MRL (sentence-transformers) model finetuned from Derify/ModChemBERT-IR-BASE on the pubchem_10m_genmol_similarity dataset. It maps SMILES to a 1024-dimensional dense vector space and can be used for molecular similarity, semantic search, database indexing, molecular classification, clustering, and more.
SentenceTransformer(
(0): Transformer({'max_seq_length': 512, 'do_lower_case': False, 'architecture': 'ModChemBertModel'})
(1): Pooling({'word_embedding_dimension': 1024, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
(2): Normalize()
)
First install the Chem-MRL library:
pip install -U chem-mrl>=0.7.3
Then you can load this model and run inference.
from chem_mrl import ChemMRL
# Download from the 🤗 Hub
model = ChemMRL(
"Derify/ChemMRL",
trust_remote_code=True,
model_kwargs={"torch_dtype": "bfloat16"},
)
# Run inference
sentences = [
'OCCCc1cc(F)cc(F)c1',
'Fc1cc(F)cc(-n2cc[o+]n2)c1',
'CCC(C)C(=O)C1(C(NN)C(C)C)CCCC1',
]
embeddings = model.backbone.encode(sentences)
print(embeddings.shape)
# [3, 1024]
# Get the similarity scores for the embeddings
similarities = model.backbone.similarity(embeddings, embeddings)
print(similarities)
# tensor([[1.0000, 0.3876, 0.0078],
# [0.3876, 1.0000, 0.0028],
# [0.0078, 0.0028, 1.0000]])
First install the Sentence Transformers library:
pip install -U sentence-transformers
Then you can load this model and run inference.
from sentence_transformers import SentenceTransformer
# Download from the 🤗 Hub
model = SentenceTransformer(
"Derify/ChemMRL",
# SentenceTransformer doesn't support tanimoto similarity natively so we set a different similarity function here
similarity_fn_name="cosine",
trust_remote_code=True,
model_kwargs={"torch_dtype": "bfloat16"},
)
# Run inference
sentences = [
'OCCCc1cc(F)cc(F)c1',
'Fc1cc(F)cc(-n2cc[o+]n2)c1',
'CCC(C)C(=O)C1(C(NN)C(C)C)CCCC1',
]
embeddings = model.encode(sentences)
print(embeddings.shape)
# [3, 1024]
# Get the similarity scores for the embeddings
similarities = model.similarity(embeddings, embeddings)
print(similarities)
# tensor([[1.0000, 0.5587, 0.0155],
# [0.5587, 1.0000, 0.0055],
# [0.0155, 0.0055, 1.0000]])
pubchem_10m_genmol_similaritychem_mrl.evaluation.EmbeddingSimilarityEvaluator.EmbeddingSimilarityEvaluator with these parameters:{
"precision": "float32"
}
| Split | Metric | Value |
|---|---|---|
| validation | spearman | 0.98914 |
| test | spearman | 0.98916 |
smiles_a, smiles_b, and label| smiles_a | smiles_b | label | |
|---|---|---|---|
| type | string | string | float |
| details |
|
|
|
| smiles_a | smiles_b | label |
|---|---|---|
COc1ccc(NC(=O)C2CC[NH+](C(C)C(=O)Nc3ccc(C(=O)Nc4ccc(F)c(F)c4)cc3C)CC2)cc1NC(=O)C1CCCCC1 |
Cc1cc(C(=O)Nc2ccc(F)c(F)c2)ccc1NC(=O)C(C)[NH+]1CCC(C(=O)Nc2cccc(NC(=O)C3CCCCC3)c2)CC1 |
0.8495575189590454 |
OCCN1CC[NH+](Cc2ccccc2OC2CC2)CC1 |
OCCN1CC[NH+](Cc2ccccc2On2cccn2)CC1 |
0.6615384817123413 |
CC1CN(C(=O)C2CC[NH+](Cc3cccc(C(N)=O)c3)CC2)CC(C)O1 |
CC1CN(C(=O)C2CC[NH+](Cc3ccccc3)CC2)CC(C)O1 |
0.7123287916183472 |
Matryoshka2dLoss with these parameters:{
"loss": "TanimotoSentLoss",
"n_layers_per_step": -1,
"last_layer_weight": 2.0,
"prior_layers_weight": 1.0,
"kl_div_weight": 0.0,
"kl_temperature": 0.0,
"matryoshka_dims": [
1024,
512,
256,
128,
64,
32,
16,
8
],
"matryoshka_weights": [
1,
1,
1,
1,
1,
1,
1,
1
],
"n_dims_per_step": -1
}
smiles_a, smiles_b, and label| smiles_a | smiles_b | label | |
|---|---|---|---|
| type | string | string | float |
| details |
|
|
|
| smiles_a | smiles_b | label |
|---|---|---|
N#CCCN(Cc1cnc(N)cn1)C1CC1 |
N#CCCN(Cc1cnc(N)cn1)C1CCCC1 |
0.8600000143051147 |
N#CCCN(Cc1cnc(N)cn1)C1CC1 |
N#CCCN(Cc1cnc(N)cn1)C1CCOCC1 |
0.7962962985038757 |
N#CCCN(Cc1cnc(N)cn1)C1CC1 |
N#CCCN(Cc1cnc(N)cn1)CC(F)F |
0.5517241358757019 |
Matryoshka2dLoss with these parameters:{
"loss": "TanimotoSentLoss",
"n_layers_per_step": -1,
"last_layer_weight": 2.0,
"prior_layers_weight": 1.0,
"kl_div_weight": 0.0,
"kl_temperature": 0.0,
"matryoshka_dims": [
1024,
512,
256,
128,
64,
32,
16,
8
],
"matryoshka_weights": [
1,
1,
1,
1,
1,
1,
1,
1
],
"n_dims_per_step": -1
}
eval_strategy: stepsper_device_train_batch_size: 192per_device_eval_batch_size: 512learning_rate: 8e-06weight_decay: 1e-05max_grad_norm: Nonelr_scheduler_type: warmup_stable_decaylr_scheduler_kwargs: {'num_decay_steps': 100943, 'warmup_type': 'linear', 'decay_type': '1-sqrt'}warmup_steps: 100943data_seed: 42bf16: Truebf16_full_eval: Truetf32: Trueoptim: stable_adamwoptim_args: decouple_lr=True,max_lr=8.0e-6gradient_checkpointing: Trueeval_on_start: Trueoverwrite_output_dir: Falsedo_predict: Falseeval_strategy: stepsprediction_loss_only: Trueper_device_train_batch_size: 192per_device_eval_batch_size: 512per_gpu_train_batch_size: Noneper_gpu_eval_batch_size: Nonegradient_accumulation_steps: 1eval_accumulation_steps: Nonetorch_empty_cache_steps: Nonelearning_rate: 8e-06weight_decay: 1e-05adam_beta1: 0.9adam_beta2: 0.999adam_epsilon: 1e-08max_grad_norm: Nonenum_train_epochs: 3max_steps: -1lr_scheduler_type: warmup_stable_decaylr_scheduler_kwargs: {'num_decay_steps': 100943, 'warmup_type': 'linear', 'decay_type': '1-sqrt'}warmup_ratio: 0.0warmup_steps: 100943log_level: passivelog_level_replica: warninglog_on_each_node: Truelogging_nan_inf_filter: Truesave_safetensors: Truesave_on_each_node: Falsesave_only_model: Falserestore_callback_states_from_checkpoint: Falseno_cuda: Falseuse_cpu: Falseuse_mps_device: Falseseed: 42data_seed: 42jit_mode_eval: Falsebf16: Truefp16: Falsefp16_opt_level: O1half_precision_backend: autobf16_full_eval: Truefp16_full_eval: Falsetf32: Truelocal_rank: 0ddp_backend: Nonetpu_num_cores: Nonetpu_metrics_debug: Falsedebug: []dataloader_drop_last: Falsedataloader_num_workers: 0dataloader_prefetch_factor: Nonepast_index: -1disable_tqdm: Falseremove_unused_columns: Truelabel_names: Noneload_best_model_at_end: Falseignore_data_skip: Falsefsdp: []fsdp_min_num_params: 0fsdp_config: {'min_num_params': 0, 'xla': False, 'xla_fsdp_v2': False, 'xla_fsdp_grad_ckpt': False}fsdp_transformer_layer_cls_to_wrap: Noneaccelerator_config: {'split_batches': False, 'dispatch_batches': None, 'even_batches': True, 'use_seedable_sampler': True, 'non_blocking': False, 'gradient_accumulation_kwargs': None}parallelism_config: Nonedeepspeed: Nonelabel_smoothing_factor: 0.0optim: stable_adamwoptim_args: decouple_lr=True,max_lr=8.0e-6adafactor: Falsegroup_by_length: Falselength_column_name: lengthproject: huggingfacetrackio_space_id: trackioddp_find_unused_parameters: Noneddp_bucket_cap_mb: Noneddp_broadcast_buffers: Falsedataloader_pin_memory: Truedataloader_persistent_workers: Falseskip_memory_metrics: Trueuse_legacy_prediction_loop: Falsepush_to_hub: Falseresume_from_checkpoint: Nonehub_model_id: Nonehub_strategy: every_savehub_private_repo: Nonehub_always_push: Falsehub_revision: Nonegradient_checkpointing: Truegradient_checkpointing_kwargs: Noneinclude_inputs_for_metrics: Falseinclude_for_metrics: []eval_do_concat_batches: Truefp16_backend: autopush_to_hub_model_id: Nonepush_to_hub_organization: Nonemp_parameters: auto_find_batch_size: Falsefull_determinism: Falsetorchdynamo: Noneray_scope: lastddp_timeout: 1800torch_compile: Falsetorch_compile_backend: Nonetorch_compile_mode: Noneinclude_tokens_per_second: Falseinclude_num_input_tokens_seen: noneftune_noise_alpha: Noneoptim_target_modules: Nonebatch_eval_metrics: Falseeval_on_start: Trueuse_liger_kernel: Falseliger_kernel_config: Noneeval_use_gather_object: Falseaverage_tokens_across_devices: Trueprompts: Nonebatch_sampler: batch_samplermulti_dataset_batch_sampler: proportionalrouter_mapping: {}learning_rate_mapping: {}| Epoch | Step | Training Loss | pubchem 10m genmol similarity loss | pubchem_10m_genmol_similarity_spearman |
|---|---|---|---|---|
| 0 | 0 | - | 297.6136 | 0.7261 |
| 0.0000 | 1 | 244.6862 | - | - |
| 0.2477 | 25000 | 161.5037 | - | - |
| 0.2500 | 25235 | - | 195.4624 | 0.9067 |
| 0.4978 | 50250 | 155.7822 | - | - |
| 0.5000 | 50470 | - | 189.4068 | 0.9655 |
| 0.7479 | 75500 | 152.7915 | - | - |
| 0.7500 | 75705 | - | 186.3661 | 0.9780 |
| 0.9981 | 100750 | 151.0411 | - | - |
| 1.0000 | 100940 | - | 184.6362 | 0.9829 |
| 1.2482 | 126000 | 149.8544 | - | - |
| 1.2500 | 126175 | - | 183.5648 | 0.9855 |
| 1.4984 | 151250 | 149.2916 | - | - |
| 1.5000 | 151410 | - | 182.8947 | 0.9868 |
| 1.7485 | 176500 | 148.7942 | - | - |
| 1.7499 | 176645 | - | 182.3662 | 0.9879 |
| 1.9987 | 201750 | 148.3459 | - | - |
| 1.9999 | 201880 | - | 181.9855 | 0.9885 |
| 2.2488 | 227000 | 148.0316 | - | - |
| 2.2499 | 227115 | - | 181.7683 | 0.9889 |
| 2.4989 | 252250 | 147.8658 | - | - |
| 2.4999 | 252350 | - | 181.6711 | 0.9890 |
| 2.7491 | 277500 | 147.9642 | - | - |
| 2.7499 | 277585 | - | 181.6077 | 0.9891 |
| 2.9992 | 302750 | 147.8874 | - | - |
| 2.9999 | 302820 | - | 181.6066 | 0.9891 |
| 3.0000 | 302829 | - | - | 0.98914 |
Carbon emissions were measured using CodeCarbon.
@inproceedings{reimers-2019-sentence-bert,
title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
author = "Reimers, Nils and Gurevych, Iryna",
booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
month = "11",
year = "2019",
publisher = "Association for Computational Linguistics",
url = "https://arxiv.org/abs/1908.10084",
}
@misc{li20242d,
title={2D Matryoshka Sentence Embeddings},
author={Xianming Li and Zongxi Li and Jing Li and Haoran Xie and Qing Li},
year={2024},
eprint={2402.14776},
archivePrefix={arXiv},
primaryClass={cs.CL}
}
@misc{kusupati2024matryoshka,
title={Matryoshka Representation Learning},
author={Aditya Kusupati and Gantavya Bhatt and Aniket Rege and Matthew Wallingford and Aditya Sinha and Vivek Ramanujan and William Howard-Snyder and Kaifeng Chen and Sham Kakade and Prateek Jain and Ali Farhadi},
year={2024},
eprint={2205.13147},
archivePrefix={arXiv},
primaryClass={cs.LG}
}
@online{kexuefm-8847,
title={CoSENT: A more efficient sentence vector scheme than Sentence-BERT},
author={Su Jianlin},
year={2022},
month={Jan},
url={https://kexue.fm/archives/8847},
}
@online{cortes-2025-tanimotosentloss,
title={TanimotoSentLoss: Tanimoto Loss for SMILES Embeddings},
author={Emmanuel Cortes},
year={2025},
month={Jan},
url={https://github.com/emapco/chem-mrl},
}
Manny Cortes ([email protected])
Base model
Derify/ModChemBERT-IR-BASE