Taykhoom/DNABERT-S

DNABERT-S

Weights and tokenizer for DNABERT-S (Zhou et al., Bioinformatics 2025), loaded with the shared MosaicBERT implementation from Taykhoom/MosaicBERT-updated.

DNABERT-S is a species-aware DNA embedding model fine-tuned from DNABERT-2 using curriculum contrastive learning. It generates embeddings that naturally cluster and segregate genomes from different species, enabling species identification, metagenomics binning, and evolutionary analysis.

This repo contains only weights and tokenizer files. The model code is loaded automatically from Taykhoom/MosaicBERT-updated via trust_remote_code=True.

Architecture

Tokenization

Uses Byte Pair Encoding (BPE) tokenization via PreTrainedTokenizerFast, identical vocabulary to DNABERT-2. No k-mer pre-processing required.

Pretraining

• —Objective: Curriculum contrastive learning (same-species pairs with i-Mix)
• —Initialization: Fine-tuned from zhihan1996/DNABERT-2-117M
• —Source checkpoint: pytorch_model.bin from zhihan1996/DNABERT-S

Parity Verification

Hidden-state representations verified identical (max abs diff = 0.00) to the original implementation at all 13 representation levels (embedding + 12 transformer layers). SDPA verified (max abs diff < 1e-4). Verified on GPU with PyTorch 2.7 / CUDA 12.9.

Related Models

See the full DNABERT collection.

Usage

Embedding generation

The current mean-pooling policy excludes padding, [CLS], and the terminal [SEP] token.

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/DNABERT-S", trust_remote_code=True)
model = AutoModel.from_pretrained("Taykhoom/DNABERT-S", trust_remote_code=True)
model.eval()

sequences = ["ACGTAGCATCGGATCTATCTATCGACACTTGG", "ATCGATCGATCGATCG"]
enc = tokenizer(sequences, return_tensors="pt", padding=True)

with torch.no_grad():
    out = model(**enc)

cls_emb = out.last_hidden_state[:, 0, :]   # (batch, 768)

pool_mask = enc["attention_mask"].clone()
pool_mask[:, 0] = 0  # exclude [CLS]
sep_positions = enc["attention_mask"].sum(dim=1) - 1
pool_mask[torch.arange(pool_mask.size(0)), sep_positions] = 0  # exclude [SEP]

pool_mask = pool_mask.unsqueeze(-1).to(out.last_hidden_state.dtype)
mean_emb = (out.last_hidden_state * pool_mask).sum(dim=1)
mean_emb = mean_emb / pool_mask.sum(dim=1).clamp_min(1)  # (batch, 768)

Attention implementation

# SDPA (default on PyTorch >= 2.0)
model = AutoModel.from_pretrained("Taykhoom/DNABERT-S", trust_remote_code=True,
                                   attn_implementation="sdpa")

# Flash Attention 2
model = AutoModel.from_pretrained("Taykhoom/DNABERT-S", trust_remote_code=True,
                                   attn_implementation="flash_attention_2",
                                   torch_dtype=torch.bfloat16)

Implementation Notes

DNABERT-S is an embedding checkpoint and does not contain a trained masked-language-model head. Load it with AutoModel; AutoModelForMaskedLM raises a clear error rather than initializing random prediction weights.

The original DNABERT-S codebase uses a Triton-based flash attention implementation (flash_attn_triton.py). This HF port uses Taykhoom/MosaicBERT-updated which replaces it with the standard flash-attn package, and also adds attn_implementation="sdpa" support. These were not part of the original codebase.

Citation

@article{zhou2025_dnaberts,
  title   = {{DNABERT}-S: Pioneering Species Differentiation with Species-Aware {DNA} Embeddings},
  author  = {Zhou, Zhihan and Wu, Weimin and Ho, Harrison and Wang, Jiayi and Shi, Lizhen and Davuluri, Ramana V. and Wang, Zhong and Liu, Han},
  journal = {Bioinformatics},
  volume  = {41},
  number  = {Supplement_1},
  pages   = {i255--i264},
  year    = {2025},
  doi     = {10.1093/bioinformatics/btaf188}
}

Credits

Original DNABERT-S model and code by Zhou et al. Source: GitHub. Hugging Face port maintained by Taykhoom Dalal.

License

Apache 2.0, following the original repository.