Materials.smi Ted

🚀 IBMの材料用基礎モデルの紹介

IBMの持続可能な材料用の大規模基礎モデルシリーズへようこそ。当社のモデルは、SMILES、SELFIES、3D原子位置、3D密度グリッド、分子グラフなど、さまざまな表現形式とモダリティをカバーしています。これらのモデルは、材料科学と化学の研究を支援し、推進するために設計されています。

GitHub: GitHubリンク

論文: arXiv:2407.20267

✨ 主な機能

• 材料科学と化学の研究を支援するための多様な表現形式とモダリティを持つ大規模基礎モデル。
• SMILESベースのTransformerエンコーダー - デコーダー（SMI - TED）モデルを提供し、量子特性予測などの複雑なタスクをサポート。
• 分類と回帰のベンチマークで競争力のある性能を発揮する事前学習済みモデル。

📦 インストール

事前学習済みモデルと学習ログ

PubChemから選りすぐった約9100万個の分子のデータセットで事前学習されたSMI - TEDモデルのチェックポイントを提供しています。事前学習済みモデルは、MoleculeNetの分類と回帰のベンチマークで競争力のある性能を示します。

必要に応じて、SMI - TEDのpre - trained weights.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

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

🔧 事前学習

事前学習には、エンコーダー部分を学習するためのマスク言語モデル手法と、SMILES再構築を改善し、生成された潜在空間を改良するためのエンコーダー - デコーダー戦略の2つの戦略を使用しています。

SMI - TEDは、PubChemから選りすぐった9100万個のSMILESサンプルのデータセットで事前学習されており、以下の制約が適用されています。

• 前処理中に化合物は最大202トークンにフィルタリングされます。
• エンコーダー学習には95/5/0の分割を使用し、デコーダーの事前学習には5％のデータを使用します。
• エンコーダーとデコーダーを直接学習するために100/0/0の分割も使用し、モデルの性能を向上させます。

事前学習コードは、小さなデータセットでのデータ処理とモデル学習の例を提供しており、8台のA100 GPUが必要です。

SMI - TEDモデルの2つのバリアントを事前学習するには、次のコマンドを実行します。

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_<measure_name>という名前のディレクトリに保存されます。

💻 使用例

基本的な使用法

このノートブック smi_ted_encoder_decoder_example.ipynb には、チェックポイントファイルを読み込み、事前学習済みモデルをエンコーダーとデコーダーのタスクに使用するコードが含まれています。また、分類と回帰のタスクの例も含まれています。

smi - tedを読み込むには、次のようにします。

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)

📚 ドキュメント

SMILESベースのTransformerエンコーダー - デコーダー（SMI - TED）

このリポジトリは、論文 "A Large Encoder - Decoder Family of Foundation Models for Chemical Language" に関連するPyTorchのソースコードを提供しています。

論文: Arxivリンク

モデルの重みを2つの形式で提供しています。

• PyTorch (.pt): smi-ted-Light_40.pt
• safetensors (.safetensors): model_weights.safetensors

詳細については、eduardo.soares@ibm.com または evital@br.ibm.com までお問い合わせください。

概要

大規模なエンコーダー - デコーダー化学基礎モデルであるSMILESベースのTransformerエンコーダー - デコーダー（SMI - TED）を紹介します。これは、PubChemから取得した9100万個のSMILESサンプル（40億個の分子トークンに相当）の精選データセットで事前学習されています。SMI - TEDは、量子特性予測を含むさまざまな複雑なタスクをサポートし、2つの主要なバリアント（289Mと8X289M）があります。複数のベンチマークデータセットを用いた実験で、さまざまなタスクに対する最先端の性能を実証しています。詳細については、eduardo.soares@ibm.com または evital@br.ibm.com までお問い合わせください。

📄 ライセンス

このプロジェクトは、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}, 
}