Materials.smi Ted

🚀 IBM材料基础模型介绍

IBM推出了一系列用于可持续材料的大型基础模型。这些模型涵盖多种表示形式和模态，包括SMILES、SELFIES、3D原子位置、3D密度网格、分子图等格式，旨在支持和推动材料科学与化学领域的研究。

GitHub: GitHub链接

论文: arXiv:2407.20267

✨ 主要特性

• 模型类型：基于SMILES的Transformer编解码器（SMI - TED）
• 评估指标：准确率
• 支持任务：特征提取、量子属性预测等多种复杂任务
• 模型权重格式：PyTorch（.pt）和safetensors（.safetensors）

📦 安装指南

复制Conda环境

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

创建并激活Conda环境

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

使用Conda安装包

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

使用Pip安装包

pip install -r requirements.txt
pip install pytorch - fast - transformers

💻 使用示例

基础用法

加载模型

model = load_smi_ted(
    folder='../inference/smi_ted_light',
    ckpt_filename='smi_ted_light.pt'
)

或

with open('model_weights.bin', 'rb') as f:
    state_dict = torch.load(f)
    model.load_state_dict(state_dict)

编码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)

📚 详细文档

预训练模型和训练日志

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

根据需要将SMI - TED的预训练权重.pt添加到inference/或finetune/目录中。目录结构应如下所示：

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

和/或：

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

预训练

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

SMI - TED在来自PubChem的经过规范化和精选的9100万个SMILES上进行预训练，有以下约束条件：

• 预处理期间，化合物过滤为最大长度202个标记。
• 编码器训练使用95/5/0分割，5%的数据用于解码器预训练。
• 也使用100/0/0分割直接训练编码器和解码器，以提高模型性能。

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

要预训练SMI - TED模型的两个变体，请运行：

bash training/run_model_light_training.sh

或

bash training/run_model_large_training.sh

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

微调

微调数据集和环境可在[finetune](https://github.com/IBM/materials/tree/main/smi - ted/finetune)目录中找到。设置好环境后，您可以使用以下命令运行微调任务：

bash finetune/smi_ted_light/esol/run_finetune_esol.sh

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

特征提取

示例笔记本[smi_ted_encoder_decoder_example.ipynb](https://github.com/IBM/materials/blob/main/smi - ted/notebooks/smi_ted_encoder_decoder_example.ipynb)包含加载检查点文件并使用预训练模型进行编码器和解码器任务的代码，还包括分类和回归任务的示例。

🔧 技术细节

SMI - TED是一个大型编解码器化学基础模型，在从PubChem获取的9100万个SMILES样本的精选数据集上进行预训练，相当于40亿个分子标记。它支持各种复杂任务，包括量子属性预测，有两种主要变体（2.89亿和8X2.89亿）。我们在多个基准数据集上的实验证明了其在各种任务上的先进性能。

📄 许可证

本项目采用Apache - 2.0许可证。

📚 引用

@misc{soares2024largeencoderdecoderfamilyfoundation,
      title={A Large Encoder-Decoder Family of Foundation Models For Chemical Language}, 
      author={Eduardo Soares and Victor Shirasuna and Emilio Vital Brazil and Renato Cerqueira and Dmitry Zubarev and Kristin Schmidt},
      year={2024},
      eprint={2407.20267},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2407.20267}, 
}

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