xxl0001/COD-PlantCAD-Rice

COD PlantCAD for Rice

COD PlantCAD Rice is a single-nucleotide DNA language model for Oryza sativa. It uses the PlantCaduceus l28 architecture and was selected from the continued rice training run at step 12,000.

The model is a masked language model with bidirectional Mamba blocks and reverse complement parameter sharing (RCPS). It accepts 512 bp DNA sequences and can be used for nucleotide logits, sequence representations, and SNV effect scoring.

Model details

Property	Value
Architecture	Caduceus masked language model
Base architecture	PlantCaduceus l28
Release checkpoint	Rice continued-training step 12,000
Layers	28
Hidden size	768 per orientation
Raw RCPS hidden size	1,536
Parameters	112,107,264
Input length	512 bp
Vocabulary	a/c/g/t plus PAD, MASK, and UNK
Reference species	Oryza sativa
Project reference assembly	R498/IRGSP-1.0 (osa1_r7.asm.chrs.fa)
Weight file	pytorch_model.bin (FP32)

No distillation-state JSON is included in this release.

Requirements

The custom Caduceus implementation requires PyTorch, Transformers, Mamba-SSM, and a compatible CUDA/Triton installation. The model uses fused Mamba/Triton operations and is intended for GPU inference.

The release was tested in the existing 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 that environment without reinstalling dependencies:

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

This repository is public, so download and inference do not require a Hugging Face token.

Download

hf download xxl0001/COD-PlantCAD-Rice \
  --local-dir ./COD-PlantCAD-Rice

The examples below can use either the Hub repository ID or the downloaded path.

Load with Transformers

Use AutoModel for sequence representations:

import torch
from transformers import AutoModel, AutoTokenizer

repo_id = "xxl0001/COD-PlantCAD-Rice"
device = torch.device("cuda")

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

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.
raw_hidden_states = outputs.last_hidden_state  # [batch, 512, 1536]

# Align and average the two orientations.
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  # [batch, 512, 768]

Use AutoModelForMaskedLM when nucleotide logits or variant scores are needed.

Score an SNV

The project uses the REF-to-ALT log-likelihood-ratio direction:

LLR(REF>ALT) = log P(ALT | sequence context) - log P(REF | sequence context)

A more negative score means the ALT allele is less supported than the REF allele in the same sequence context. The following function supports biallelic A/C/G/T SNVs and either of the project's two scoring protocols:

• masked=True: replace the variant position with [MASK]. This is the protocol used for the rice labeled-variant PR-AUC evaluation.
• masked=False: retain the reference nucleotide. This is the unmasked LLR protocol used for the genome-wide score collection.

Do not compare or combine scores from the two protocols as if they were the same distribution.

import torch
from transformers import AutoModelForMaskedLM, AutoTokenizer

repo_id = "xxl0001/COD-PlantCAD-Rice"
device = torch.device("cuda")

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(sequence_512, variant_index, ref, alt, masked=True):
    """Return ALT-logit minus REF-logit for one biallelic SNV.

    sequence_512: reference DNA sequence containing exactly 512 bases
    variant_index: zero-based SNV position inside sequence_512
    ref/alt: one of A, C, G, or T
    masked: use the center-masked protocol when True
    """
    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}"
        )

    input_ids = 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"].to(device)
    if input_ids.shape != (1, 512):
        raise RuntimeError(f"Unexpected tokenized shape: {tuple(input_ids.shape)}")

    if masked:
        input_ids[0, variant_index] = tokenizer.mask_token_id

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

    site_logits = logits[0, variant_index].float()
    return (
        site_logits[base_token_ids[alt]]
        - site_logits[base_token_ids[ref]]
    ).item()


sequence = "ACGT" * 128  # A at zero-based index 256
masked_llr = score_snv(sequence, 256, ref="A", alt="G", masked=True)
unmasked_llr = score_snv(sequence, 256, ref="A", alt="G", masked=False)
print({"masked_llr": masked_llr, "unmasked_llr": unmasked_llr})

For this release checkpoint, the example returns approximately:

masked_llr   = -7.037059
unmasked_llr = -9.671406

Score a rice genomic coordinate

The project reference FASTA is osa1_r7.asm.chrs.fa, with chromosome names Chr1 through Chr12. VCF positions are one-based. This helper extracts a centered 512 bp reference window with the SNV at zero-based index 256, matching the masked PR-AUC workflow:

from pyfaidx import Fasta


def get_centered_window(fasta, chrom, pos, window_size=512):
    pos0 = int(pos) - 1
    variant_index = window_size // 2
    start = pos0 - variant_index
    end = start + window_size
    chrom_length = len(fasta[chrom])

    if start < 0 or end > chrom_length:
        raise ValueError("Variant is too close to a chromosome boundary")

    sequence = str(fasta[chrom][start:end]).upper()
    if len(sequence) != window_size:
        raise RuntimeError("Failed to extract a 512 bp window")
    return sequence, variant_index


fasta = Fasta(
    "/root/private_data/xlxiang/MTEDF/huggingface/"
    "mtedf-rice-vep/reference_genome/osa1_r7.asm.chrs.fa",
    as_raw=True,
    sequence_always_upper=True,
)

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

sequence_512, variant_index = get_centered_window(fasta, chrom, pos)
llr = score_snv(sequence_512, variant_index, ref, alt, masked=True)
print({"chrom": chrom, "pos": pos, "ref": ref, "alt": alt, "llr": llr})

Keep the reference-match check enabled. A mismatch usually means the coordinate base, chromosome naming, or reference assembly does not match the scoring data.

Score direction and allele-frequency folding

• The released score direction is logit(ALT) - logit(REF).
• Lower values indicate less model support for ALT relative to REF.
• For minor-allele analyses only, negate the score when ALT allele frequency is greater than 0.5, 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 512 bp context and masked/unmasked protocol.

Limitations

• This checkpoint supports 512 bp inputs; the examples do not pad shorter input.
• The documented variant function supports biallelic A/C/G/T SNVs only.
• The model was evaluated against the R498/IRGSP-1.0 project reference; verify coordinates and REF alleles before using another rice assembly.
• An LLR is a relative model score, not a calibrated probability of a phenotype, pathogenicity, or fitness effect.
• trust_remote_code=True executes code from this repository. Pin a reviewed revision for production use.

Citation

COD PlantCAD Rice 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}
}