README.md

---
library_name: transformers
pipeline_tag: fill-mask
tags:
  - genomics
  - dna
  - masked-lm
  - ntv3
  - long-range
license: other
language:
  - code
model_parameter_count: 595169467
---

<div style="background-color: rgba(255, 68, 68, 0.15); padding: 5px; border: 2px solid #ff4444; border-radius: 3px;">
<h3>⚠️ WARNING: Ablation Models Ahead</h3>
<p>This 5-downsample model structure is <strong>experimental</strong> and intended solely for exploration related to the model structure ablation studies.</p>
<p><strong>They are NOT the main, recommended NTv3 models for results.</strong></p>
</div>

## 🧬 NTv3: A Foundation Model for Genomics

NTv3 is a series of foundational models designed to understand and generate genomic sequences. It unifies representation learning, functional prediction, and controllable sequence generation within a single, efficient U-Net-like architecture. It also enables the modeling of long-range dependencies, up to 1 Mb of context, at nucleotide resolution. Pretrained on 9 trillion base pairs, NTv3 excels at functional-track prediction and genome annotation across 24 animal and plant species. It can also be fine-tuned into a controllable generative model for genomic sequence design. This repository contains the MLM pre-trained models and weights. For more details, please refer to the [NTv3 paper placeholder]. 

## ⚖️ License Summary

1. The Licensed Models are **only** available under this License for Non-Commercial Purposes.
2. You are permitted to reproduce, publish, share and adapt the Output generated by the Licensed Model only for Non-Commercial Purposes and in accordance with this License.
3. You may **not** use the Licensed Models or any of its Outputs in connection with:
   1. any Commercial Purposes, unless agreed by Us under a separate licence;
   2. to train, improve or otherwise influence the functionality or performance of any other third-party derivative model that is commercial or intended for a Commercial Purpose and is similar to the Licensed Models;
   3. to create models distilled or derived from the Outputs of the Licensed Models, unless such models are for Non-Commercial Purposes and open-sourced under the same license as the Licensed Models; or
   4. in violation of any applicable laws and regulations.

## 📋 Model Summary

- Architecture: U-Net style conv tower → Transformer stack → deconv tower → LM head
- Tokenizer: character-level over A T C G N + specials (`<unk>` `<pad>` `<mask>` `<cls>` `<eos>` `<bos>`)
- Selective intermediate outputs: use config to save specific layers
- Dependencies: needs transformers >= 4.55.0
- Input size: input sequence length need to be a multiple of 128
- Note: custom code → use `trust_remote_code=True`

## 🚀 Quickstart

```python
from transformers import AutoTokenizer, AutoModelForMaskedLM

repo = "InstaDeepAI/NTv3_5downsample_pre"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained(repo, trust_remote_code=True)

batch = tok(["ATCGNATCG", "ACGT"], add_special_tokens=False, padding=True, pad_to_multiple_of=128, return_tensors="pt")
out = model(**batch)

print(out.logits.shape)  # (B, L, V = 11)
```

## 🔤 Tokenization

```python
enc = tok("ATCGNATCG", add_special_tokens=False)
print(enc["input_ids"])  # char-level IDs
```

## 🔍 Getting hidden states and attentions

To get all hidden states and attention weights from all layers:

```python
out = model(**batch, output_hidden_states=True, output_attentions=True)

# Access all hidden states (tuple of tensors, one per layer)
hidden_states = out.hidden_states
print(len(hidden_states))  # Number of layers
print(hidden_states[0].shape)  # (B, L, 1536)

# Access all attention weights (tuple of tensors, one per transformer layer)
attentions = out.attentions
print(len(attentions))  # Number of transformer layers
print(attentions[0].shape)  # (B, H = 24, L, L)

# Get final embedding (after deconv tower)
final_emb = out.hidden_states[-1]  # shape (B, L, 1536)
```

## 🛠️ Selective intermediate outputs

You can also save specific intermediate outputs with custom keys:

```python
from ntv3_huggingface_new import Ntv3PreTrainedConfig

config = Ntv3PreTrainedConfig.from_pretrained(repo)
# Save embeddings from specific transformer layers
config.embeddings_layers_to_save = (1, 2)
# Save attention maps from specific layers/heads
config.attention_maps_to_save = [(1, 0), (2, 1)]  # (layer, head)
# Save embeddings from specific deconv layers
config.deconv_layers_to_save = (1, 2)

model = AutoModelForMaskedLM.from_pretrained(repo, config=config, trust_remote_code=True)
# Access via core's output dict (these are saved in addition to hidden_states/attentions)
core_out = model.core(**batch, output_hidden_states=True, output_attentions=True)
emb_1 = core_out['embeddings_1']  # Transformer layer 1
attn_1_0 = core_out['attention_map_layer_1_number_0']  # Layer 1, head 0
deconv_1 = core_out['embeddings_deconv_1']  # Deconv layer 1
```

## 📝 Getting input embeddings

```python
emb_layer = model.get_input_embeddings()  # nn.Embedding(V = 11, D = 16)
```

## 🎯 Masked LM training

```python
import torch
inputs = tok(["ATCGNATCG"], add_special_tokens=False, padding=True, pad_to_multiple_of=128, return_tensors="pt")
labels = inputs["input_ids"].clone(); labels[:] = -100
mask_id = tok.mask_token_id
inputs["input_ids"][0, 2] = mask_id
labels[0, 2] = tok.convert_tokens_to_ids("C")
out = model(**inputs, labels=labels)
print(out.loss.item())
```

## 📊 Shapes & config summary

| Parameter | Value |
|-----------|-------|
| Vocab size | 11 |
| Token embedding dim | 16 |
| Model (hidden) dim | 1536 |
| FFN dim | 6144 |
| Attention heads | 24 |
| Transformer layers | 12 |
| Downsample stages | 5 |


## ⚡ Mixed precision

This model was originally trained with mixed precision (bf16) in JAX and later ported to Torch. During JAX training, all weights maintained full fp32 precision at all times, but certain inferences were performed in bf16 for efficiency. This repo will be loaded with full precision (fp32) inference by default to ensure numerical stability. However, it can be used with mixed precision (bf16) for efficient long range training and inferences. Do note, to support bfloat16 precision, you need to use a GPU with bfloat16 support (e.g. A100, H100, etc.). Also, loading the model with mixed precision would introduce numerical instability, including small differences to the original JAX model. The difference is usually insignificant, but be aware of it when using the model.

To load the model with mixed precision, use the following code:

```python
from transformers import AutoTokenizer, AutoModelForMaskedLM

repo = "InstaDeepAI/NTv3_5downsample_pre"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained(
    repo, trust_remote_code=True,
    stem_compute_dtype='bfloat16',
    down_convolution_compute_dtype='bfloat16',
    transformer_qkvo_compute_dtype='bfloat16',
    transformer_ffn_compute_dtype='bfloat16',
    up_convolution_compute_dtype='bfloat16',
    modulation_compute_dtype='bfloat16',
)
```