Txgemma 2b Predict

🚀 TxGemma模型卡片

TxGemma是基於Gemma 2構建的輕量級、最先進的開放語言模型集合，針對治療開發進行了微調。它有2B、9B和27B三種規模，能夠處理和理解與各種治療方式和靶點相關的信息，可用於藥物發現等多個領域。

🚀 快速開始

本地運行模型

以下是一些示例代碼片段，可幫助你在GPU上快速本地運行模型。如果你想使用該模型對大量輸入進行推理，建議你使用Model Garden創建生產版本。

治療任務的提示格式化

import json
from huggingface_hub import hf_hub_download

# Load prompt template for tasks from TDC
tdc_prompts_filepath = hf_hub_download(
    repo_id="google/txgemma-2b-predict",
    filename="tdc_prompts.json",
)
with open(tdc_prompts_filepath, "r") as f:
    tdc_prompts_json = json.load(f)

# Set example TDC task and input
task_name = "BBB_Martins"
input_type = "{Drug SMILES}"
drug_smiles = "CN1C(=O)CN=C(C2=CCCCC2)c2cc(Cl)ccc21"

# Construct prompt using template and input drug SMILES string
TDC_PROMPT = tdc_prompts_json[task_name].replace(input_type, drug_smiles)
print(TDC_PROMPT)

在預測任務上運行模型

# pip install accelerate transformers
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load model directly from Hugging Face Hub
tokenizer = AutoTokenizer.from_pretrained("google/txgemma-2b-predict")
model = AutoModelForCausalLM.from_pretrained(
    "google/txgemma-2b-predict",
    device_map="auto",
)

# Formatted TDC prompt (see "Formatting prompts for therapeutic tasks" section above)
prompt = TDC_PROMPT

# Prepare tokenized inputs
input_ids = tokenizer(prompt, return_tensors="pt").to("cuda")

# Generate response
outputs = model.generate(**input_ids, max_new_tokens=8)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

或者，你可以使用pipeline API，它提供了一種簡單的方法來運行推理，同時抽象掉了加載和使用模型及分詞器的複雜細節：

# pip install transformers
from transformers import pipeline

# Instantiate a text generation pipeline using the model
pipe = pipeline(
    "text-generation",
    model="google/txgemma-2b-predict",
    device="cuda",
)

# Formatted TDC prompt (see "Formatting prompts for therapeutic tasks" section above)
prompt = TDC_PROMPT

# Generate response
outputs = pipe(prompt, max_new_tokens=8)
response = outputs[0]["generated_text"]
print(response)

示例

有關如何使用TxGemma的示例，請參閱以下Colab筆記本：

• 若要快速嘗試該模型，使用Hugging Face的權重在本地運行，請參閱Colab中的快速入門筆記本，其中包含一些來自TDC的示例評估任務。
• 有關如何在Hugging Face中微調TxGemma的演示，請參閱我們的Colab中的微調筆記本。
• 有關如何將TxGemma用作由Gemini 2驅動的更大代理工作流的一部分的工具的演示，請參閱Colab中的代理工作流筆記本。

✨ 主要特性

模型描述

TxGemma是基於Gemma 2構建的輕量級、最先進的開放語言模型集合，針對治療開發進行了微調。它有3種規模，分別為2B、9B和27B。

TxGemma模型旨在處理和理解與各種治療方式和靶點相關的信息，包括小分子、蛋白質、核酸、疾病和細胞系。TxGemma在屬性預測等任務中表現出色，可作為進一步微調的基礎，或作為藥物發現的交互式對話代理。該模型使用從治療學數據公共庫（TDC）精心策劃的多樣化指令調優數據集，從Gemma 2進行微調。

TxGemma既提供期望特定形式提示的預測模型，對於9B和27B版本，還提供更靈活的對話模型，可用於多輪交互，包括解釋其預測背後的原理。這種對話模型在一定程度上犧牲了原始預測性能。更多信息請參閱我們的論文。

關鍵特性

• 多功能性：在廣泛的治療任務中表現出色，在大量基準測試中超越或達到同類最佳性能。
• 數據效率：與較大模型相比，即使在數據有限的情況下也能顯示出有競爭力的性能，相較於其前身有所改進。
• 對話能力（TxGemma - Chat）：包括能夠進行自然語言對話並解釋其預測推理的對話變體。
• 微調基礎：可作為預訓練基礎用於特定用例。

潛在應用

TxGemma可以成為以下領域研究人員的寶貴工具：

• 加速藥物發現：通過預測治療方法和靶點的屬性，簡化治療開發過程，適用於各種任務，包括靶點識別、藥物 - 靶點相互作用預測和臨床試驗批准預測。

📦 安裝指南

文檔未提及具體安裝步驟，故跳過此章節。

