GENA-LM

GENA-LM 是一系列用于长DNA序列的开源基础模型。

GENA-LM 模型是在人类DNA序列上训练的transformer掩码语言模型。

我们的GENA-LM模型的主要特点：

• 使用BPE分词而非k-mers（DNABERT、Nucleotide Transformer）
• 最大输入序列长度范围从4.5k到36k碱基对，相比之下DNABERT为512碱基对，Nucleotide Transformer为1000碱基对
• 在最新的T2T人类基因组组装上进行预训练，而非GRCh38/hg38

预训练模型

模型	架构	最大序列长度，标记（碱基对）	参数数	分词器数据	训练数据
bert-base	BERT-12L	512(4500)	110M	T2T 拆分 v1	T2T 拆分 v1
bert-base-t2t	BERT-12L	512(4500)	110M	T2T+1000G SNPs+多物种	T2T+1000G SNPs
bert-base-lastln-t2t	BERT-12L	512(4500)	110M	T2T+1000G SNPs+多物种	T2T+1000G SNPs
bert-base-t2t-multi	BERT-12L	512(4500)	110M	T2T+1000G SNPs+多物种	T2T+1000G SNPs+多物种
bert-large-t2t	BERT-24L	512(4500)	336M	T2T+1000G SNPs+多物种	T2T+1000G SNPs
bigbird-base-sparse	BERT-12L, DeepSpeed稀疏操作, RoPE	4096(36000)	110M	T2T 拆分 v1	T2T 拆分 v1
bigbird-base-sparse-t2t	BERT-12L, DeepSpeed稀疏操作, RoPE	4096(36000)	110M	T2T+1000G SNPs+多物种	T2T+1000G SNPs
bigbird-base-t2t	BERT-12L, HF BigBird	4096(36000)	110M	T2T+1000G SNPs+多物种	T2T+1000G SNPs

T2T 拆分 v1 指的是使用未增强的T2T人类基因组组装拆分的初步模型。基于BERT的模型采用预层归一化，而lastln明确表示层归一化也应用于最后一层。RoPE表示使用旋转位置嵌入代替BERT类型的绝对位置嵌入。

对于我们的首批模型（gena-lm-bert-base和gena-lm-bigbird-base-sparse），我们保留人类染色体22和Y（CP068256.2和CP086569.2）作为掩码语言建模任务的测试数据集。对于所有其他模型，我们保留人类染色体7和10（CP068271.2和CP068268.2）；这些模型的名称中带有后缀"t2t"。其他数据用于训练。仅人类模型在预处理的Human T2T v2基因组组装及其1000基因组SNP增强数据上进行训练，总计约480 x 10^9碱基对。多物种模型在仅人类和多物种数据上进行训练，总计约1072 x 10^9碱基对。

下游任务上的预训练模型

模型	任务	任务序列长度	指标	HF 分支名称
gena-lm-bert-base-t2t	启动子	300bp	74.56+-0.36 F1	promoters_300_run_1
gena-lm-bert-large-t2t	启动子	300bp	76.44+-0.16 F1	promoters_300_run_1
gena-lm-bert-large-t2t	启动子	2000bp	93.70+-0.44 F1	promoters_2000_run_1
gena-lm-bert-base-t2t	剪接位点	15000bp	92.63+-0.09 PR AUC	spliceai_run_1
gena-lm-bert-large-t2t	剪接位点	15000bp	93.59+-0.11 PR AUC	spliceai_run_1

要获取下游任务上的预训练模型，请用表格中的值替换model_name和branch_name。表格中的指标是多次运行的平均值。因此，每个检查点的值可能与此处报告的有所不同。

from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained(f'AIRI-Institute/{model_name}')
model = AutoModel.from_pretrained(f'AIRI-Institute/{model_name}', revision=branch_name, trust_remote_code=True)

示例

如何加载预训练的GENA-LM用于掩码语言建模

from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained('AIRI-Institute/gena-lm-bert-base-t2t')
model = AutoModel.from_pretrained('AIRI-Institute/gena-lm-bert-base-t2t', trust_remote_code=True)

如何加载预训练的GENA-LM以在分类任务上进行微调

从GENA-LM仓库获取模型类:

git clone https://github.com/AIRI-Institute/GENA_LM.git

from GENA_LM.src.gena_lm.modeling_bert import BertForSequenceClassification
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained('AIRI-Institute/gena-lm-bert-base-t2t')
model = BertForSequenceClassification.from_pretrained('AIRI-Institute/gena-lm-bert-base-t2t')

或者你可以直接下载modeling_bert.py并将其放在你的代码附近。

或者你可以从HuggingFace AutoModel获取模型类:

from transformers import AutoTokenizer, AutoModel
model = AutoModel.from_pretrained('AIRI-Institute/gena-lm-bert-base-t2t', trust_remote_code=True)
gena_module_name = model.__class__.__module__
print(gena_module_name)
import importlib
# 可用的类名:
# - BertModel, BertForPreTraining, BertForMaskedLM, BertForNextSentencePrediction,
# - BertForSequenceClassification, BertForMultipleChoice, BertForTokenClassification,
# - BertForQuestionAnswering
# 查看 https://huggingface.co/docs/transformers/model_doc/bert
cls = getattr(importlib.import_module(gena_module_name), 'BertForSequenceClassification')
print(cls)
model = cls.from_pretrained('AIRI-Institute/gena-lm-bert-base-t2t', num_labels=2)

GENA-LM bigbird-base-t2t 模型使用HuggingFace BigBird实现。因此，可以使用Transformers库中的默认类:

