Segment Nt Multi Species

🚀 SegmentNT-multi-species

SegmentNT-multi-speciesは、Nucleotide Transformer (NT) DNA基礎モデルを活用したセグメンテーションモデルです。このモデルは、シングルヌクレオチド解像度で配列内のいくつかのタイプのゲノム要素の位置を予測します。これは、SegmentNTモデルを、ヒトゲノムだけでなく、マウス、ニワトリ、ハエ、ゼブラフィッシュ、ワームの5つの選択された種のゲノムを含むデータセットでファインチューニングした結果です。

マルチ種ゲノムでのファインチューニングのために、私たちはSegmentNTのトレーニングに使用される注釈のサブセットのデータセットを作成しました。主な理由は、これらの種に関してはこの注釈のサブセットのみが利用可能だからです。したがって、注釈はEnsemblから利用可能な7つの主要な遺伝子要素、すなわちタンパク質コード遺伝子、5’UTR、3’UTR、イントロン、エクソン、スプライスアクセプターサイト、およびドナーサイトに関するものです。

開発者: InstaDeep

🚀 クイックスタート

モデルソース

• リポジトリ: Nucleotide Transformer
• 論文: Segmenting the genome at single-nucleotide resolution with DNA foundation models

使い方

次のリリースまでは、モデルを使用するために、transformersライブラリを以下のコマンドでソースからインストールする必要があります。

pip install --upgrade git+https://github.com/huggingface/transformers.git

ダミーのDNA配列からロジットと埋め込みを取得するための小さなコードスニペットを以下に示します。

⚠️ 重要提示

最大シーケンス長は、デフォルトでトレーニング長の30,000ヌクレオチド、または5001トークン（先頭のCLSトークンを含む）に設定されています。ただし、SegmentNTは最大50,000bpのシーケンスまで汎化することが示されています。30kbpから50kbpのシーケンスで推論する必要がある場合は、設定のrescaling_factor引数をnum_dna_tokens_inference / max_num_tokens_ntに変更してください。ここで、num_dna_tokens_inferenceは推論時のトークン数（例えば、40008塩基対のシーケンスの場合は6669）であり、max_num_tokens_ntはバックボーンのnucleotide-transformerがトレーニングされた最大トークン数、つまり2048です。

# モデルとトークナイザーをロード
from transformers import AutoTokenizer, AutoModel
import torch

tokenizer = AutoTokenizer.from_pretrained("InstaDeepAI/segment_nt_multi_species", trust_remote_code=True)
model = AutoModel.from_pretrained("InstaDeepAI/segment_nt_multi_species", trust_remote_code=True)

# 入力シーケンスをパディングする長さを選択します。デフォルトでは、
# モデルの最大長が選択されますが、埋め込みを取得するのにかかる時間が大幅に増加するため、
# 自由に減らすことができます。
# DNAトークンの数（先頭に付加されるCLSトークンを除く）は、
# ダウンサンプリングブロックの数の2の累乗、つまり4で割り切れる必要があります。
max_length = 12 + 1

assert (max_length - 1) % 4 == 0, (
    "DNAトークンの数（先頭に付加されるCLSトークンを除く）は、"
    "ダウンサンプリングブロックの数の2の累乗、つまり4で割り切れる必要があります。")

# ダミーのDNAシーケンスを作成し、トークン化します
sequences = ["ATTCCGATTCCGATTCCG", "ATTTCTCTCTCTCTCTGAGATCGATCGATCGAT"]
tokens = tokenizer.batch_encode_plus(sequences, return_tensors="pt", padding="max_length", max_length = max_length)["input_ids"]

# 推論
attention_mask = tokens != tokenizer.pad_token_id
outs = model(
    tokens,
    attention_mask=attention_mask,
    output_hidden_states=True
)

# ゲノム特徴に関するロジットを取得
logits = outs.logits.detach()
# 確率に変換
probabilities = torch.nn.functional.softmax(logits, dim=-1)
print(f"確率の形状: {probabilities.shape}")

# イントロンに関連する確率を取得
idx_intron = model.config.features.index("intron")
probabilities_intron = probabilities[:,:,idx_intron]
print(f"イントロンの確率の形状: {probabilities_intron.shape}")

