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