license: apache-2.0 metrics:

• accuracy pipeline_tag: feature-extraction tags:
• chemistry
• foundation models
• AI4Science
• materials
• molecules
• safetensors
• pytorch
• transformer
• diffusers library_name: transformers

基于SMILES的状态空间编码器-解码器（SMI-SSED）- MoLMamba

本代码库提供了与我们的论文《基于Mamba的化学基础模型》相关的PyTorch源代码。

论文NeurIPS AI4Mat 2024: Arxiv链接

更多信息请联系：eduardo.soares@ibm.com 或 evital@br.ibm.com。

简介

我们提出了一种基于Mamba的编码器-解码器化学基础模型——基于SMILES的状态空间编码器-解码器（SMI-SSED），该模型在从PubChem收集的9100万个SMILES样本（相当于40亿个分子标记）上进行了预训练。SMI-SSED支持多种复杂任务，包括量子属性预测，并提供了两种主要变体（$336$和$8 \times 336M$）。我们在多个基准数据集上的实验表明，该模型在各种任务中均达到了最先进的性能。

我们提供了两种格式的模型权重：

• PyTorch（.pt）：smi_ssed_130.pt
• safetensors（.bin）：smi_ssed_130.bin

更多信息请联系：eduardo.soares@ibm.com 或 evital@br.ibm.com。

目录

1. 快速开始
   1. 预训练模型与训练日志
   2. 复制Conda环境
2. 预训练
3. 微调
4. 特征提取

快速开始

本代码及环境已在Nvidia V100和Nvidia A100上测试通过

预训练模型与训练日志

我们提供了SMI-SSED模型的检查点，该模型在从PubChem收集的约9100万个分子数据集上进行了预训练。预训练模型在MoleculeNet的分类和回归基准测试中表现出色。

将SMI-SSED的预训练权重.pt添加到inference/或finetune/目录中，具体根据需求而定。目录结构应如下所示：

inference/
├── smi_ssed
│   ├── smi_ssed.pt
│   ├── bert_vocab_curated.txt
│   └── load.py

和/或：

finetune/
├── smi_ssed
│   ├── smi_ssed.pt
│   ├── bert_vocab_curated.txt
│   └── load.py

复制Conda环境

按照以下步骤复制我们的Conda环境并安装必要的库：

创建并激活Conda环境

conda create --name smi-ssed-env python=3.10
conda activate smi-ssed-env

使用Conda安装包

conda install pytorch=2.1.0 pytorch-cuda=11.8 -c pytorch -c nvidia

使用Pip安装包

pip install -r requirements.txt

预训练

在预训练阶段，我们采用了两种策略：使用掩码语言模型方法训练编码器部分，以及使用编码器-解码器策略优化SMILES重构并改进生成的潜在空间。

SMI-SSED在PubChem的9100万个经过规范化和筛选的SMILES上进行预训练，并遵循以下约束：

• 在预处理过程中，化合物的最大长度限制为202个标记。
• 使用95/5/0的数据划分进行编码器训练，其中5%的数据用于解码器预训练。
• 还使用100/0/0的数据划分直接训练编码器和解码器，以提升模型性能。

预训练代码提供了在较小数据集上进行数据处理和模型训练的示例，需要8块A100 GPU。

运行以下命令进行SMI-SSED模型的预训练：

bash training/run_model_training.sh

使用train_model_D.py仅训练解码器，或使用train_model_ED.py同时训练编码器和解码器。

微调

微调数据集和环境可在finetune目录中找到。设置好环境后，可以通过以下命令运行微调任务：

bash finetune/smi_ssed/esol/run_finetune_esol.sh

微调训练/检查点资源将保存在名为checkpoint_<measure_name>的目录中。

特征提取

示例笔记本smi_ssed_encoder_decoder_example.ipynb包含加载检查点文件和使用预训练模型进行编码器与解码器任务的代码。其中还包括分类和回归任务的示例。模型权重：HuggingFace链接

加载smi-ssed的代码如下：

model = load_smi_ssed(
    folder='../inference/smi_ssed',
    ckpt_filename='smi_ssed.pt'
)

将SMILES编码为嵌入向量的代码如下：

with torch.no_grad():
    encoded_embeddings = model.encode(df['SMILES'], return_torch=True)

对于解码器，可以使用以下函数从嵌入向量还原为SMILES字符串：

with torch.no_grad():
    decoded_smiles = model.decode(encoded_embeddings)