Txgemma 9b Predict

🚀 TxGemma model card

TxGemma is a collection of lightweight, state-of-the-art, open language models fine - tuned for therapeutic development. It can process and understand therapeutic information, and is useful for drug discovery and related research.

🚀 Quick Start

Below are some example code snippets to help you quickly get started running the model locally on GPU. If you want to use the model to run inference on a large number of inputs, we recommend that you create a production version using Model Garden.

Basic Usage

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-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)

Advanced Usage

# pip install accelerate transformers
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load model directly from Hugging Face Hub
tokenizer = AutoTokenizer.from_pretrained("google/txgemma-9b-predict")
model = AutoModelForCausalLM.from_pretrained(
    "google/txgemma-9b-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))

Alternatively, you can use the 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-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)

✨ Features

• Versatility: Exhibits strong performance across a wide range of therapeutic tasks, outperforming or matching best - in - class performance on a significant number of benchmarks.
• Data Efficiency: Shows competitive performance even with limited data compared to larger models, offering improvements over its predecessors.
• Conversational Capability (TxGemma - Chat): Includes conversational variants that can engage in natural language dialogue and explain the reasoning behind their predictions.
• Foundation for Fine - tuning: Can be used as a pre - trained foundation for specialized use cases.

📚 Documentation

Model information

Description

TxGemma is a collection of lightweight, state - of - the - art, open language models built upon Gemma 2, fine - tuned for therapeutic development. It comes in 3 sizes, 2B, 9B, and 27B.

TxGemma models are designed to process and understand information related to various therapeutic modalities and targets, including small molecules, proteins, nucleic acids, diseases, and cell lines. TxGemma excels at tasks such as property prediction, and can serve as a foundation for further fine - tuning or as an interactive, conversational agent for drug discovery. The model is fine - tuned from Gemma 2 using a diverse set of instruction - tuning datasets, curated from the Therapeutics Data Commons (TDC).

TxGemma is offered as both a prediction model that expects a narrow form of prompting and for the 9B and 27B version, conversational models that are more flexible and can be used in multi - turn interactions, including to explain its rationale behind a prediction. This conversational model comes at the expense of some raw prediction performance. See our manuscript for more information.

Potential Applications

TxGemma can be a valuable tool for researchers in the following areas:

• Accelerated Drug Discovery: Streamline the therapeutic development process by predicting properties of therapeutics and targets for a wide variety of tasks including target identification, drug - target interaction prediction, and clinical trial approval prediction.

Examples

See the following Colab notebooks for examples of how to use TxGemma:

• To give the model a quick try, running it locally with weights from Hugging Face, see [Quick start notebook in Colab](https://colab.research.google.com/github/google - gemini/gemma - cookbook/blob/main/TxGemma/%5BTxGemma%5DQuickstart_with_Hugging_Face.ipynb), which includes some example eval tasks from TDC.
• For a demo of how to fine - tune TxGemma in Hugging Face, see our [Fine - tuning notebook in Colab](https://colab.research.google.com/github/google - gemini/gemma - cookbook/blob/main/TxGemma/%5BTxGemma%5DFinetune_with_Hugging_Face.ipynb).
• For a demo of how TxGemma can be used as a tool as part of a larger agentic workflow powered by Gemini 2 see the [Agentic workflow notebook in Colab](https://colab.research.google.com/github/google - gemini/gemma - cookbook/blob/main/TxGemma/%5BTxGemma%5DAgentic_Demo_with_Hugging_Face.ipynb).

Model architecture overview

• TxGemma is based on the Gemma 2 family of lightweight, state - of - the - art open LLMs. It utilizes a decoder - only transformer architecture.
• Base Model: Gemma 2 (2B, 9B, and 27B parameter versions).
• Fine - tuning Data: Therapeutics Data Commons, a collection of instruction - tuning datasets covering diverse therapeutic modalities and targets.
• Training Approach: Instruction fine - tuning using a mixture of therapeutic data (TxT) and, for conversational variants, general instruction - tuning data.
• Conversational Variants: TxGemma - Chat models (9B and 27B) are trained with a mixture of therapeutic and general instruction - tuning data to maintain conversational abilities.

Technical Specifications

Performance & Validation

TxGemma's performance has been validated on a comprehensive benchmark of 66 therapeutic tasks derived from TDC.

Key performance metrics

• Aggregated Improvement: Improves over the original Tx - LLM paper on 45 out of 66 therapeutic tasks.
• Best - in - Class Performance: Surpasses or matches best - in - class performance on 50 out of 66 tasks, exceeding specialist models on 26 tasks. See Table A.11 of the TxGemma paper for the full breakdown.

Inputs and outputs

• Input: Text. For best performance, text prompts should be formatted according to the TDC structure, including instructions, context, question, and, optionally, few - shot examples. Inputs can include SMILES strings, amino acid sequences, nucleotide sequences, and natural language text.
• Output: Text.

Dataset details

Training dataset

• Therapeutics Data Commons: A curated collection of instruction - tuning datasets covering 66 tasks spanning the discovery and development of safe and effective medicine. This includes over 15 million data points across different biomedical entities. Released TxGemma models are only trained on datasets with commercial licenses, whereas models in our publication are also trained on datasets with non - commercial licenses.
• General Instruction - Tuning Data: Used for TxGemma - Chat in combination with TDC.

Evaluation dataset

Therapeutics Data Commons: The same 66 tasks used for training are used for evaluation, following TDC's recommended methodologies for data splits (random, scaffold, cold - start, combination, and temporal).

Implementation information

Software

Training was done using [JAX](https://github.com/jax - ml/jax).

JAX allows researchers to take advantage of the latest generation of hardware, including TPUs, for faster and more efficient training of large models.

Use and limitations

Intended use

• Research and development of therapeutics.

Benefits

TxGemma provides a versatile and powerful tool for accelerating therapeutic development. It offers:

• Strong performance across a wide range of tasks.
• Data efficiency compared to larger models.
• A foundation for further fine - tuning from private data.
• Integration into agentic workflows.

Limitations

• Trained on public data from TDC.
• Task - specific validation remains an important aspect of downstream model development by the end user.
• As with any research, developers should ensure that any downstream application is validated to understand performance using data that is appropriately representative of the intended use setting for the specific application (e.g., age, sex, gender, condition, scanner, etc.).

🔧 Technical Details

TxGemma is fine - tuned from Gemma 2 using a diverse set of instruction - tuning datasets from the Therapeutics Data Commons. It uses a decoder - only transformer architecture, and for conversational variants, general instruction - tuning data is also incorporated. The model has been validated on 66 therapeutic tasks from TDC, showing significant performance improvements over its predecessors.

📄 License

The use of TxGemma is governed by the [Health AI Developer Foundations terms of use](https://developers.google.com/health - ai - developer - foundations/terms).

📚 Citation

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

Find the paper here.