InstaDeepAI/segment_nt

1---2license: cc-by-nc-sa-4.03widget:4- text: ACCTGA<mask>TTCTGAGTC5tags:6- DNA7- biology8- genomics9- segmentation10---11# SegmentNT12 13SegmentNT is a segmentation model leveraging the [Nucleotide Transformer](https://huggingface.co/InstaDeepAI/nucleotide-transformer-v2-500m-multi-species) (NT) DNA foundation model to predict the location of several types of genomics 14elements in a sequence at a single nucleotide resolution. It was trained on 14 different classes of human genomics elements in input sequences up to 30kb. These 15include gene (protein-coding genes, lncRNAs, 5’UTR, 3’UTR, exon, intron, splice acceptor and donor sites) and regulatory (polyA signal, tissue-invariant and 16tissue-specific promoters and enhancers, and CTCF-bound sites) elements.17 18 19**Developed by:** [InstaDeep](https://huggingface.co/InstaDeepAI)20 21### Model Sources22 23<!-- Provide the basic links for the model. -->24 25- **Repository:** [Nucleotide Transformer](https://github.com/instadeepai/nucleotide-transformer)26- **Paper:** [Segmenting the genome at single-nucleotide resolution with DNA foundation models](https://www.biorxiv.org/content/biorxiv/early/2024/03/15/2024.03.14.584712.full.pdf) 27 28### How to use29 30<!-- Need to adapt this section to our model. Need to figure out how to load the models from huggingface and do inference on them -->31Until its next release, the `transformers` library needs to be installed from source with the following command in order to use the models:32```bash33pip install --upgrade git+https://github.com/huggingface/transformers.git34```35 36A small snippet of code is given here in order to retrieve both logits and embeddings from a dummy DNA sequence.37 38⚠️ The maximum sequence length is set by default at the training length of 30,000 nucleotides, or 5001 tokens (accounting for the CLS token). However,39SegmentNT-multi-species has been shown to generalize up to sequences of 50,000 bp. In case you need to infer on sequences between 30kbp and 50kbp, make sure to change40the `rescaling_factor` of the Rotary Embedding layer in the esm model  `num_dna_tokens_inference / max_num_tokens_nt` where `num_dna_tokens_inference` is the number of tokens at inference41(i.e 6669 for a sequence of 40008 base pairs) and `max_num_tokens_nt` is the max number of tokens on which the backbone nucleotide-transformer was trained on, i.e `2048`. 42 43[![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/#fileId=https%3A//huggingface.co/InstaDeepAI/segment_nt/blob/main/inference_segment_nt.ipynb)44The `./inference_segment_nt.ipynb` can be run in Google Colab by clicking on the icon and shows how to handle inference on sequence lengths require changing 45the rescaling factor and sequence lengths that do not. One can run the notebook and reproduce Fig.1 and Fig.3 from the SegmentNT paper. 46 47```python48# Load model and tokenizer49from transformers import AutoTokenizer, AutoModel50import torch51 52tokenizer = AutoTokenizer.from_pretrained("InstaDeepAI/segment_nt", trust_remote_code=True)53model = AutoModel.from_pretrained("InstaDeepAI/segment_nt", trust_remote_code=True)54 55# Choose the length to which the input sequences are padded. By default, the 56# model max length is chosen, but feel free to decrease it as the time taken to 57# obtain the embeddings increases significantly with it.58# The number of DNA tokens (excluding the CLS token prepended) needs to be dividible by59# 2 to the power of the number of downsampling block, i.e 4.60max_length = 12 + 161 62assert (max_length - 1) % 4 == 0, (63    "The number of DNA tokens (excluding the CLS token prepended) needs to be dividible by"64     "2 to the power of the number of downsampling block, i.e 4.")65 66# Create a dummy dna sequence and tokenize it67sequences = ["ATTCCGATTCCGATTCCG", "ATTTCTCTCTCTCTCTGAGATCGATCGATCGAT"]68tokens = tokenizer.batch_encode_plus(sequences, return_tensors="pt", padding="max_length", max_length = max_length)["input_ids"]69 70# Infer71attention_mask = tokens != tokenizer.pad_token_id72outs = model(73    tokens,74    attention_mask=attention_mask,75    output_hidden_states=True76)77 78# Obtain the logits over the genomic features79logits = outs.logits.detach()80# Transform them in probabilities81probabilities = torch.nn.functional.softmax(logits, dim=-1)82print(f"Probabilities shape: {probabilities.shape}")83 84# Get probabilities associated with intron85idx_intron = model.config.features.index("intron")86probabilities_intron = probabilities[:,:,idx_intron]87print(f"Intron probabilities shape: {probabilities_intron.shape}")88 89 90```91 92 93## Training data94 95The **SegmentNT** model was trained on all human chromosomes except for chromosomes 20 and 21, kept as test set, and chromosome 22, used as a validation set.96During training, sequences are randomly sampled in the genome with associated annotations. However, we keep the sequences in the validation and test set fixed by 97using a sliding window of length 30,000 over the chromosomes 20 and 21. The validation set was used to monitor training and for early stopping.98 99## Training procedure100 101### Preprocessing102 103The DNA sequences are tokenized using the Nucleotide Transformer Tokenizer, which tokenizes sequences as 6-mers tokens as described in the [Tokenization](https://github.com/instadeepai/nucleotide-transformer#tokenization-abc) section of the associated repository. This tokenizer has a vocabulary size of 4105. The inputs of the model are then of the form:104 105```106<CLS> <ACGTGT> <ACGTGC> <ACGGAC> <GACTAG> <TCAGCA>107```108 109### Training110 111The model was trained on a DGXH100 node with 8 GPUs on a total of 23B tokens for 3 days. The model was trained on 3kb, 10kb, 20kb and finally 30kb sequences, at each time with an effective batch size of 256 sequences. 112 113 114### Architecture115 116The model is composed of the [nucleotide-transformer-v2-500m-multi-species](https://huggingface.co/InstaDeepAI/nucleotide-transformer-v2-500m-multi-species) encoder, from which we removed 117the language model head and replaced it by a 1-dimensional U-Net segmentation head [4] made of 2 downsampling convolutional blocks and 2 upsampling convolutional blocks. Each of these 118blocks is made of 2 convolutional layers with 1, 024 and 2, 048 kernels respectively. This additional segmentation head accounts for 53 million parameters, bringing the total number of parameters119to 562M.120 121### BibTeX entry and citation info122 123```bibtex124@article{de2024segmentnt,125  title={SegmentNT: annotating the genome at single-nucleotide resolution with DNA foundation models},126  author={de Almeida, Bernardo P and Dalla-Torre, Hugo and Richard, Guillaume and Blum, Christopher and Hexemer, Lorenz and Gelard, Maxence and Pandey, Priyanka and Laurent, Stefan and Laterre, Alexandre and Lang, Maren and others},127  journal={bioRxiv},128  pages={2024--03},129  year={2024},130  publisher={Cold Spring Harbor Laboratory}131}132 133```