Ntv3 650m Pre 8kb

🧬 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

from transformers import AutoTokenizer, AutoModelForMaskedLM

repo = "InstaDeepAI/NTv3_650M_pre_8kb"
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

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:

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:

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

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

🎯 Masked LM training

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

⚡ 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:

from transformers import AutoTokenizer, AutoModelForMaskedLM

repo = "InstaDeepAI/NTv3_650M_pre_8kb"
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',
)