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