Rna Torsionbert

🚀 RNA-TorsionBERT

RNA-TorsionBERT 是一个基于BERT的语言模型，参数大小为86.9 MB，可根据RNA序列预测其扭转角和伪扭转角。该模型在约4200个RNA结构上进行了预训练，为RNA相关研究提供了有力支持。

🚀 快速开始

使用以下代码片段，即可开始使用 RNA-TorsionBERT 生成文本：

from transformers import AutoModel, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("sayby/rna_torsionbert", trust_remote_code=True)
model = AutoModel.from_pretrained("sayby/rna_torsionbert", trust_remote_code=True)

sequence = "ACG CGG GGT GTT"
params_tokenizer = {
    "return_tensors": "pt",
    "padding": "max_length",
    "max_length": 512,
    "truncation": True,
}
inputs = tokenizer(sequence, **params_tokenizer)
output = model(inputs)["logits"]

⚠️ 重要提示

该模型是从DNABERT - 3模型微调而来，因此分词器与DNABERT使用的分词器相同。输入序列中的核苷酸 U 应替换为 T。

输出是每个角度的正弦和余弦值。角度顺序如下：alpha、beta、gamma、delta、epsilon、zeta、chi、eta、theta、eta'、theta'、v0、v1、v2、v3、v4。

若要将预测结果转换为角度，可使用以下代码片段：

import transformers
from transformers import AutoModel, AutoTokenizer
import numpy as np
import pandas as pd
from typing import Optional, Dict
import os

os.environ["TOKENIZERS_PARALLELISM"] = "false"

transformers.logging.set_verbosity_error()


BACKBONE = [
    "alpha",
    "beta",
    "gamma",
    "delta",
    "epsilon",
    "zeta",
    "chi",
    "eta",
    "theta",
    "eta'",
    "theta'",
    "v0",
    "v1",
    "v2",
    "v3",
    "v4",
]


class RNATorsionBERTHelper:
    def __init__(self):
        self.model_name = "sayby/rna_torsionbert"
        self.tokenizer = AutoTokenizer.from_pretrained(
            self.model_name, trust_remote_code=True
        )
        self.params_tokenizer = {
            "return_tensors": "pt",
            "padding": "max_length",
            "max_length": 512,
            "truncation": True,
        }
        self.model = AutoModel.from_pretrained(self.model_name, trust_remote_code=True)

    def predict(self, sequence: str):
        sequence_tok = self.convert_raw_sequence_to_k_mers(sequence)
        inputs = self.tokenizer(sequence_tok, **self.params_tokenizer)
        outputs = self.model(inputs)["logits"]
        outputs = self.convert_sin_cos_to_angles(
            outputs.cpu().detach().numpy(), inputs["input_ids"]
        )
        output_angles = self.convert_logits_to_dict(
            outputs[0, :], inputs["input_ids"][0, :].cpu().detach().numpy()
        )
        output_angles.index = list(sequence)[:-2]  # Because of the 3-mer representation
        return output_angles

    def convert_raw_sequence_to_k_mers(self, sequence: str, k_mers: int = 3):
        """
        Convert a raw RNA sequence into sequence readable for the tokenizer.
        It converts the sequence into k-mers, and replace U by T
        :return: input readable by the tokenizer
        """
        sequence = sequence.upper().replace("U", "T")
        k_mers_sequence = [
            sequence[i : i + k_mers]
            for i in range(len(sequence))
            if len(sequence[i : i + k_mers]) == k_mers
        ]
        return " ".join(k_mers_sequence)

    def convert_sin_cos_to_angles(
        self, output: np.ndarray, input_ids: Optional[np.ndarray] = None
    ):
        """
        Convert the raw predictions of the RNA-TorsionBERT into angles.
        It converts the cos and sinus into angles using:
            alpha = arctan(sin(alpha)/cos(alpha))
        :param output: Dictionary with the predictions of the RNA-TorsionBERT per angle
        :param input_ids: the input_ids of the RNA-TorsionBERT. It allows to only select the of the sequence,
            and not the special tokens.
        :return: a np.ndarray with the angles for the sequence
        """
        if input_ids is not None:
            output[
                (input_ids == 0)
                | (input_ids == 2)
                | (input_ids == 3)
                | (input_ids == 4)
            ] = np.nan
        pair_indexes, impair_indexes = np.arange(0, output.shape[-1], 2), np.arange(
            1, output.shape[-1], 2
        )
        sin, cos = output[:, :, impair_indexes], output[:, :, pair_indexes]
        tan = np.arctan2(sin, cos)
        angles = np.degrees(tan)
        return angles

    def convert_logits_to_dict(self, output: np.ndarray, input_ids: np.ndarray) -> Dict:
        """
        Convert the raw predictions into dictionary format.
        It removes the special tokens and only keeps the predictions for the sequence.
        :param output: predictions from the models in angles
        :param input_ids: input ids from the tokenizer
        :return: a dictionary with the predictions for each angle
        """
        index_start, index_end = (
            np.where(input_ids == 2)[0][0],
            np.where(input_ids == 3)[0][0],
        )
        output_non_pad = output[index_start + 1 : index_end, :]
        output_angles = {
            angle: output_non_pad[:, angle_index]
            for angle_index, angle in enumerate(BACKBONE)
        }
        out = pd.DataFrame(output_angles)
        return out


if __name__ == "__main__":
    sequence = "AGGGCUUUAGUCUUUGGAG"
    rna_torsionbert_helper = RNATorsionBERTHelper()
    output_angles = rna_torsionbert_helper.predict(sequence)
    print(output_angles)

✨ 主要特性

• 扭转角和伪扭转角预测：能够精准预测RNA的扭转角和伪扭转角。
• 长序列处理能力：可处理长度达512个核苷酸的序列。

📚 详细文档

模型描述

RNA-TorsionBERT 是一个参数为86.9 MB的基于BERT的语言模型，可根据RNA序列预测其扭转角和伪扭转角。它是在约4200个RNA结构上进行预训练的DNABERT模型。

在测试集（由RNA - Puzzles和CASP - RNA组成）上，与之前的先进模型（如 SPOT - RNA - 1D）或现有方法推断的角度相比，RNA-TorsionBERT 在 MCQ 指标上有显著提升。

📄 许可证

Other