COD PlantCAD for Arabidopsis

COD PlantCAD is an Arabidopsis thaliana DNA language model based on the PlantCaduceus_l28 architecture. This repository contains the project-selected final checkpoint from full-genome Arabidopsis distillation.

The model operates at single-nucleotide resolution with a maximum input length of 512 bp. It is a masked language model with bidirectional Mamba blocks and reverse complement parameter sharing (RCPS).

Model details

This release does not include a distillation-state JSON file. The checkpoint is the final Arabidopsis model selected for release by this project.

Tested environment

The examples below were tested in the existing server environment:

/root/private_data/miniconda3/envs/plantgenoann
Python 3.8.20
PyTorch 2.2.2+cu121
Transformers 4.38.1
Mamba-SSM 1.2.0.post1
NVIDIA RTX 4090

Activate it without reinstalling dependencies:

source /root/private_data/miniconda3/bin/activate plantgenoann

This repository is public. Downloading and loading the model does not require a Hugging Face account, access request, or authentication token.

Download the model

Download a complete local copy with the Hugging Face CLI:

/root/.local/bin/hf download \
  xxl0001/COD-PlantCAD-Arabidopsis \
  --local-dir ./COD-PlantCAD-Arabidopsis

The examples can use either the Hub repository ID or the downloaded directory:

repo_id = "xxl0001/COD-PlantCAD-Arabidopsis"
# Or: repo_id = "./COD-PlantCAD-Arabidopsis"

Standard Transformers import

Use AutoModel for sequence representations:

import torch
from transformers import AutoModel, AutoTokenizer

repo_id = "xxl0001/COD-PlantCAD-Arabidopsis"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
model = AutoModel.from_pretrained(repo_id, trust_remote_code=True)
model.to(device).eval()

# Exactly 512 DNA bases. The tokenizer adds no CLS or SEP token.
sequences = ["ACGT" * 128, "TGCA" * 128]
encoded = tokenizer(
    [sequence.lower() for sequence in sequences],
    add_special_tokens=False,
    padding=False,
    truncation=False,
    return_attention_mask=False,
    return_token_type_ids=False,
    return_tensors="pt",
)

input_ids = encoded["input_ids"].to(device)
assert input_ids.shape == (2, 512)

with torch.inference_mode():
    outputs = model(input_ids=input_ids, return_dict=True)

# RCPS concatenates forward and reverse-complement representations.
# Shape: [batch, 512, 1536]
raw_hidden_states = outputs.last_hidden_state

# Align reverse-complement coordinates and average both orientations.
# Shape: [batch, 512, 768]
hidden_size = raw_hidden_states.shape[-1] // 2
forward = raw_hidden_states[..., :hidden_size]
reverse_complement = raw_hidden_states[..., hidden_size:].flip(dims=(1,))
hidden_states = (forward + reverse_complement) / 2

print(raw_hidden_states.shape)  # torch.Size([2, 512, 1536])
print(hidden_states.shape)      # torch.Size([2, 512, 768])

This repository was also tested with AutoModelForMaskedLM.from_pretrained using trust_remote_code=True. Use the MLM class, rather than AutoModel, when nucleotide logits or variant scores are needed.

Score one SNV in a 512 bp sequence

COD PlantCAD's standard project score is an unmasked log-likelihood ratio (LLR):

LLR(REF>ALT) = log P(ALT | reference 512 bp sequence)
             - log P(REF | reference 512 bp sequence)

The input remains the reference sequence; the variant position is not replaced by the MASK token for this score. A more negative LLR means that the ALT allele is less supported than the REF allele in the same sequence context. This is the scoring protocol used for the project's genome-wide unmasked LLR results.

import torch
from transformers import AutoModelForMaskedLM, AutoTokenizer

repo_id = "xxl0001/COD-PlantCAD-Arabidopsis"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained(
    repo_id,
    trust_remote_code=True,
).to(device).eval()

base_token_ids = {
    base: tokenizer.convert_tokens_to_ids(base.lower())
    for base in "ACGT"
}


def score_snv_unmasked(sequence_512, variant_index, ref, alt):
    """Return log P(ALT) - log P(REF) for one biallelic SNV.

    sequence_512: reference DNA sequence containing exactly 512 bases
    variant_index: zero-based position of the SNV inside sequence_512
    ref/alt: one of A, C, G, or T
    """
    sequence = sequence_512.upper()
    ref = ref.upper()
    alt = alt.upper()

    if len(sequence) != 512:
        raise ValueError(f"Expected 512 bp, received {len(sequence)} bp")
    if not 0 <= variant_index < 512:
        raise ValueError("variant_index must be in [0, 511]")
    if ref not in base_token_ids or alt not in base_token_ids:
        raise ValueError("Only biallelic A/C/G/T SNVs are supported")
    if sequence[variant_index] != ref:
        raise ValueError(
            f"Reference mismatch: sequence has {sequence[variant_index]}, ref={ref}"
        )

    encoded = tokenizer(
        sequence.lower(),
        add_special_tokens=False,
        padding=False,
        truncation=False,
        return_attention_mask=False,
        return_token_type_ids=False,
        return_tensors="pt",
    )
    input_ids = encoded["input_ids"].to(device)
    if input_ids.shape != (1, 512):
        raise RuntimeError(f"Unexpected tokenized shape: {tuple(input_ids.shape)}")

    with torch.inference_mode():
        logits = model(input_ids=input_ids, return_dict=True).logits
        log_probs = logits[0, variant_index].float().log_softmax(dim=-1)

    ref_log_prob = log_probs[base_token_ids[ref]]
    alt_log_prob = log_probs[base_token_ids[alt]]
    return (alt_log_prob - ref_log_prob).item()


