Views
No views yet
mamba-ssm and causal-conv1d libraries for the core backbone. You can retrieve both genomic feature probabilities and sequence embeddings using the following snippet:1import torch
2from transformers import AutoTokenizer, AutoModel
3
4# Load model and tokenizer
5repo_id = "qzzhang/PlantGeneAnn-v1.5-flower-plants"
6tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
7model = AutoModel.from_pretrained(repo_id, trust_remote_code=True)
8
9# The number of DNA tokens (excluding the [CLS] and [SEP] token) needs to be divisible by 16
10# as required by the U-Net downsampling blocks.
11sequences = ["AATTCCGGAA"*4096,"AAAATTTTCC"*4096]
12
13# Tokenize
14tokenized_sequences = tokenizer(
15 sequences,
16 return_tensors="pt",
17 padding="longest"
18)["input_ids"]
19
20# Infer
21model.to("cuda")
22model.eval()
23with torch.no_grad():
24 outs = model(input_ids=tokenized_sequences.to("cuda"), return_dict=True)
25
26# Obtain the logits over the genomic features
27# Shape: [batch, sequence_length - 2 * flank_length, 2 * num_features]
28logits = outs.logits
29
30# Get the sequence embeddings
31# Shape: [batch, sequence_length - 2 * flank_length, 512]
32hidden_states = outs.hidden_states.detach()1@article{zhang2026plantgeneann,
2 title={PlantGeneAnn: a strand-specific genome foundation model for ab initio gene structure annotation of plant genomes},
3 author={Zhang, Qizhe and Zhang, Zhengyang and Lin, Kepeng and Wang, Jing and Deng, Kaixuan and Xiang, Xianglei and Xu, Wei and Hu, Xuehai},
4 journal={bioRxiv},
5 year={2026},
6 doi={10.64898/2026.06.25.733695},
7 url={[https://doi.org/10.64898/2026.06.25.733695](https://doi.org/10.64898/2026.06.25.733695)}
8}