This repository is publicly accessible, but you have to accept the conditions to access its files and content.
To access TxGemma on Hugging Face, you're required to review and agree to Health AI Developer Foundation's terms of use. To do this, please ensure you're logged in to Hugging Face and click below. Requests are processed immediately.
Log in or Sign Up to review the conditions and access this model content.
Model documentation: TxGemma
Resources:
Terms of use: Health AI Developer Foundations terms of use
Author: Google
This section describes the TxGemma model and how to use it.
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.
Key Features
Potential Applications
TxGemma can be a valuable tool for researchers in the following areas:
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.
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)
The resulting prompt is in the format expected by the model:
Instructions: Answer the following question about drug properties.
Context: As a membrane separating circulating blood and brain extracellular fluid, the blood-brain barrier (BBB) is the protection layer that blocks most foreign drugs. Thus the ability of a drug to penetrate the barrier to deliver to the site of action forms a crucial challenge in development of drugs for central nervous system.
Question: Given a drug SMILES string, predict whether it
(A) does not cross the BBB (B) crosses the BBB
Drug SMILES: CN1C(=O)CN=C(C2=CCCCC2)c2cc(Cl)ccc21
Answer:
# 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))
Alternatively, you can use the pipeline API, which provides a simple way to
run inference while abstracting away complex details of loading and using the
model and tokenizer:
# 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)
TxGemma-Chat models use a chat template that must be adhered to for conversational use. The easiest way to apply it is using the tokenizer`s built-in chat template, as shown in the following snippet.
Let's load the model and apply the chat template to a conversation. In this example, we'll start with a single user interaction:
# 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)
At this point, the prompt contains the following text:
<bos><start_of_turn>user
Instructions: Answer the following question about drug properties.
Context: As a membrane separating circulating blood and brain extracellular fluid, the blood-brain barrier (BBB) is the protection layer that blocks most foreign drugs. Thus the ability of a drug to penetrate the barrier to deliver to the site of action forms a crucial challenge in development of drugs for central nervous system.
Question: Given a drug SMILES string, predict whether it
(A) does not cross the BBB (B) crosses the BBB
Drug SMILES: CN1C(=O)CN=C(C2=CCCCC2)c2cc(Cl)ccc21
Answer:<end_of_turn>
<start_of_turn>model
As you can see, each turn is preceded by a <start_of_turn> delimiter and then
the role of the entity (either user, for content supplied by the user, or
model, for LLM responses). Turns finish with the <end_of_turn> token.
You can follow this format to build the prompt manually if you need to do it without the tokenizer's chat template.
After the prompt is ready, generation can be performed as follows:
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)
For multiple interactions, append the model response and user prompt for an additional turn to the chat message history and perform generation in the same way:
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)
Alternatively, you can use the pipeline API, which abstracts away details such
as applying the chat template using the tokenizer:
# 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())
See the following Colab notebooks for examples of how to use TxGemma:
TxGemma's performance has been validated on a comprehensive benchmark of 66 therapeutic tasks derived from TDC.
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.
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.
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.
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).
The use of TxGemma is governed by the Health AI Developer Foundations terms of use.
This section contains details about the model internals.
Training was done using JAX.
JAX allows researchers to take advantage of the latest generation of hardware, including TPUs, for faster and more efficient training of large models.
TxGemma provides a versatile and powerful tool for accelerating therapeutic development. It offers:
@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.