# Example: 512 bp with an A at zero-based index 256.
sequence = "ACGT" * 128
llr = score_snv_unmasked(
    sequence_512=sequence,
    variant_index=256,
    ref="A",
    alt="G",
)
print(f"A>G LLR: {llr:.6f}")

In the tested environment and current checkpoint, this example returns approximately -8.001844.

Score a genomic coordinate from TAIR10

The function below extracts the same type of non-overlapping 512 bp block used by the project's genome-wide unmasked scoring workflow. Input pos is one-based, as in VCF. The FASTA chromosome names must match chrom. The bundled TAIR10 FASTA uses 1 through 5 for nuclear chromosomes.

from pyfaidx import Fasta


def get_block_aligned_window(fasta, chrom, pos, window_size=512):
    """Return (reference_sequence, zero_based_index_in_window)."""
    pos0 = int(pos) - 1
    chrom_length = len(fasta[chrom])
    if not 0 <= pos0 < chrom_length:
        raise ValueError(f"Position {pos} is outside chromosome {chrom}")

    start = (pos0 // window_size) * window_size
    end = start + window_size

    # The genome-wide workflow uses the final 512 bp when the terminal block is short.
    if end > chrom_length:
        start = max(0, chrom_length - window_size)
        end = chrom_length

    sequence = str(fasta[chrom][start:end]).upper()
    if len(sequence) != window_size:
        raise ValueError("Chromosome is shorter than the requested 512 bp window")
    return sequence, pos0 - start


fasta = Fasta(
    "/root/private_data/xlxiang/MTEDF/huggingface/"
    "mtedf-arabidopsis-vep/reference_genome/TAIR10_Arabidopsis_thaliana.fa",
    as_raw=True,
    sequence_always_upper=True,
)

chrom = "1"
pos = 1000001
ref = "A"  # Replace with REF from the VCF.
alt = "G"  # Replace with ALT from the VCF.

sequence_512, variant_index = get_block_aligned_window(fasta, chrom, pos)
llr = score_snv_unmasked(sequence_512, variant_index, ref, alt)
print({"chrom": chrom, "pos": pos, "ref": ref, "alt": alt, "llr": llr})

Always keep the reference-match check enabled. A mismatch usually indicates a zero-based/one-based coordinate error, a chromosome naming mismatch, or a different reference-genome build.

Optional center-masked score

A center-masked LLR can be calculated by replacing the variant token with the MASK token before inference. This answers a different question:

masked LLR = log P(ALT | 512 bp with the variant position masked)
           - log P(REF | 512 bp with the variant position masked)

Insert these lines after tokenization in the preceding scoring function:

input_ids = encoded["input_ids"].to(device)
input_ids[0, variant_index] = tokenizer.mask_token_id

Do not mix masked and unmasked scores in one benchmark or percentile distribution. For the example above, the masked score is approximately -6.239229, which is not expected to equal the unmasked score.

Score direction and allele-frequency folding

• LLR = log P(ALT) - log P(REF) is oriented from REF to ALT.
• Lower or more negative values indicate less model support for ALT relative to REF.
• If an analysis must orient every score to the minor allele and ALT allele frequency is greater than 0.5, use minor_allele_llr = -llr because REF is then the minor allele.
• Do not apply that sign flip when retaining the original REF-to-ALT direction.
• Scores depend on the selected 512 bp context. Use the same window construction when comparing against cached genome-wide scores.

Limitations

• The documented variant function supports biallelic A/C/G/T SNVs only.
• Inputs longer than 512 bp are unsupported by this checkpoint.
• The model is a DNA language model; an LLR is not a calibrated probability of pathogenicity or fitness effect.
• Numerical scores from masked and unmasked protocols are not interchangeable.
• trust_remote_code=True executes the modeling files in this repository. Pin a reviewed revision for production workflows when reproducibility is required.

Citation

COD PlantCAD uses the PlantCaduceus architecture. Please cite the PlantCAD work:

@article{Zhai2025CrossSpecies,
  author  = {Zhai, Jingjing and Gokaslan, Aaron and Schiff, Yoni and Berthel,
             Alexander and Liu, Z. Y. and Lai, W. L. and Miller, Z. R. and
             Scheben, Armin and Stitzer, Michelle C. and Romay, Maria C. and
             Buckler, Edward S. and Kuleshov, Volodymyr},
  title   = {Cross-species modeling of plant genomes at single nucleotide
             resolution using a pretrained DNA language model},
  journal = {Proceedings of the National Academy of Sciences},
  year    = {2025},
  volume  = {122},
  number  = {24},
  pages   = {e2421738122},
  doi     = {10.1073/pnas.2421738122}
}