📦 インストール

次のリリースまでは、モデルを使用するために、transformersライブラリを以下のコマンドでソースからインストールする必要があります。

pip install --upgrade git+https://github.com/huggingface/transformers.git

💻 使用例

基本的な使用法

# モデルとトークナイザーをロード
from transformers import AutoTokenizer, AutoModel
import torch

tokenizer = AutoTokenizer.from_pretrained("InstaDeepAI/segment_nt_multi_species", trust_remote_code=True)
model = AutoModel.from_pretrained("InstaDeepAI/segment_nt_multi_species", trust_remote_code=True)

# 入力シーケンスをパディングする長さを選択します。デフォルトでは、
# モデルの最大長が選択されますが、埋め込みを取得するのにかかる時間が大幅に増加するため、
# 自由に減らすことができます。
# DNAトークンの数（先頭に付加されるCLSトークンを除く）は、
# ダウンサンプリングブロックの数の2の累乗、つまり4で割り切れる必要があります。
max_length = 12 + 1

assert (max_length - 1) % 4 == 0, (
    "DNAトークンの数（先頭に付加されるCLSトークンを除く）は、"
    "ダウンサンプリングブロックの数の2の累乗、つまり4で割り切れる必要があります。")

# ダミーのDNAシーケンスを作成し、トークン化します
sequences = ["ATTCCGATTCCGATTCCG", "ATTTCTCTCTCTCTCTGAGATCGATCGATCGAT"]
tokens = tokenizer.batch_encode_plus(sequences, return_tensors="pt", padding="max_length", max_length = max_length)["input_ids"]

# 推論
attention_mask = tokens != tokenizer.pad_token_id
outs = model(
    tokens,
    attention_mask=attention_mask,
    output_hidden_states=True
)

# ゲノム特徴に関するロジットを取得
logits = outs.logits.detach()
# 確率に変換
probabilities = torch.nn.functional.softmax(logits, dim=-1)
print(f"確率の形状: {probabilities.shape}")

# イントロンに関連する確率を取得
idx_intron = model.config.features.index("intron")
probabilities_intron = probabilities[:,:,idx_intron]
print(f"イントロンの確率の形状: {probabilities_intron.shape}")

📚 ドキュメント

トレーニングデータ

segment-nt-multi-speciesモデルは、ヒト、マウス、ニワトリ、ハエ、ゼブラフィッシュ、およびワームのゲノムでファインチューニングされました。各種について、染色体のサブセットはトレーニングモニタリングのための検証用と、最終評価のためのテスト用として保持されています。

トレーニング手順

前処理

DNA配列は、Nucleotide Transformer Tokenizerを使用してトークン化されます。このトークナイザーは、関連するリポジトリのTokenizationセクションで説明されているように、配列を6マートークンとしてトークン化します。このトークナイザーの語彙サイズは4105です。したがって、モデルの入力は次の形式になります。

<CLS> <ACGTGT> <ACGTGC> <ACGGAC> <GACTAG> <TCAGCA>

トレーニング

モデルは、8つのGPUを備えたDGXH100ノードで、合計80億トークンに対して3日間ファインチューニングされました。

アーキテクチャ

モデルは、nucleotide-transformer-v2-500m-multi-speciesエンコーダで構成されています。ここから、言語モデルヘッドを削除し、2つのダウンサンプリング畳み込みブロックと2つのアップサンプリング畳み込みブロックからなる1次元U-Netセグメンテーションヘッド[4]に置き換えました。これらの各ブロックは、それぞれ1,024と2,048のカーネルを持つ2つの畳み込み層で構成されています。この追加のセグメンテーションヘッドは5300万のパラメータを占め、総パラメータ数を5億6200万に増やします。

BibTeXエントリと引用情報

@article{de2024segmentnt,
  title={SegmentNT: annotating the genome at single-nucleotide resolution with DNA foundation models},
  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},
  journal={bioRxiv},
  pages={2024--03},
  year={2024},
  publisher={Cold Spring Harbor Laboratory}
}

📄 ライセンス

このプロジェクトは、CC BY-NC-SA 4.0ライセンスの下で公開されています。