license: apache-2.0 metrics:

• accuracy pipeline_tag: feature-extraction tags:
• 化学
• 基础模型
• 科学人工智能
• 材料
• 分子
• 安全张量格式
• PyTorch
• 变压器架构
• 扩散模型 library_name: transformers

IBM材料科学基础模型导览

欢迎来到IBM可持续材料系列大型基础模型库。我们的模型涵盖多种分子表示形式与模态，包括SMILES字符串、SELFIES编码、3D原子坐标、3D密度网格、分子图结构等。这些模型旨在支持和推动材料科学与化学领域的研究进展。

GitHub仓库: 项目链接

研究论文: arXiv:2407.20267

基于SMILES的变压器编解码模型(SMI-TED)

本代码库提供与论文《化学语言基础模型的大型编解码器家族》相关的PyTorch实现源码。

论文预印本: 下载链接

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

• PyTorch格式(.pt): smi-ted-Light_40.pt
• 安全张量格式(.safetensors): model_weights.safetensors

技术咨询请联系: eduardo.soares@ibm.com 或 evital@br.ibm.com

模型概述

我们提出了基于SMILES的大型化学基础编解码模型SMI-TED，该模型在来自PubChem的9100万SMILES样本（相当于40亿分子标记）上进行预训练。SMI-TED支持包括量子属性预测在内的多种复杂任务，提供289M和8×289M两种参数量版本。在多个基准数据集上的实验表明，该模型在各类任务中均达到最先进性能。

目录结构

1. 快速开始
   1. 预训练模型与训练日志
   2. 复现Conda环境
2. 预训练流程
3. 微调方法
4. 特征提取
5. 文献引用

快速开始

本代码环境已在Nvidia V100和A100显卡完成验证

预训练模型与训练日志

我们提供了基于PubChem约9100万分子数据预训练的SMI-TED模型检查点。该预训练模型在MoleculeNet的分类与回归基准测试中展现出竞争优势。

根据需求将预训练权重文件.pt放入对应目录：

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

复现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

预训练流程

我们采用两种预训练策略：掩码语言模型训练编码器部分，以及编解码策略优化SMILES重构能力。SMI-TED在预处理时设置以下约束：

• 分子最大长度限制为202个标记
• 采用95/5/0的数据划分策略（5%数据用于解码器预训练）
• 同时使用100/0/0划分直接训练编解码器

预训练代码提供了小规模数据集的处理示例，需要8块A100显卡运行。

启动不同规格模型预训练：

bash training/run_model_light_training.sh

或

bash training/run_model_large_training.sh

专用训练脚本：

• train_model_D.py：仅训练解码器
• train_model_ED.py：联合训练编解码器

微调方法

微调数据集与环境配置详见finetune目录。环境就绪后，运行示例：

bash finetune/smi_ted_light/esol/run_finetune_esol.sh

训练过程检查点将保存在checkpoint_<指标名称>目录。

特征提取

示例笔记本smi_ted_encoder_decoder_example.ipynb演示了：

• 加载检查点文件
• 预训练模型在编解码任务的应用
• 分类与回归任务实例

模型加载方式：

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)

分子解码示例：

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

文献引用

@misc{soares2024largeencoderdecoderfamilyfoundation,
      title={化学语言基础模型的大型编解码器家族}, 
      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}, 
}