TxGravity-30B-A5B-Preview

Superseded. This is the July 2026 preview (single 2-epoch SFT run). The current release is TxGravity-30B-A5B, evaluated on full TDC test sets against a re-run TxGemma-27B baseline. Please use that model for new work.

Model Summary

TxGravity-30B-A5B-Preview is a therapeutics-focused language model fine-tuned from the Gravity-30B-A5B-base. It is trained to predict a broad range of therapeutic properties — small-molecule ADMET, toxicity, drug–target interaction, protein–protein and peptide–MHC interaction, and more — following the instruction format of the Therapeutic Data Commons (TDC) benchmark.

The model is supervised-fine-tuned on the TDC therapeutic instruction-tuning data (single-turn instruction → answer), covering 57 tasks. Answers are formatted as (A)/(B) for binary classification or a normalized 000–1000 bin for regression.

⚠️ This is a task-specialized property predictor, not a general instruction-tuned or safety-aligned assistant. Its outputs are intended for the TDC-style therapeutic prediction prompts it was trained on. Predictions may be inaccurate, biased, or incomplete and must be independently verified before any experimental, clinical, or decision-making use.

Training

• Base / init: Gravity 30B base checkpoint.
• Data: TDC therapeutic instruction-tuning set (single-turn), 57 tasks. Nine tasks without a paper-reported comparison / generation-only tasks were excluded from training.
• Objective: supervised fine-tuning (next-token, assistant span only), sequence packing at seq_len 8,192.
• Epochs: 2.
• Optimizer: Muon, lr 2e-5, cosine schedule (min-ratio 0.1), warmup 50 steps, grad clip 1.0, weight decay 0.
• Compute: 1× node of 8 GPUs, FSDP (dp_shard 8), full activation checkpointing, bf16.

Evaluation

Evaluated on the Therapeutic Data Commons (TDC) test sets with a corrected, generation-based pipeline: classification tasks are scored by generating an answer and reading the first-token log-probability of (A)/(B) (→ AUROC / AUPRC / Accuracy); regression tasks parse the predicted normalized bin (→ PCC / MAE / Spearman). Both columns are measured with this same pipeline.

Nine TDC tasks (disgenet, gdsc1, gdsc2, mirtarbase, phase1/2/3, uspto, uspto_yields) were excluded from training, so they are omitted here.

On the 57 evaluated TDC tasks, TxGravity-30B-A5B outperforms TxGemma-27B on 53/57 tasks.

Best value per row in bold (↑ = higher is better, ↓ = lower is better).

Quickstart

Installation

pip install "transformers>=4.45" accelerate safetensors

Using Transformers

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "trillionlabs/TxGravity-30B-A5B-Preview"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, torch_dtype=torch.bfloat16, device_map="auto"
)

# TxGravity expects TDC-style therapeutic prompts (Instruction / Context / Question).
prompt = (
    "Instructions: Answer the following question about drug properties.\n"
    "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 "
    "from reaching the brain.\n"
    "Question: Given a drug SMILES string, predict whether it\n"
    "(A) does not cross the BBB (B) crosses the BBB\n"
    "Drug SMILES: <smiles>CC(C)Cc1ccc(cc1)C(C)C(=O)O</smiles>\n"
    "Answer:"
)

messages = [{"role": "user", "content": prompt}]
inputs = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True, return_tensors="pt"
).to(model.device)

out = model.generate(inputs, max_new_tokens=8, do_sample=False)
print(tokenizer.decode(out[0][inputs.shape[-1]:], skip_special_tokens=True))

Limitations

• Task-specialized: it targets the TDC therapeutic prediction prompts it was trained on and is not a general conversational assistant.
• Predictions are model estimates and may be wrong; do not use for clinical or experimental decisions without independent validation.

Acknowledgements

This work was supported by the AI Specialized Foundation Model Project (인공지능 특화 파운데이션 모델 프로젝트), funded by the Ministry of Science and ICT (과학기술정보통신부, MSIT) and managed by the National IT Industry Promotion Agency (NIPA, 정보통신산업진흥원).

License

This model is released under the Apache License 2.0.

Citation

@misc{trillionlabs2026txgravity,
  title  = {TxGravity-30B-A5B-Preview},
  author = {Trillion Labs},
  year   = {2026},
  url    = {https://huggingface.co/trillionlabs/TxGravity-30B-A5B-Preview}
}

Contact

For questions, please contact Trillion Labs.