Evo: 从分子到基因组尺度的DNA基础建模

Evo

Evo是一个能够进行长上下文建模和设计的生物学基础模型。 Evo使用StripedHyena架构，能够以单核苷酸、字节级分辨率对序列进行建模，计算和内存相对于上下文长度近乎线性扩展。 Evo拥有70亿参数，在OpenGenome上训练，这是一个含有约3000亿个标记的原核生物全基因组数据集。

我们在论文"用Evo从分子到基因组尺度进行序列建模和设计"和配套博客文章中描述了Evo。

我们提供以下模型检查点：

检查点名称	描述
`evo-1-8k-base`	使用8,192上下文预训练的模型。我们将此模型用作分子尺度微调任务的基础模型。
`evo-1-131k-base`	使用`evo-1-8k-base`作为基础模型，用131,072上下文预训练的模型。我们使用此模型在基因组尺度推理和生成序列。

新闻

我们发现并修复了一个与某些投影错误排列相关的问题，该问题影响生成质量。要在HuggingFace上使用新的模型修订版，请按如下方式加载：

config = AutoConfig.from_pretrained(model_name, trust_remote_code=True, revision="1.1_fix")
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    config=config,
    trust_remote_code=True,
    revision="1.1_fix"
)

设置

要求

Evo基于StripedHyena。

Evo使用FlashAttention-2，可能不适用于所有GPU架构。请查阅FlashAttention GitHub仓库获取当前支持的GPU列表。

确保在您的系统上安装正确的PyTorch版本。

安装

您可以使用pip安装Evo

pip install evo-model

或直接从GitHub源代码安装

git clone https://github.com/evo-design/evo.git
cd evo/
pip install .

我们建议您先安装PyTorch库，然后再安装所有其他依赖项（由于flash-attn库的依赖问题；参见例如这个问题）。

我们的一个示例脚本，演示如何从用Evo生成序列到折叠蛋白质（scripts/generation_to_folding.py），还需要安装prodigal。我们为此创建了一个environment.yml文件：

conda env create -f environment.yml
conda activate evo-design

使用

以下是如何下载Evo并通过Python API在本地使用它的示例。

from evo import Evo
import torch

device = 'cuda:0'

evo_model = Evo('evo-1-131k-base')
model, tokenizer = evo_model.model, evo_model.tokenizer
model.to(device)
model.eval()

sequence = 'ACGT'
input_ids = torch.tensor(
    tokenizer.tokenize(sequence),
    dtype=torch.int,
).to(device).unsqueeze(0)

with torch.no_grad():
    logits, _ = model(input_ids) # (batch, length, vocab)

print('Logits: ', logits)
print('Shape (batch, length, vocab): ', logits.shape)

批量推理的示例可以在scripts/example_inference.py中找到。

我们提供了一个示例脚本，演示如何提示模型并根据提示采样一组序列。

python -m scripts.generate \
    --model-name 'evo-1-131k-base' \
    --prompt ACGT \
    --n-samples 10 \
    --n-tokens 100 \
    --temperature 1. \
    --top-k 4 \
    --device cuda:0

我们还提供了一个示例脚本，用于使用模型对一组序列的对数似然进行评分。

python -m scripts.score \
    --input-fasta examples/example_seqs.fasta \
    --output-tsv scores.tsv \
    --model-name 'evo-1-131k-base' \
    --device cuda:0

HuggingFace

Evo已集成到HuggingFace。

from transformers import AutoConfig, AutoModelForCausalLM

model_name = 'togethercomputer/evo-1-8k-base'

model_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True, revision="1.1_fix")
model_config.use_cache = True

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    config=model_config,
    trust_remote_code=True,
    revision="1.1_fix"
)

Together API

Evo可通过Together AI的网页界面使用，您可以在类似聊天的界面中生成DNA序列。

对于更详细或批量工作流程，您可以使用以下简单示例调用Together API。

import openai
import os

# 在此填写您的API信息。
client = openai.OpenAI(
  api_key=TOGETHER_API_KEY,
  base_url='https://api.together.xyz',
)

chat_completion = client.chat.completions.create(
  messages=[
    {
      "role": "system",
      "content": ""
    },
    {
      "role": "user",
      "content": "ACGT", # 用序列提示模型。
    }
  ],
  model="togethercomputer/evo-1-131k-base",
  max_tokens=128, # 采样一些新标记。
  logprobs=True
)
print(
    chat_completion.choices[0].logprobs.token_logprobs,
    chat_completion.choices[0].message.content
)

引用

引用Evo时，请引用以下预印本。

@article {nguyen2024sequence,
    author = {Eric Nguyen and Michael Poli and Matthew G Durrant and Armin W Thomas and Brian Kang and Jeremy Sullivan and Madelena Y Ng and Ashley Lewis and Aman Patel and Aaron Lou and Stefano Ermon and Stephen A Baccus and Tina Hernandez-Boussard and Christopher Ré and Patrick D Hsu and Brian L Hie},
    title = {Sequence modeling and design from molecular to genome scale with Evo},
    year = {2024},
    doi = {10.1101/2024.02.27.582234},
    publisher = {Cold Spring Harbor Laboratory},
    URL = {https://www.biorxiv.org/content/early/2024/02/27/2024.02.27.582234},
    journal = {bioRxiv}
}