Txgemma 9b Chat

🚀 TxGemma模型卡片

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

🚀 快速開始

若要在Hugging Face上訪問TxGemma，你需要查看並同意Health AI Developer Foundation的使用條款。請確保你已登錄Hugging Face，然後點擊下方按鈕。請求將立即處理。

訪問按鈕

Acknowledge license

資源鏈接

• 模型文檔：TxGemma
• Google Cloud Model Garden上的模型：TxGemma
• Hugging Face上的模型：TxGemma
• GitHub倉庫：TxGemma
• 快速開始筆記本：notebooks/quick_start
• 支持信息：請查看Contact

使用條款

Health AI Developer Foundations使用條款

作者

Google

✨ 主要特性

模型描述

TxGemma是基於Gemma 2構建的輕量級、最先進的開放語言模型集合，針對治療開發進行了微調。它有2B、9B和27B三種規模，旨在處理和理解與各種治療方式和靶點相關的信息，包括小分子、蛋白質、核酸、疾病和細胞系等。TxGemma在屬性預測等任務中表現出色，可作為進一步微調的基礎，也可作為藥物發現的交互式對話代理。該模型使用從Therapeutics Data Commons (TDC)精選的多樣化指令調整數據集，從Gemma 2進行微調。

關鍵特性

• 多功能性：在廣泛的治療任務中表現出色，在大量基準測試中超越或達到同類最佳性能。
• 數據效率：與大型模型相比，即使在數據有限的情況下也能展現出有競爭力的性能，相較於其前身有所改進。
• 對話能力（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-9b-chat",
    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-9b-chat")
model = AutoModelForCausalLM.from_pretrained(
    "google/txgemma-9b-chat",
    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-9b-chat",
    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)

高級用法

應用聊天模板進行對話使用

# pip install accelerate transformers
from transformers import AutoTokenizer, AutoModelForCausalLM

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

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

# Format prompt in the conversational format
messages = [
    { "role": "user", "content": prompt}
]

# Apply the tokenizer's built-in chat template
chat_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

生成響應

inputs = tokenizer.encode(chat_prompt, add_special_tokens=False, return_tensors="pt")
outputs = model.generate(input_ids=inputs.to("cuda"), max_new_tokens=8)
response = tokenizer.decode(outputs[0, len(inputs[0]):], skip_special_tokens=True)
print(response)

多輪交互

messages.extend([
    { "role": "assistant", "content": response },
    { "role": "user", "content": "Explain your reasoning based on the molecule structure." },
])

chat_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer.encode(chat_prompt, add_special_tokens=False, return_tensors="pt")
outputs = model.generate(input_ids=inputs.to("cuda"), max_new_tokens=512)
response = tokenizer.decode(outputs[0, len(inputs[0]):], skip_special_tokens=True)
print(response)

使用pipeline API進行多輪交互

# pip install transformers
from transformers import pipeline

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

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

# Format prompt in the conversational format for initial turn
messages = [
    { "role": "user", "content": prompt}
]

# Generate response for initial turn
outputs = pipe(messages, max_new_tokens=8)
print(outputs[0]["generated_text"][-1]["content"].strip())

# Append user prompt for an additional turn
messages = outputs[0]["generated_text"]
messages.append(
    { "role": "user", "content": "Explain your reasoning based on the molecule structure." }
)

# Generate response for additional turn
outputs = pipe(messages, max_new_tokens=512)
print(outputs[0]["generated_text"][-1]["content"].strip())

示例筆記本

• 快速嘗試模型：Colab中的快速開始筆記本，包含一些來自TDC的示例評估任務。
• 在Hugging Face中微調TxGemma的演示：Colab中的微調筆記本。
• 將TxGemma作為由Gemini 2驅動的更大代理工作流的一部分的演示：Colab中的代理工作流筆記本。

📚 詳細文檔

模型架構概述

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

技術規格

性能與驗證

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

關鍵性能指標

• 綜合改進：在66個治療任務中的45個任務上，相較於原始的Tx - LLM論文有所改進。
• 同類最佳性能：在66個任務中的50個任務上超越或達到同類最佳性能，在26個任務上超過專業模型。詳情見TxGemma論文的表A.11。

輸入和輸出

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

數據集詳情

訓練數據集

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

評估數據集

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

實現信息

軟件

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

使用和限制

預期用途

• 治療方法的研究和開發。

優點

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

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

限制

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

🔧 技術細節

文檔未提供更詳細的技術實現細節相關內容。

📄 許可證

TxGemma的使用受Health AI Developer Foundations使用條款的約束。

📖 引用

@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},
}

論文鏈接：TxGemma: Efficient and Agentic LLMs for Therapeutics