🚀 IBM材料基礎模型介紹
IBM推出了一系列用於可持續材料的大型基礎模型。這些模型涵蓋多種表示形式和模態,包括SMILES、SELFIES、3D原子位置、3D密度網格、分子圖等格式,旨在支持和推動材料科學與化學領域的研究。
GitHub: GitHub鏈接
論文: arXiv:2407.20267
✨ 主要特性
- 模型類型:基於SMILES的Transformer編解碼器(SMI - TED)
- 評估指標:準確率
- 支持任務:特徵提取、量子屬性預測等多種複雜任務
- 模型權重格式:PyTorch(
.pt
)和safetensors(.safetensors
)
屬性 |
詳情 |
模型類型 |
基於SMILES的Transformer編解碼器(SMI - TED) |
訓練數據 |
來自PubChem的9100萬個SMILES樣本,相當於40億個分子標記 |
📦 安裝指南
複製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。