biomed-multi-alignment (MAMMAL) 完整教學
Repository: https://github.com/BiomedSciAI/biomed-multi-alignment Stars: 109 | Tags: IBM, multimodal, foundation License: Apache-2.0 | Language: Jupyter Notebook / Python arXiv: 2410.22367 Model Weights: ibm/biomed.omics.bl.sm.ma-ted-458m
2. 核心架構
2.1 整體架構圖
graph TB
subgraph Input["輸入模態 (Input Modalities)"]
P["Protein Sequence
蛋白質序列
AA tokenizer"]
S["Small Molecule
小分子 SMILES
SMILES tokenizer"]
G["Gene Expression
基因表現
Gene tokenizer + Scalars"]
end
subgraph MT["Modular Tokenizer
模組化分詞器"]
M1["AA Tokenizer"]
M2["SMILES Tokenizer"]
M3["Gene Tokenizer"]
M4["Unified ID Space
統一 ID 空間"]
end
subgraph Prompt["Task Prompt Syntax
任務提示語法"]
TP["<@TOKENIZER-TYPE=AA>
<TASK_TAG>
<MOLECULAR_ENTITY>...</MOLECULAR_ENTITY>
<EOS>"]
end
subgraph Model["MAMMAL Model (458M params)"]
ENC["T5 Encoder
編碼器"]
DEC["T5 Decoder
解碼器"]
EH["Encoder Head
ClassifierMLP"]
SH["Scalars Prediction Head
標量預測頭"]
end
subgraph Output["輸出模式"]
CLS["Classification
分類 (token generation)"]
REG["Regression
回歸 (scalar prediction)"]
EMB["Embeddings
嵌入向量 (vLLM pooling)"]
end
P --> M1
S --> M2
G --> M3
M1 --> M4
M2 --> M4
M3 --> M4
M4 --> TP
TP --> ENC
ENC -->|"Encoder-Decoder mode"| DEC --> CLS
ENC -->|"Encoder-only mode"| EH --> REG
ENC -->|"Encoder-only mode"| SH --> REG
ENC -->|"vLLM pooling"| EMB
2.2 Task Prompt 格式
MAMMAL 的核心設計哲學是用統一的 prompt 格式將不同模態與任務編碼為 token 序列。以下是各任務的 prompt 結構:
1PPI (蛋白質-蛋白質交互作用):
2<@TOKENIZER-TYPE=AA><BINDING_AFFINITY_CLASS><SENTINEL_ID_0>
3<MOLECULAR_ENTITY><MOLECULAR_ENTITY_GENERAL_PROTEIN>
4<SEQUENCE_NATURAL_START>{protein_1}<SEQUENCE_NATURAL_END>
5<MOLECULAR_ENTITY><MOLECULAR_ENTITY_GENERAL_PROTEIN>
6<SEQUENCE_NATURAL_START>{protein_2}<SEQUENCE_NATURAL_END><EOS>
7
8DTI (藥物-標靶交互作用):
9<@TOKENIZER-TYPE=AA><DRUG_TARGET_INTERACTION><SENTINEL_ID_0>
10<MOLECULAR_ENTITY><MOLECULAR_ENTITY_SMALL_MOLECULE>{smiles}</MOLECULAR_ENTITY>
11<MOLECULAR_ENTITY><MOLECULAR_ENTITY_GENERAL_PROTEIN>
12<SEQUENCE_NATURAL_START>{protein}<SEQUENCE_NATURAL_END><EOS>
2.3 專案目錄結構
1biomed-multi-alignment/
2├── mammal/ # 核心套件
3│ ├── model.py # MammalConfig + Mammal 模型類別
4│ ├── task.py # MammalTask 基礎類別
5│ ├── keys.py # 資料流 key 常數定義
6│ ├── lora.py # LoRA 微調整合 (peft)
7│ ├── losses.py # CE loss + Scalars loss
8│ ├── metrics.py # Accuracy / Regression metrics
9│ ├── lr_schedulers.py # 學習率排程
10│ ├── main_finetune.py # Hydra-based 微調入口
11│ └── examples/ # 下游任務範例
12│ ├── carcinogenicity/ # 藥物致癌性預測
13│ ├── cell_line_drug_response/ # 細胞株藥物反應 (IC50)
14│ ├── dti_bindingdb_kd/ # 藥物-標靶交互作用 (pKd)
15│ ├── protein_solubility/ # 蛋白質溶解度
16│ ├── scrna_cell_type/ # scRNA 細胞類型標註
17│ ├── tcr_epitope_binding/ # TCR-抗原表位結合
18│ ├── molnet/ # MolNet benchmark
19│ └── tests/ # 整合測試
20├── mammal_mcp/ # MCP Server (AI Agent 整合)
21│ ├── server.py # FastMCP server 實作
22│ ├── dependencies.py # 模型載入與資源管理
23│ └── util.py # 輸出處理工具
24├── mammal_vllm/ # vLLM 外掛 (高效推論)
25│ ├── vllm_mammal_plugin/ # 自動發現外掛
26│ └── examples/ # 離線 / 線上推論範例
27└── tutorials/ # Jupyter 教學
28 ├── begginer_inference.ipynb # Google Colab 入門
29 └── advanced_create_new_task.ipynb # 自定義任務
3. 安裝與設定
3.1 基本安裝 (conda + pip)
1# 1. 建立 conda 環境
2conda create -n mammal_env python=3.10 -y
3conda activate mammal_env
4
5# 2. 安裝 PyTorch (依 CUDA 版本調整)
6conda install pytorch pytorch-cuda=12.1 -c pytorch -c nvidia
7
8# 3a. 從 PyPI 安裝 (推薦一般使用者)
9pip install biomed-multi-alignment[examples]
10
11# 3b. 從原始碼安裝 (推薦開發者 / 需自定義任務)
12git clone https://github.com/BiomedSciAI/biomed-multi-alignment.git
13pip install -e ./biomed-multi-alignment[examples]
3.2 MCP Server 安裝 (AI Agent 整合)
1cd biomed-multi-alignment/mammal_mcp
2
3# 複製環境設定
4cp .env.example .env
5
6# 編輯 .env,啟用需要的任務模態
7# ENABLE_PPI=true
8# ENABLE_DTI=true
9# ENABLE_SOLUBILITY=true
10
11# 首次執行會下載模型 (~2GB)
12uv run python -m server
3.3 vLLM 外掛安裝 (高效 Embedding 服務)
1cd biomed-multi-alignment/mammal_vllm
2pip install -e .
3
4# 啟動 vLLM embedding server
5vllm serve ibm-research/biomed.omics.bl.sm.ma-ted-458m \
6 --runner pooling \
7 --trust-remote-code \
8 --skip_tokenizer_init \
9 --gpu_memory_utilization 0.4 \
10 --enforce_eager \
11 --no-enable-prefix-caching
3.4 系統需求
| 項目 | 最低需求 | 建議配置 |
|---|---|---|
| Python | >= 3.10 | 3.10 - 3.11 |
| PyTorch | >= 2.0 | 2.1+ with CUDA 12.1 |
| GPU VRAM | 4 GB (inference) | 16 GB (fine-tuning) |
| 磁碟空間 | ~2 GB (model weights) | ~10 GB (含 datasets) |
4. 使用方式與程式碼範例
範例 1:Protein-Protein Interaction (蛋白質-蛋白質交互作用預測)
此範例示範如何用預訓練模型直接推論兩個蛋白質是否會結合。適用於 target identification (標靶鑑定) 階段判斷蛋白質間的交互關係。
1import torch
2from fuse.data.tokenizers.modular_tokenizer.op import ModularTokenizerOp
3from mammal.model import Mammal
4from mammal.keys import *
5
6# === 1. 載入預訓練模型與分詞器 ===
7model = Mammal.from_pretrained("ibm/biomed.omics.bl.sm.ma-ted-458m")
8model.eval()
9tokenizer_op = ModularTokenizerOp.from_pretrained(
10 "ibm/biomed.omics.bl.sm.ma-ted-458m"
11)
12
13# === 2. 準備蛋白質序列 ===
14# Calmodulin (鈣調蛋白) — 廣泛的鈣信號傳導蛋白
15protein_calmodulin = (
16 "MADQLTEEQIAEFKEAFSLFDKDGDGTITTKELGTVMRSLGQNPTEAELQDMISELDQDG"
17 "FIDKEDLHDGDGKISFEEFLNLVNKEMTADVDGDGQVNYEEFVTMMTSK"
18)
19# Calcineurin (鈣調神經磷酸酶) — 已知與 calmodulin 結合
20protein_calcineurin = (
21 "MSSKLLLAGLDIERVLAEKNFYKEWDTWIIEAMNVGDEEVDRIKEFKEDEIFEEAKTLGTA"
22 "EMQEYKKQKLEEAIEGAFDIFDKDGNGYISAAELRHVMTNLGEKLTDEEVDEMIRQMWDQN"
23 "GDWDRIKELKFGEIKKLSAKDTRGTIFIKVFENLGTGVDSEYEDVSKYMLKHQ"
24)
25
26# === 3. 建構 task prompt ===
27sample_dict = dict()
28sample_dict[ENCODER_INPUTS_STR] = (
29 f"<@TOKENIZER-TYPE=AA><BINDING_AFFINITY_CLASS><SENTINEL_ID_0>"
30 f"<MOLECULAR_ENTITY><MOLECULAR_ENTITY_GENERAL_PROTEIN>"
31 f"<SEQUENCE_NATURAL_START>{protein_calmodulin}<SEQUENCE_NATURAL_END>"
32 f"<MOLECULAR_ENTITY><MOLECULAR_ENTITY_GENERAL_PROTEIN>"
33 f"<SEQUENCE_NATURAL_START>{protein_calcineurin}<SEQUENCE_NATURAL_END>"
34 f"<EOS>"
35)
36
37# === 4. 分詞 (Tokenization) ===
38tokenizer_op(
39 sample_dict=sample_dict,
40 key_in=ENCODER_INPUTS_STR,
41 key_out_tokens_ids=ENCODER_INPUTS_TOKENS,
42 key_out_attention_mask=ENCODER_INPUTS_ATTENTION_MASK,
43)
44sample_dict[ENCODER_INPUTS_TOKENS] = torch.tensor(
45 sample_dict[ENCODER_INPUTS_TOKENS]
46)
47sample_dict[ENCODER_INPUTS_ATTENTION_MASK] = torch.tensor(
48 sample_dict[ENCODER_INPUTS_ATTENTION_MASK]
49)
50
51# === 5. 生成預測 ===
52batch_dict = model.generate(
53 [sample_dict],
54 output_scores=True,
55 return_dict_in_generate=True,
56 max_new_tokens=5,
57)
58
59# === 6. 解碼結果 ===
60generated_output = tokenizer_op._tokenizer.decode(batch_dict[CLS_PRED][0])
61print(f"PPI 預測結果: {generated_output}")
62# 預期輸出: "<SENTINEL_ID_0><1>" 表示有交互作用
範例 2:Drug-Target Interaction 微調 + 推論 (藥物-標靶交互作用)
此範例示範完整的 fine-tuning 到 inference 流程。DTI 任務預測藥物與標靶蛋白的結合親和力 (binding affinity),以 pKd 值表示,是 lead optimization (先導化合物最佳化) 的關鍵步驟。
1# === Step 1: Fine-tuning (微調) ===
2# 使用 Hydra 組態系統,自動從 TDC BindingDB 下載資料
3python mammal/main_finetune.py \
4 --config-name config.yaml \
5 --config-path examples/dti_bindingdb_kd
6
7# === Step 2: Inference (推論) ===
8# 預測 Imatinib (伊馬替尼) 與 BCR-ABL 的結合親和力
9python mammal/examples/dti_bindingdb_kd/main_infer.py \
10 ./dti_finetune_output \
11 "MGCGCSSHPEDDWMENIDVCENCHYPIVPLDGKGTLLRNGSEVRDVRGSGIFGTVYKVAVK..." \
12 "CC1=C(C=C(C=C1)NC(=O)C2=CC=C(C=C2)CN3CCN(CC3)C)NC4=NC=CC(=N4)C5=CN=CC=C5" \
13 6.45 \
14 1.32
15
16# === Step 3: Evaluation (評估) ===
17python mammal/main_finetune.py \
18 --config-name config.yaml \
19 --config-path examples/dti_bindingdb_kd \
20 evaluate=True \
21 model.pretrained_kwargs.pretrained_model_name_or_path=./dti_finetune_output/best_epoch.ckpt
DTI 任務的 config.yaml 結構解析:
1# examples/dti_bindingdb_kd/config.yaml (關鍵欄位)
2model:
3 pretrained_kwargs:
4 pretrained_model_name_or_path: ibm/biomed.omics.bl.sm.ma-ted-458m
5 use_lora: true # 啟用 LoRA 降低 VRAM 需求
6 lora_kwargs:
7 r: 8 # LoRA rank
8 lora_alpha: 8 # LoRA alpha
9 target_modules: # 僅微調注意力層
10 - q
11 - v
12
13task:
14 name: dti_bindingdb_kd
15 norm_y_mean: 6.45 # pKd 正規化均值
16 norm_y_std: 1.32 # pKd 正規化標準差
範例 3:scRNA Cell Type Annotation (scRNA 細胞類型標註)
此範例展示 MAMMAL 處理單細胞基因表現資料的能力——對標靶鑑定 (target identification) 和生物標記發現 (biomarker discovery) 至關重要。模型將每個細胞的基因表現量排序編碼為 token 序列。
1# === Step 1: 資料準備 ===
2# 下載 Zheng68k PBMC 資料集 (需從 10x Genomics 網站取得)
3cd mammal/examples/scrna_cell_type/data
4python Zheng68k_to_anndata.py
5
6# 前處理:過濾 + 正規化 + log 轉換 + binning (離散化)
7python process_h5ad_data.py \
8 --input data.h5ad \
9 --output data_processed.h5ad \
10 --min_genes 200 \
11 --n_bins 10
12
13# === Step 2: Fine-tuning ===
14python mammal/main_finetune.py \
15 --config-name config.yaml \
16 --config-path examples/scrna_cell_type
17
18# === Step 3: Inference (對新細胞預測細胞類型) ===
19python mammal/examples/scrna_cell_type/scRNA_infer.py \
20 --model_path ./scrna_finetune_output/best_epoch.ckpt \
21 --h5ad_file new_cells.h5ad
scRNA 編碼策略 (Geneformer-inspired Ordered Gene Encoding):
1輸入: 基因表現矩陣 [cells x genes]
2 ↓ 篩選 (min 200 genes/cell)
3 ↓ 正規化 (total count → 1.0)
4 ↓ log(x+1) 轉換
5 ↓ Binning (10 bins, 0=最高表現)
6 ↓ 依表現量降序排列
7 ↓ 編碼為 token 序列:
8 <@TOKENIZER-TYPE=GENE><CELL_TYPE><SENTINEL_ID_0>
9 <GENE_SYMBOL_START>TP53<GENE_SYMBOL_END><BIN_3>
10 <GENE_SYMBOL_START>BRCA1<GENE_SYMBOL_END><BIN_5>
11 ...<EOS>
5. 在蛋白質設計/基因組模擬生態系中的定位
5.1 Domain 5 生態系定位
graph LR
subgraph ProteinDesign["A. Protein Design (蛋白質設計)"]
PLM["Protein LMs
ESM / ProstT5 / ProGen2"]
DIFF["Structure Diffusion
RFdiffusion / Chroma"]
AB["Antibody Design
DiffAb / AbLang"]
PRED["Structure Prediction
ColabFold"]
end
subgraph GenomicsSim["B. Genomics Simulation (基因組模擬)"]
SCFM["scRNA Foundation
scGPT / scvi-tools"]
SIM["scRNA Simulation
Splatter / scgen"]
GSEQ["Genome Sequence
wgsim / HyenaDNA"]
FEAT["Feature Extraction
PyFeat"]
end
subgraph MAMMAL_POS["MAMMAL (本專案)"]
MA["Multi-modal Bridge
蛋白質 + 小分子 + 單細胞
統一 task prompt"]
end
PLM -.->|"蛋白質序列嵌入"| MA
SCFM -.->|"基因表現嵌入"| MA
MA -->|"PPI / DTI / Cell Response"| PIPELINE["藥物開發管線"]
DIFF -.->|"設計候選蛋白"| MA
AB -.->|"抗體序列"| MA
style MA fill:#e1f5fe,stroke:#0288d1,stroke-width:3px
5.2 獨特定位
MAMMAL 在 Domain 5 中扮演 「多模態橋樑」(Multi-modal Bridge) 的角色:
| 維度 | MAMMAL 的定位 | 生態系互補 |
|---|---|---|
| 輸入模態 | 蛋白質 + 小分子 + 單細胞 (三模態) | ESM/scGPT 各僅處理單一模態 |
| 任務類型 | 分類 + 回歸 + 生成 (unified) | 多數工具僅支援單一任務類型 |
| 跨域推論 | DTI / Cell Line Response 天然跨域 | 需串接多個工具才能達成 |
| 規模 | 458M 參數 (輕量) | ESM-2 (650M-15B)、scGPT (30M-80M) |
5.3 Blue Ocean 機會
- 抗體藥物結合預測:MAMMAL 的 PPI 能力可延伸至 antibody-antigen binding,目前尚未有專用範例,但 prompt 架構天然支援
- 細胞藥物反應合成資料:結合 scRNA 模態 + DTI 能力,可為 cell line drug response 產生 synthetic data augmentation (合成資料增強)
- 多組學藥物標靶排名:整合蛋白質表現量 (proteomics) + 基因表現量 (transcriptomics) + 小分子親和力,建立 multi-evidence target ranking (多證據標靶排名)
6. 與其他工具的整合
6.1 Pre-IND 藥物開發管線整合 (WP1-WP7)
| 工作包 | MAMMAL 應用 | 整合工具 |
|---|---|---|
| WP1 標靶鑑定 | scRNA cell type annotation → 鑑定疾病相關細胞群 | scGPT (合成資料) + MAMMAL (標註) |
| WP2 標靶驗證 | PPI prediction → 驗證 target-pathway 關係 | STRING DB + MAMMAL (PPI scoring) |
| WP3 先導化合物發現 | DTI prediction → 虛擬篩選候選分子 | RDKit (分子庫) + MAMMAL (DTI) |
| WP4 先導化合物最佳化 | Drug carcinogenicity + solubility → 安全性/藥性篩選 | TDC benchmark + MAMMAL (ADMET) |
| WP5 候選藥物 | Cell line drug response → IC50 預測 | GDSC data + MAMMAL (response) |
| WP7 生物製劑 | PPI + antibody sequence → 抗體工程 | DiffAb (設計) + MAMMAL (評估) |
6.2 MCP Agent 整合
MAMMAL 提供 FastMCP server,可直接被 Claude Desktop / MCPHost / 任何 MCP 相容的 AI Agent 呼叫:
1{
2 "mcpServers": {
3 "mammal": {
4 "command": "uv",
5 "args": [
6 "--directory", "/path/to/mammal_mcp",
7 "run", "server.py"
8 ]
9 }
10 }
11}
MCP 提供的工具 (Tools):
gene_name_to_amino_acid_sequence— 查詢基因名對應的蛋白質序列 (via UniProt)protein_protein_interaction_prediction— PPI 預測drug_target_interaction_prediction— DTI 結合親和力預測protein_solubility_prediction— 蛋白質溶解度預測tcr_epitope_binding_prediction— TCR 抗原表位結合預測
6.3 vLLM Embedding 服務
用於大規模蛋白質/分子嵌入 (embedding) 計算,可與向量資料庫整合做 similarity search (相似性搜尋):
1from vllm import LLM
2from vllm_mammal_plugin.tokenization import tokenize_mammal
3
4# 初始化
5model = LLM(
6 "ibm-research/biomed.omics.bl.sm.ma-ted-458m",
7 runner="pooling",
8 skip_tokenizer_init=True,
9)
10
11# 取得蛋白質嵌入向量
12prompt = "<@TOKENIZER-TYPE=AA><SEQUENCE_NATURAL_START>MADQLT...<SEQUENCE_NATURAL_END><EOS>"
13tokens = tokenize_mammal(prompt)
14embeddings = model.embed([tokens])
15# → L2-normalized dense vector,可直接用於 cosine similarity
6.4 AIKT 管線整合建議
1[paper-search] → 找到目標蛋白 / 藥物相關文獻
2 ↓
3[tu-plan-generator] → ChEMBL/DrugBank 查詢候選分子
4 ↓
5[MAMMAL DTI] → 預測候選分子與標靶的 pKd
6 ↓
7[MAMMAL Cell Response] → 預測 IC50 on disease cell lines
8 ↓
9[paper-qa-lite] → 交叉驗證預測結果與文獻數據
10 ↓
11[quarkdown] → 彙整為研究報告 HTML
7. 優缺點分析
優點
| 面向 | 評價 |
|---|---|
| 多模態統一 | 單一模型處理蛋白質 + 小分子 + 基因表現,大幅降低工具串接成本 |
| Task prompt 可擴充 | 新增任務僅需定義 prompt 格式 + Task 子類別,不需改動模型架構 |
| LoRA 微調友善 | 內建 peft 整合,8GB VRAM 即可微調特定任務 |
| MCP + vLLM 生態 | 原生支援 AI Agent 整合 (MCP) 與高效推論 (vLLM),部署成熟度高 |
| Apache-2.0 授權 | 商用友善,無限制 |
| IBM 長期維護 | 持續更新(最近仍在新增 MCP / vLLM / cell line drug response 功能) |
| PyPI 可安裝 | pip install biomed-multi-alignment 一行搞定 |
缺點
| 面向 | 評價 |
|---|---|
| 模型規模偏小 | 458M 參數,遠小於 ESM-2 (15B) / scGPT 系列,單模態深度可能不足 |
| 預訓練資料不透明 | 論文未完整公開 20 億樣本的組成細節與品質控制流程 |
| scRNA 支援初階 | 細胞類型標註僅一個範例 (Zheng68k),缺乏 batch correction / trajectory 等進階功能 |
| 無 3D 結構 | 完全基於序列,不處理蛋白質 3D 結構資訊,與 RFdiffusion / ColabFold 互補而非替代 |
| 回歸任務精度 | DTI / IC50 的 benchmark 表現為 “practical demonstrations”,未宣稱 SOTA |
| 文件品質不一 | scRNA 範例有完整 README,但其他範例依賴 README.md 中簡短說明 |
| GPU 依賴 | 推論可用 CPU 但速度極慢,實務上需 GPU |
適用場景與限制摘要
1最適合:
2 - 需要同時處理蛋白質 + 小分子 + 基因表現的跨域任務
3 - 快速建立 DTI / PPI / 藥物毒性的 baseline 預測
4 - 透過 MCP 讓 AI Agent 直接呼叫生醫預測工具
5 - LoRA 微調至特定治療領域的 custom dataset
6
7不適合:
8 - 需要蛋白質 3D 結構生成/預測 → 用 RFdiffusion / ColabFold
9 - 需要高精度單模態蛋白質嵌入 → 用 ESM-2
10 - 需要完整 scRNA 分析流程 (QC / batch correction / trajectory) → 用 scvi-tools
11 - 需要從頭設計全新蛋白質序列 → 用 EvoDiff / ProGen2
一行摘要:MAMMAL 是 IBM 開發的 458M 參數多模態生醫基礎模型,以統一 task prompt 語法橫跨蛋白質/小分子/單細胞三大模態,內建 LoRA 微調 + MCP Agent + vLLM 推論整合,特別適合藥物開發管線中需要同時處理跨域生物資料的場景。
Comments