README.md

---
license: mit
library_name: transformers
inference: false
tags:
  - genomics
  - human-genetics
  - genotype
  - snp
  - modernbert
  - masked-language-modeling
  - uk-biobank
  - allele-dosage
---

# SNPBERT

Self-supervised transformer that learns individual-level representations from unlabeled
genotype data. Each person is encoded as an ordered sequence of **rsID tokens** — one per
locus carrying at least one alternate allele — and the model is pretrained by masked
language modeling (masked rsID prediction) on **485,923 UK Biobank participants across
13,615 SNPs**. Built on the [ModernBERT](https://huggingface.co/docs/transformers/model_doc/modernbert)
architecture; its alternating local/global attention captures both short-range linkage
disequilibrium and genome-wide structure in a single forward pass.

**Code / reproduction:** https://github.com/BioLab-D/SNPBERT

## Two backbones in this repo

This repo ships **two pretrained backbones** as subfolders:

| Subfolder | Model | Genotype encoding | Load with | Masked-SNP acc. |
|---|---|---|---|---|
| [`dominant/`](./dominant) | `ModernBertForMaskedLM` | presence/absence (≥1 alt allele) | `AutoModel` + `subfolder="dominant"` | 0.785 |
| [`dosage/`](./dosage) | `ModernBertWithDosage` (base + dosage embedding) | het (1) vs. hom-alt (2) | **custom class** + `subfolder="dosage"` | 0.787 |

The **dosage** backbone adds an additive per-token dosage embedding on top of the dominant
architecture, so it distinguishes heterozygous from homozygous-alternate genotypes.

- **Architecture:** ModernBERT — 22 layers, hidden 768, 12 heads, intermediate 1152,
  global attention every 3rd layer, max position 9936, vocab 13,620.
- **Precision:** bfloat16.

> **Precision note:** weights are bfloat16 (trained on GPU). On **CPU**, cast to float32
> (`.float()`) — CPU LayerNorm does not support bfloat16. On GPU, use
> `torch_dtype=torch.bfloat16`.

## Tokenizer conventions (shared by both)

The vocabulary reserves the **first 5 ids** for special tokens, then SNP column *i* maps
to token id `i + 5`:

```
0=[UNK]  1=[SEP]  2=[PAD]  3=[CLS]  4=[MASK]   5..=SNP tokens (e.g. rs11260596_T)
```

Each SNP token is `rsID_allele`. An individual is a sequence
`[CLS] <SNP tokens for loci with ≥1 alt allele> [SEP]`.

---

## Usage — dominant backbone

Standard `ModernBertForMaskedLM`; loads directly with `AutoModel`.

```python
import torch
from transformers import AutoModel, AutoTokenizer

device = "cuda" if torch.cuda.is_available() else "cpu"
tok = AutoTokenizer.from_pretrained("ccjsj/SNPBERT", subfolder="dominant")
model = AutoModel.from_pretrained("ccjsj/SNPBERT", subfolder="dominant",
                                  attn_implementation="sdpa").to(device)
model = model.bfloat16() if device == "cuda" else model.float()
model.eval()

# One individual: [CLS] + SNP tokens (loci with >=1 alt allele) + [SEP]
snps = ["rs11260596_T", "rs6603782_T", "rs11523819_A"]
ids = [tok.cls_token_id] + tok.convert_tokens_to_ids(snps) + [tok.sep_token_id]
input_ids = torch.tensor([ids], device=device)
attention_mask = torch.ones_like(input_ids)

with torch.no_grad():
    out = model(input_ids=input_ids, attention_mask=attention_mask)
cls_embedding = out.last_hidden_state[:, 0]   # (1, 768)
```

These `[CLS]` embeddings feed the paper's downstream tasks (population-structure
clustering, disease-risk AMIL, variant identification).

---

## Usage — dosage backbone

### ⚠️ Not a plain `AutoModel` — custom class required

