DNABERT-4mer

DNABERT-4mer

Minimal HuggingFace port of the 4-mer variant of DNABERT -- a BERT-base masked language model pre-trained on the human reference genome using overlapping 4-mer tokenization.

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

Architecture

Standard BERT-base with a 4-mer DNA vocabulary.

Tokenization

Input sequences must be pre-split into overlapping 4-mers (stride 1) with spaces between tokens before calling the tokenizer. For example:

ATCGATG  ->  ATCG TCGA CGAT GATG

def seq_to_kmers(seq, k=4):
    return " ".join(seq[i:i+k] for i in range(len(seq) - k + 1))

Pretraining

• Objective: Masked Language Modeling
• Data: Human reference genome (GRCh38)
• Source checkpoint: pytorch_model.bin from zhihan1996/DNA_bert_4

Parity Verification

All 13 representation levels (embedding + 12 transformer layers) verified against the source implementation (max abs diff = 9.79e-5); MLM logits match with max abs diff = 2.52e-4. The source dnabert_layer.BertModel is a direct subclass of transformers.BertModel with no modifications. Verified on GPU with PyTorch 2.7.1 / CUDA 12.9.

Related Models

See the full DNABERT collection.

Usage

Embedding generation

import torch
from transformers import AutoTokenizer, AutoModel

def seq_to_kmers(seq, k=4):
    return " ".join(seq[i:i+k] for i in range(len(seq) - k + 1))

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

sequences = ["ATCGATCGATCG", "GCTAGCTAGCTA"]
kmer_seqs = [seq_to_kmers(s) for s in sequences]
enc = tokenizer(kmer_seqs, return_tensors="pt", padding=True)

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

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

# Mean-pool DNA k-mers only (exclude CLS, SEP, and padding)
content_mask = enc["attention_mask"].bool()
content_mask[:, 0] = False
sep_positions = enc["attention_mask"].sum(dim=1) - 1
batch_indices = torch.arange(len(sequences), device=content_mask.device)
content_mask[batch_indices, sep_positions] = False
mean_emb = (
    token_emb * content_mask.unsqueeze(-1)
).sum(dim=1) / content_mask.sum(dim=1, keepdim=True)

# Intermediate layers
out_all = model(**enc, output_hidden_states=True)
layer6_emb = out_all.hidden_states[6]

Sequences shorter than 4 nucleotides contain no k-mer tokens and therefore cannot be mean-pooled or assigned a k-mer pseudo-likelihood.

MLM logits

from transformers import AutoModelForMaskedLM

model = AutoModelForMaskedLM.from_pretrained(
    "Taykhoom/DNABERT-4mer", trust_remote_code=True
)
tokens = seq_to_kmers("ATCGATCG", k=4).split()
tokens[2] = tokenizer.mask_token
enc = tokenizer(" ".join(tokens), return_tensors="pt")

with torch.no_grad():
    logits = model(**enc).logits   # (1, seq_len, 261)

Faster attention backends

# SDPA (PyTorch 2.0+)
model = AutoModel.from_pretrained("Taykhoom/DNABERT-4mer", trust_remote_code=True,
                                   attn_implementation="sdpa")

# Flash Attention 2 (requires flash-attn)
model = AutoModel.from_pretrained("Taykhoom/DNABERT-4mer", trust_remote_code=True,
                                   attn_implementation="flash_attention_2",
                                   dtype=torch.float16)

Fine-tuning

For sequence-level tasks, mean-pool only k-mer positions as above or use the CLS token embedding as input to a prediction head.

Implementation Notes

The original DNABERT codebase has BertModel as a thin subclass of transformers.BertModel with no modifications. This HF port uses Taykhoom/BERT-updated which adds attn_implementation="sdpa" and attn_implementation="flash_attention_2" support — these were not part of the original codebase.

Citation

@article{ji2021_dnabert,
  title   = {{DNABERT}: pre-trained Bidirectional Encoder Representations from Transformers model for {DNA}-language in genome},
  author  = {Ji, Yanrong and Zhou, Zhihan and Liu, Han and Davuluri, Ramana V},
  journal = {Bioinformatics},
  volume  = {37},
  number  = {15},
  pages   = {2112--2120},
  year    = {2021},
  doi     = {10.1093/bioinformatics/btab083}
}

Credits

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

License

Apache License 2.0, following the original repository.