from transformers import AutoTokenizer, BigBirdForSequenceClassification
tokenizer = AutoTokenizer.from_pretrained('AIRI-Institute/gena-lm-bigbird-base-t2t')
model = BigBirdForSequenceClassification.from_pretrained('AIRI-Institute/gena-lm-bigbird-base-t2t')

笔记本

• 使用GENA-LM和Huggingface Transformers进行序列分类

• 使用GENA LM生成的DNA嵌入进行聚类

• 探索在Enformer数据集上微调的GENA-LM模型用于基因表达

引用

@article {GENA_LM,
	author = {Veniamin Fishman and Yuri Kuratov and Maxim Petrov and Aleksei Shmelev and Denis Shepelin and Nikolay Chekanov and Olga Kardymon and Mikhail Burtsev},
	title = {GENA-LM: A Family of Open-Source Foundational Models for Long DNA Sequences},
	elocation-id = {2023.06.12.544594},
	year = {2023},
	doi = {10.1101/2023.06.12.544594},
	publisher = {Cold Spring Harbor Laboratory},
	URL = {https://www.biorxiv.org/content/early/2023/06/13/2023.06.12.544594},
	eprint = {https://www.biorxiv.org/content/early/2023/06/13/2023.06.12.544594.full.pdf},
	journal = {bioRxiv}
}

下游任务

用于模型评估的下游任务包括预测启动子和增强子活性、剪接位点、染色质谱和多聚腺苷酸化位点强度。 查看downstream_tasks文件夹以获取我们使用的代码和数据预处理脚本:

• [启动子预测](https://github.com/AIRI-Institute/GENA_LM/blob/main/./downstream_tasks/promoter_prediction/
• [剪接位点预测 (SpliceAI)](https://github.com/AIRI-Institute/GENA_LM/blob/main/./downstream_tasks/SpliceAI/
• [果蝇增强子预测 (DeepSTARR)](https://github.com/AIRI-Institute/GENA_LM/blob/main/./downstream_tasks/DeepSTARR/
• [染色质谱分析 (DeepSea)](https://github.com/AIRI-Institute/GENA_LM/blob/main/./downstream_tasks/DeepSea/
• 多聚腺苷酸化位点预测 (APARENT)

预训练数据

下载和预处理数据

要下载人类基因组，请运行以下脚本:

./download_data.sh human

对于预处理，执行以下脚本：

python src/gena_lm/genome_tools/create_corpus.py --input_file data/ncbi_dataset/data/GCA_009914755.4/GCA_009914755.4_T2T-CHM13v2.0_genomic.fna --output_dir data/processed/human/

安装

对于具有稀疏注意力的模型（gena-lm-bigbird-base-sparse、gena-lm-bigbird-base-sparse-t2t），需要FP16支持和DeepSpeed。

用于FP16的APEX

安装APEX https://github.com/NVIDIA/apex#quick-start

git clone https://github.com/NVIDIA/apex
cd apex
# 最新的提交可能会构建失败
git checkout 2386a912164b0c5cfcd8be7a2b890fbac5607c82
# 如果pip版本 >= 23.1（参考：https://pip.pypa.io/en/stable/news/#v23-1），支持多个相同键的`--config-settings`... 
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
# 否则
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option="--cpp_ext" --global-option="--cuda_ext" ./

用于稀疏操作的DeepSpeed

需要安装DeepSpeed以使用语言模型的稀疏注意力版本。DeepSpeed稀疏注意力仅支持计算能力 >= 7的GPU（V100、T4、A100），CUDA 10.1、10.2、11.0或11.1，并且仅在FP16模式下运行（截至DeepSpeed 0.6.0）。

可以使用pytorch.org提供的CUDA 10.2/11.0/11.1版本的PyTorch>=1.7.1,<=1.10.1轮子。 但是，使用CUDA 11.1 PyTorch轮子的稀疏操作需要在系统上安装CUDA 11.3/11.4。 稀疏操作也可以与PyTorch==1.12.1 CUDA 11.3轮子一起使用，但运行DeepSpeed稀疏操作测试需要修改它们，因为它们检查Torch CUDA版本 <=11.1。 用于Triton 1.1.1的DeepSpeed分支已经更新了测试。

Triton 1.0.0和1.1.1需要python<=3.9。

pip install triton==1.0.0
DS_BUILD_SPARSE_ATTN=1 pip install deepspeed==0.6.0 --global-option="build_ext" --global-option="-j8" --no-cache

并使用以下命令检查安装：

ds_report

Triton 1.1.1

Triton 1.1.1为A100上的稀疏操作带来了2倍的速度提升，但DeepSpeed（0.6.5）目前仅支持triton 1.0.0。 在需要这种速度提升的情况下，可以使用支持triton 1.1.1的DeepSpeed分支：

pip install triton==1.1.1
git clone https://github.com/yurakuratov/DeepSpeed.git
cd DeepSpeed
DS_BUILD_SPARSE_ATTN=1 pip install -e . --global-option="build_ext" --global-option="-j8" --no-cache

并使用以下命令运行稀疏操作测试：

cd tests/unit
pytest -v test_sparse_attention.py

使用lm-experiments-tools进行微调

我们使用lm-experiments-tools仓库中的Trainer和多GPU训练作为我们微调脚本的基础。但是，您也可以使用HF Transformers Trainer、PyTorch Lightning或Accelerate和PyTorch与自定义训练循环来代替。

根据https://github.com/yurakuratov/t5-experiments#install-only-lm_experiments_tools安装lm-experiments-tools：

git clone https://github.com/yurakuratov/t5-experiments
cd t5-experiments
pip install -e .