Agro Nucleotide Transformer 1b

🚀 AgroNT - 植物基因组DNA语言模型

AgroNT是一个主要基于可食用植物基因组训练的DNA语言模型。它运用Transformer架构和自注意力机制，通过掩码语言建模目标，利用48种不同植物的基因型数据，学习核苷酸序列的通用表示。该模型拥有10亿个参数，上下文窗口为1024个标记。

许可证

本项目采用CC BY-NC-SA 4.0许可证。

模型信息

🚀 快速开始

模型概述

AgroNT是一个主要在可食用植物基因组上训练的DNA语言模型。具体而言，AgroNT使用具有自注意力机制的Transformer架构和掩码语言建模目标，利用来自48种不同植物物种的高可用性基因型数据，学习核苷酸序列的通用表示。AgroNT包含10亿个参数，上下文窗口为1024个标记。

AgroNT使用非重叠的6聚体分词器将基因组核苷酸序列转换为标记。因此，1024个标记大约对应6144个碱基对。

📦 安装指南

文档未提及具体安装步骤，可参考代码示例从Hugging Face仓库加载模型和分词器。

💻 使用示例

基础用法

from transformers import AutoModelForMaskedLM, AutoTokenizer
import torch


model_name = 'agro-nucleotide-transformer-1b'

# fetch model and tokenizer from InstaDeep's hf repo
agro_nt_model = AutoModelForMaskedLM.from_pretrained(f'InstaDeepAI/{model_name}')
agro_nt_tokenizer = AutoTokenizer.from_pretrained(f'InstaDeepAI/{model_name}')

print(f"Loaded the {model_name} model with {agro_nt_model.num_parameters()} parameters and corresponding tokenizer.")

# example sequence and tokenization
sequences = ['ATATACGGCCGNC','GGGTATCGCTTCCGAC']

batch_tokens = agro_nt_tokenizer(sequences,padding="longest")['input_ids']
print(f"Tokenzied sequence: {agro_nt_tokenizer.batch_decode(batch_tokens)}")

torch_batch_tokens = torch.tensor(batch_tokens)
attention_mask = torch_batch_tokens != agro_nt_tokenizer.pad_token_id

# inference
outs = agro_nt_model(
    torch_batch_tokens,
    attention_mask=attention_mask,
    encoder_attention_mask=attention_mask,
    output_hidden_states=True
)

# get the final layer embeddings and language model head logits
embeddings = outs['hidden_states'][-1].detach().numpy()
logits = outs['logits'].detach().numpy()

📚 详细文档

预训练

数据

我们的预训练数据集主要基于Ensembl Plants数据库中的可食用植物参考基因组构建。该数据集包含48个不同物种的约1050万个基因组序列。

处理

每个物种的所有参考基因组被组装成一个单一的FASTA文件。在这个FASTA文件中，除A、T、C、G之外的所有核苷酸都被替换为N。使用分词器将字母字符串转换为标记序列。

分词器的字母表由通过组合A、T、C、G获得的4⁶ = 4096种可能的6聚体组合，以及代表独立A、T、C、G和N的五个额外标记组成。它还包括三个特殊标记：填充标记[PAD]、掩码标记[MASK]和分类标记[CLS]。这导致词汇表包含4104个标记。为了对输入序列进行分词，分词器从分类标记开始，然后从左到右转换序列，尽可能匹配6聚体标记，必要时使用独立标记（例如，当存在字母N或序列长度不是6的倍数时）。

分词示例

核苷酸序列: ATCCCGGNNTCGACACN 标记: <CLS> <ATCCCG> <G> <N> <N> <TCGACA> <C> <N>

训练

使用掩码语言建模（MLM）目标以自监督方式对AgroNT进行预训练。在自监督学习设置中，不需要每个序列的注释（监督），因为我们可以对序列的一部分进行掩码，并使用未掩码部分的信息来预测掩码位置。

这使我们能够利用大量未标记的基因组测序数据。具体而言，输入序列中15%的标记被选择进行增强，其中80%被掩码标记替换，10%被词汇表中的另一个标记随机替换，最后10%保持相同的标记。

分词后的序列通过模型，对掩码标记计算交叉熵损失。预训练使用1024个标记的序列长度和150万个标记的有效批量大小进行315000次更新步骤，导致模型总共训练了4725亿个标记。

硬件

模型预训练使用Google TPU - V4加速器进行，具体是包含512个设备的TPU v4 - 1024。我们总共训练了大约四天。

BibTeX引用

@article{mendoza2023foundational,
  title={A Foundational Large Language Model for Edible Plant Genomes},
  author={Mendoza-Revilla, Javier and Trop, Evan and Gonzalez, Liam and Roller, Masa and Dalla-Torre, Hugo and de Almeida, Bernardo P and Richard, Guillaume and Caton, Jonathan and Lopez Carranza, Nicolas and Skwark, Marcin and others},
  journal={bioRxiv},
  pages={2023--10},
  year={2023},
  publisher={Cold Spring Harbor Laboratory}
}