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.103.10 - 3.11
PyTorch>= 2.02.1+ with CUDA 12.1
GPU VRAM4 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 機會

  1. 抗體藥物結合預測:MAMMAL 的 PPI 能力可延伸至 antibody-antigen binding,目前尚未有專用範例,但 prompt 架構天然支援
  2. 細胞藥物反應合成資料:結合 scRNA 模態 + DTI 能力,可為 cell line drug response 產生 synthetic data augmentation (合成資料增強)
  3. 多組學藥物標靶排名:整合蛋白質表現量 (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] → 交叉驗證預測結果與文獻數據
1011[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 推論整合,特別適合藥物開發管線中需要同時處理跨域生物資料的場景。