README.md

---
license: apache-2.0
library_name: transformers
pipeline_tag: fill-mask
base_model: kuleshov-group/PlantCaduceus_l28
tags:
- biology
- genomics
- dna
- oryza-sativa
- rice
- variant-effect-prediction
- masked-language-modeling
- caduceus
- mamba
---

# 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:

```text
/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:

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

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

## Download

```bash
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:

```python
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:

```text
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.

```python
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:

```text
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:

```python
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:

```bibtex
@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}
}
```