💻 使用示例

基礎用法

# 治療任務的提示格式化
import json
from huggingface_hub import hf_hub_download

# Load prompt template for tasks from TDC
tdc_prompts_filepath = hf_hub_download(
    repo_id="google/txgemma-2b-predict",
    filename="tdc_prompts.json",
)
with open(tdc_prompts_filepath, "r") as f:
    tdc_prompts_json = json.load(f)

# Set example TDC task and input
task_name = "BBB_Martins"
input_type = "{Drug SMILES}"
drug_smiles = "CN1C(=O)CN=C(C2=CCCCC2)c2cc(Cl)ccc21"

# Construct prompt using template and input drug SMILES string
TDC_PROMPT = tdc_prompts_json[task_name].replace(input_type, drug_smiles)
print(TDC_PROMPT)

高級用法

# 在預測任務上運行模型
# pip install accelerate transformers
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load model directly from Hugging Face Hub
tokenizer = AutoTokenizer.from_pretrained("google/txgemma-2b-predict")
model = AutoModelForCausalLM.from_pretrained(
    "google/txgemma-2b-predict",
    device_map="auto",
)

# Formatted TDC prompt (see "Formatting prompts for therapeutic tasks" section above)
prompt = TDC_PROMPT

# Prepare tokenized inputs
input_ids = tokenizer(prompt, return_tensors="pt").to("cuda")

# Generate response
outputs = model.generate(**input_ids, max_new_tokens=8)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

📚 詳細文檔

模型架構概述

• 基礎模型：TxGemma基於Gemma 2系列輕量級、最先進的開放大語言模型。它採用僅解碼器的Transformer架構。
• 微調數據：治療學數據公共庫，這是一組涵蓋各種治療方式和靶點的指令調優數據集。
• 訓練方法：使用治療數據（TxT）的混合進行指令微調，對於對話變體，還使用通用指令調優數據。
• 對話變體：TxGemma - Chat模型（9B和27B）使用治療和通用指令調優數據的混合進行訓練，以保持對話能力。

技術規格

性能與驗證

TxGemma的性能已在從TDC派生的66個治療任務的綜合基準測試中得到驗證。

關鍵性能指標

• 綜合改進：在66個治療任務中的45個任務上比原始Tx - LLM論文有所改進。
• 同類最佳性能：在66個任務中的50個任務上超越或達到同類最佳性能，在26個任務上超過專業模型。完整明細請參閱TxGemma論文的表A.11。

輸入和輸出

• 輸入：文本。為獲得最佳性能，文本提示應根據TDC結構進行格式化，包括指令、上下文、問題以及可選的少樣本示例。輸入可以包括SMILES字符串、氨基酸序列、核苷酸序列和自然語言文本。
• 輸出：文本。

數據集詳情

訓練數據集

• 治療學數據公共庫：一個精心策劃的指令調優數據集集合，涵蓋66個任務，涉及安全有效藥物的發現和開發。這包括跨不同生物醫學實體的超過1500萬個數據點。已發佈的TxGemma模型僅在具有商業許可證的數據集上進行訓練，而我們論文中的模型還在具有非商業許可證的數據集上進行訓練。
• 通用指令調優數據：與TDC結合用於TxGemma - Chat。

評估數據集

治療學數據公共庫：使用與訓練相同的66個任務進行評估，遵循TDC推薦的數據分割方法（隨機、支架、冷啟動、組合和時間）。

使用和限制

預期用途

• 治療學的研究和開發。

優點

TxGemma為加速治療開發提供了一個多功能且強大的工具。它具有以下優點：

• 在廣泛的任務中表現出色。
• 與較大模型相比，數據效率更高。
• 可作為從私有數據進行進一步微調的基礎。
• 可集成到代理工作流中。

限制

• 在來自TDC的公共數據上進行訓練。
• 特定任務的驗證仍然是最終用戶進行下游模型開發的重要方面。
• 與任何研究一樣，開發人員應確保任何下游應用都經過驗證，以瞭解使用適合特定應用預期使用場景的數據（例如，年齡、性別、狀況、掃描儀等）時的性能。

🔧 技術細節

軟件

訓練使用JAX進行。

JAX允許研究人員利用最新一代的硬件，包括TPU，以更快、更高效地訓練大型模型。

📄 許可證

TxGemma的使用受健康人工智能開發者基礎使用條款的約束。

📖 引用

@article{wang2025txgemma,
    title={TxGemma: Efficient and Agentic LLMs for Therapeutics},
    author={Wang, Eric and Schmidgall, Samuel and Jaeger, Paul F. and Zhang, Fan and Pilgrim, Rory and Matias, Yossi and Barral, Joelle and Fleet, David and Azizi, Shekoofeh},
    year={2025},
}

論文鏈接：here