Gemma4EC-2B-Predict

LoRA adapter for google/txgemma-2b-predict that predicts the Enzyme Commission (EC) number of a biochemical reaction from its SMILES, down to the sub-subclass level (class.subclass.sub-subclass).

• Input: a reaction SMILES, e.g. CCO.O=O>>CC=O.OO
• Output: an EC number, e.g. 1.01.03
• Code and benchmark: github.com/PlanesLab/Gemma4EC

Performance

Benchmarked on 10 deduplicated, EC-stratified re-splits of ECReact (57,464 reaction–EC pairs, 80/10/10 per seed). Support-weighted MCC at the EC sub-subclass level:

	MCC
Gemma4EC-2B, mean ± SD over the 10 seeds	0.903 ± 0.005
These weights (seed 6)	0.912 (precision 0.912, recall 0.917)

Gemma4EC outperforms BEC-Pred (0.879), Theia (0.877), CLAIRE (0.864), SelenzymeRF (0.675), SIMMER (0.610) and few-shot Gemini 3.1 Flash-Lite (0.655) and GPT-4.1-mini (0.234) on this benchmark. Full tables and significance tests are in the GitHub repository.

Usage

The base model is gated: request access to google/txgemma-2b-predict once, then huggingface-cli login.

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

base = "google/txgemma-2b-predict"
tokenizer = AutoTokenizer.from_pretrained(base)
model = AutoModelForCausalLM.from_pretrained(base, torch_dtype=torch.bfloat16, device_map="auto")
model = PeftModel.from_pretrained(model, "PlanesLab/Gemma4EC-2B-Predict").eval()

PROMPT = (
    "Instructions: Given a reaction SMILES, predict the first three parts of its "
    "Enzyme Commission (EC) number (class.subclass.sub-subclass).\n"
    "EC numbers are four-part hierarchical codes (e.g., 1.3.2.4) describing the type "
    "of enzyme and its reaction mechanism.\n"
    "Only output the first three digits (e.g., 1.03.02) representing the enzyme's "
    "general class, subclass, and sub-subclass.\n"
    "Reaction SMILES: {reaction_smiles}\nAnswer:"
)

prompt = PROMPT.format(reaction_smiles="CCO.O=O>>CC=O.OO")
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=8, do_sample=False)
print(tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip())

The prompt above is the one used in training; a different prompt will degrade accuracy. For batch prediction, explanations of the predictions and the full benchmark, use predict.py and chat/explain.sh in the repository.

Training

• Parameter-efficient fine-tuning with LoRA (r = 16, alpha = 16, dropout = 0.1) on the attention and MLP projections (q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj).
• Trained on the training split of seed 6 of the stratified ECReact benchmark (~45.9k reaction–EC pairs), with early stopping on validation MCC.
• The base model was loaded in bfloat16 during training.

Citation

Please cite the Gemma4EC paper (see the repository) and TxGemma.

License

The adapter weights are released under Apache-2.0. The base TxGemma weights are subject to the Health AI Developer Foundations terms.