README.md

---
license: apache-2.0
tags:
  - biology
  - genomics
  - dna
  - masked-lm
library_name: transformers
pipeline_tag: fill-mask
---

# NucEngram

A **genomic language model (GLM)** for DNA sequences. It is a character-level
(single-nucleotide) model over the alphabet `A / C / G / T / N`, pretrained with
masked language modeling on genomic sequence, built on a ModernBERT encoder with
an **8192-nucleotide context window**.

It produces a **per-nucleotide contextual embedding** that you can pool and use
as features for downstream genomics tasks — promoter / splice-site / enhancer /
regulatory-element classification, sequence property prediction, etc. — either as
a frozen feature extractor or by fine-tuning.

The custom architecture ships with the repo, so load it with
`trust_remote_code=True`.

## Available sizes

Four sizes are released. They share the **same architecture, tokenizer and 8192-nt
context** and differ only in the transformer's **width (hidden size)** and
**depth (layers)** — i.e. capacity and compute:

| variant | hidden | layers | heads | parameters | download | how to load |
|---------|:------:|:------:|:-----:|:----------:|:--------:|-------------|
| `mini`  | 384    | 8      | 6     | ~150M      | 0.60 GB  | `subfolder="mini"` |
| `base`  | 512    | 22     | 16    | ~229M      | 0.92 GB  | *(default — repo root)* |
| `pro`   | 768    | 24     | 12    | ~364M      | 1.46 GB  | `subfolder="pro"` |
| `max`   | 1024   | 24     | 16    | ~542M      | 2.17 GB  | `subfolder="max"` |

Rule of thumb: **`mini`** is the fastest / lightest and a good default for large
screens or limited GPU memory; **`max`** gives the strongest representations at
the highest compute cost; `base` / `pro` sit in between. Same API for all.

```python
from transformers import AutoModel

# base (default, repo root)
base = AutoModel.from_pretrained("FreakingPotato/NucEngram", trust_remote_code=True)

# any other size via subfolder
mini = AutoModel.from_pretrained("FreakingPotato/NucEngram", subfolder="mini", trust_remote_code=True)
pro  = AutoModel.from_pretrained("FreakingPotato/NucEngram", subfolder="pro",  trust_remote_code=True)
maxm = AutoModel.from_pretrained("FreakingPotato/NucEngram", subfolder="max",  trust_remote_code=True)
```

## Install

```bash
pip install "transformers>=4.44" torch safetensors
```

## Quick start — embeddings

```python
from transformers import AutoModel

model = AutoModel.from_pretrained("FreakingPotato/NucEngram",
                                  trust_remote_code=True).eval()

# convenience helper: sequence(s) -> pooled embedding [B, hidden]
emb = model.embed(["ACGTACGTACGTGGTAAGT", "TTGCCGCGCGATCGATCG"])
print(emb.shape)   # torch.Size([2, 512])   (512 = base hidden size)
```

Per-nucleotide hidden states (for token-level tasks):

```python
ids, attention_mask = model.encode("ACGT...")     # char-level tokenizer, pad id 0
out = model(ids, attention_mask)
h = out.last_hidden_state                          # [B, T, hidden]
```

## Downstream task — fine-tuning

Add a pooling + linear head and fine-tune (or freeze `base` for linear probing).
Swap the `subfolder=` argument to choose a size:

```python
import torch, torch.nn as nn
from transformers import AutoModel

class SequenceClassifier(nn.Module):
    def __init__(self, n_classes, size=None, freeze_base=False):
        super().__init__()
        kw = {"trust_remote_code": True}
        if size:                       # None -> base (root); else "mini"/"pro"/"max"
            kw["subfolder"] = size
        self.base = AutoModel.from_pretrained("FreakingPotato/NucEngram", **kw)
        hidden = self.base.config.hidden_size
        if freeze_base:
            for p in self.base.parameters():
                p.requires_grad_(False)
        self.head = nn.Linear(hidden, n_classes)

    def forward(self, input_ids, attention_mask):
        h = self.base(input_ids, attention_mask).last_hidden_state   # [B, T, hidden]
        m = attention_mask.unsqueeze(-1).float()
        pooled = (h * m).sum(1) / m.sum(1).clamp(min=1.0)            # mean-pool
        return self.head(pooled)

clf = SequenceClassifier(n_classes=2, size="mini").train()
ids, am = clf.base.encode(["ACGT...", "GGGT..."])   # your batch of sequences
logits = clf(ids, am)
# ... standard cross-entropy training loop on your labelled dataset ...
```

For masked-LM scoring / filling:

```python
from transformers import AutoModelForMaskedLM
mlm = AutoModelForMaskedLM.from_pretrained("FreakingPotato/NucEngram",
                                           trust_remote_code=True).eval()
ids, am = mlm.encode("ACGTACGT")
logits = mlm(ids, am).logits        # [B, T, 9] over A/C/G/T/N + special tokens
```

## Details

| | |
|---|---|
| Backbone | ModernBERT encoder (see the size table above) |
| Context | up to 8192 nucleotides |
| Vocabulary | 9 tokens (A, C, G, T, N + pad/bos/eos/mask), char-level |
| Attention | `sdpa` by default (no flash-attn required) |
| Precision | fp32 weights (cast with `.half()` / `.bfloat16()` as you like) |

Input sequences are uppercase DNA strings; the built-in `encode()` maps
characters to ids and pads with id 0. Use `attention_mask` to ignore padding.