The `dosage/` checkpoint is **not** a standard `ModernBertForMaskedLM`. Its flat state
dict is:

```
base.*              # a ModernBertForMaskedLM  (138 tensors)
dosage_emb.weight   # nn.Embedding(3, 768)     (the dosage embedding)
```

`config.json` still declares `architectures: ["ModernBertForMaskedLM"]`, so loading with
`AutoModel.from_pretrained(...)` will **silently drop `dosage_emb.weight`**. You must
rebuild the `ModernBertWithDosage` wrapper. Its forward adds the dosage embedding to the
token embeddings **before** the encoder:

```
inputs_embeds = base.get_input_embeddings()(input_ids) + dosage_emb(dosage_ids)
```

`dosage_ids` is per-token in `{0, 1, 2}` — `0` at special positions, `1` (heterozygous) /
`2` (homozygous-alternate) at SNP positions.

### Option A — install the SNPBERT package

```bash
pip install "git+https://github.com/BioLab-D/SNPBERT.git"
```

```python
import torch
from huggingface_hub import snapshot_download
from snpbert.model import load_dosage_model

root = snapshot_download("ccjsj/SNPBERT")     # downloads both subfolders
ckpt = f"{root}/dosage"
device = "cuda" if torch.cuda.is_available() else "cpu"

model = load_dosage_model(ckpt, ckpt, num_dosage=3, attn_impl="sdpa").to(device)
model = model.bfloat16() if device == "cuda" else model.float()  # CPU needs fp32
model.eval()
```

### Option B — self-contained (no repo install)

```python
import torch, torch.nn as nn
from transformers import AutoConfig, AutoModelForMaskedLM, AutoTokenizer
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

class ModernBertWithDosage(nn.Module):
    def __init__(self, base, num_dosage=3):
        super().__init__()
        self.base = base
        self.dosage_emb = nn.Embedding(num_dosage, base.config.hidden_size)
    def forward(self, input_ids, dosage_ids=None, attention_mask=None, labels=None):
        emb = self.base.get_input_embeddings()(input_ids)
        if dosage_ids is not None:
            emb = emb + self.dosage_emb(dosage_ids)
        return self.base(inputs_embeds=emb, attention_mask=attention_mask, labels=labels)

repo = "ccjsj/SNPBERT"
cfg = AutoConfig.from_pretrained(repo, subfolder="dosage")
cfg._attn_implementation = "sdpa"
model = ModernBertWithDosage(AutoModelForMaskedLM.from_config(cfg), num_dosage=3)
sd = load_file(hf_hub_download(repo, "model.safetensors", subfolder="dosage"))
model.load_state_dict(sd, strict=True)        # keys: base.* + dosage_emb.weight
model.float().eval()
```

### Forward with dosage

```python
tok = AutoTokenizer.from_pretrained("ccjsj/SNPBERT", subfolder="dosage")
snps = ["rs11260596_T", "rs6603782_T", "rs11523819_A"]
ids  = [tok.cls_token_id] + tok.convert_tokens_to_ids(snps) + [tok.sep_token_id]
input_ids  = torch.tensor([ids])
attn       = torch.ones_like(input_ids)
# per-token dosage: 0 at [CLS]/[SEP]; 1=het, 2=hom-alt at SNP positions
dosage_ids = torch.tensor([[0, 1, 2, 1, 0]])

with torch.no_grad():
    out = model(input_ids=input_ids, dosage_ids=dosage_ids, attention_mask=attn)
logits = out.logits              # (1, seq_len, 13620)
```

---

## Downstream use

The frozen backbone feeds three analyses in the paper: population-structure clustering
(matching PCA), disease-risk modeling with a sparse-gated attention-MIL head, and
attention-based variant prioritization. Pipelines and a runnable synthetic-data demo are
in the [GitHub repo](https://github.com/BioLab-D/SNPBERT).

## License

MIT. Trained on UK Biobank data under approved access; use of derived representations
must comply with UK Biobank terms.