Diffusers documentation
RFD4ProteinaTransformer
RFD4ProteinaTransformer
RFD4-Proteínaのデノイザーです。サンプラーが積分する2つのモダリティ、すなわちバックボーンのCA座標とトークンごとのローカル潜在変数について、準備済みの分子テンソルから速度を予測します。入力には名前付きテンソルと、チェックポイントの設定によってその内容が選択される条件付けマッピングが含まれます。構造の解析はこのクラスの範囲外です。
from diffusers import RFD4ProteinaTransformer
transformer = RFD4ProteinaTransformer.from_pretrained(
"nvidia/RFD4-Proteina-Diffusers", subfolder="transformer"
)RFD4ProteinaTransformer
class diffusers.RFD4ProteinaTransformer
< source >( nlayers: int = 14token_dim: int = 768nheads: int = 12dim_cond: int = 256pair_repr_dim: int = 384latent_dim: int = 32token_type_emb_dim: int = 32cath_emb_dim: int = 32t_emb_dim: int = 256seq_sep_dim: int = 64xt_pair_dist_dim: int = 30xt_pair_dist_min: float = 0.1xt_pair_dist_max: float = 3.0feats_seq: tuple[str, ...] = ('xt_bb_ca', 'xt_local_latents', 'token_type', 'sequence_probabilities', 'cond_a37coors_nm', 'cond_a37coors_internal_nm', 'rasa_binned', 'secondary_structure', 'hbond_modernized', 'pistack', 'pication', 'saltbridge', 'hydrophobic', 'metal_coord', 'na_ss_state', 'hotspot', 'atom_charge', 'atomic_element', 'coord_group_seq', 'chirality', 'chain_centroid_xyz')feats_cond_seq: tuple[str, ...] = ('time_emb_bb_ca', 'time_emb_local_latents', 'cath_code', 'recycle_remaining')feats_pair_repr: tuple[str, ...] = ('xt_bb_ca_pair_dists', 'bond_mask', 'bond_order', 'is_aromatic', 'index_group_pair', 'molecule_idx_pair', 'molecule_entity_pair', 'token_bond_path_length', 'canonical_polymer_ordering', 'na_ss_pair')feats_pair_cond: tuple[str, ...] = ('time_emb_bb_ca', 'time_emb_local_latents')update_pair_repr_every_n: int = 2tri_mult_c: int = 256num_sequence_classes: int = 147num_token_types: int = 7recycle_conditioning_dim: int = 1 )
パラメータ
- nlayers (
int、デフォルトは14) — トランク層の数。 - token_dim (
int、デフォルトは768) — トークン表現の幅。 - nheads (
int, デフォルトは12) — 各レイヤーあたりのアテンションヘッド数。 - dim_cond (
int、デフォルトは256) — 条件ベクトル(time、recycle index、CATH label)の幅。 - pair_repr_dim (
int、デフォルトは384) — ペア表現の幅。 - latent_dim (
int、デフォルトは32) — トークンごとのローカル潜在変数の幅。 - token_type_emb_dim (
int、デフォルトは32) — トークンタイプ埋め込みの幅。 - cath_emb_dim (
int、デフォルトは32) — 各 CATH クラス/アーキテクチャ/トポロジー埋め込みの幅。 - t_emb_dim (
int、デフォルトは256) — 正弦波時間埋め込みの幅。 - seq_sep_dim (
int、デフォルトは64) — 結合パス長のビン数。 - xt_pair_dist_dim (
int、デフォルトは30) — バックボーンのペア距離のビン数。 - xt_pair_dist_min (
float、デフォルトは0.1)、xt_pair_dist_max (float、デフォルトは3.0) — ナノメートル単位のペア距離の範囲。 - feats_seq, feats_cond_seq, feats_pair_repr, feats_pair_cond (
list[str]) — トークン入力、条件入力、ペア入力、ペア条件入力に射影される特徴の名前。 - update_pair_repr_every_n (
int、デフォルトは2) — レイヤー間のトークンからペア表現を更新する頻度です。 - tri_mult_c (
int、デフォルトは256) — 該当する更新における三角形乗法の隠れ幅。 - num_token_types (
int、デフォルトは7) — トークンタイプ語彙のサイズ: protein、RNA、DNA、atomized、unknown、glycan、および非標準アミノ酸。 - recycle_conditioning_dim (
int、デフォルトは1) —feats_cond_seqにrecycle_remainingのone-hot特徴が含まれている場合の幅。 - num_sequence_classes (
int、デフォルトは147) — シーケンス読み出しが予測する残基語彙のサイズ。
ペアされたバックボーンとローカル・ラテンントのモダリティを備えたRFD4速度デノイザー。
モデルは独自の dtype で実行されます: 低精度推論のために dtype=torch.bfloat16 で読み込みます。forward に渡されるサンプリング状態は float32 のままです。特徴量埋め込みは float32 で計算され、入力投影時にキャストされます。
forward
< source >( bb_ca: Tensorlocal_latents: Tensorbb_ca_time: Tensorlocal_latents_time: Tensorpadding_mask: Tensortoken_type: typing.Optional[torch.Tensor] = Noneatom37_mask: typing.Optional[torch.Tensor] = Noneconditioning: dict[str, typing.Any] | None = Nonerecycle_depth: int = 0attention_kwargs: dict[str, typing.Any] | None = Nonereturn_dict: bool = True ) → RFD4ProteinaTransformerOutput or tuple
パラメータ
- bb_ca (
torch.Tensor) — ナノメートル単位のノイジーなバックボーン座標(batch, tokens, 3)。 - local_latents (
torch.Tensor) — ノイズ付きのローカル潜在変数(batch, tokens, latent_dim)。 - bb_ca_time (
torch.Tensor) —(batch,)[0, 1] 範囲の backbone time。 - local_latents_time (
torch.Tensor) —(batch,)[0, 1] 範囲のローカル・ラテンツの時間。 - padding_mask (
torch.Tensor) —(batch, tokens)。トークンがパディングではなく実データであることを示すマスク。 - token_type (
torch.Tensor, 任意) —(batch, tokens)のトークンタイプコード。 - atom37_mask (
torch.Tensor, 任意) —(batch, tokens, 37)。アトムスロットが使用されている箇所が true になります。 - conditioning (
dict[str, Any]、省略可) — プリプロセッサが出力するキーの下にある、残りのトポロジーおよび条件付けテンソルです。読み取られるキーは、チェックポイントのfeats_seq、feats_cond_seq、feats_pair_repr、feats_pair_condによって固定されています。設定されたフィーチャーが要求するキーが存在しない場合、エラーを発生させるのではなく、そのフィーチャーは条件付けなしのままになります。 - recycle_depth (
int、デフォルトは0) — トランクのパス数。0でも1回のパスが実行されます。 - attention_kwargs (
dict、省略可) — この呼び出しの LoRA スケールとして、scaleのみが読み取られます。 - return_dict (
bool、デフォルトはTrue) — 型付きの出力を返すか、None 以外のフィールドのタプルを返すかを指定します。
戻り値
RFD4ProteinaTransformerOutput または tuple
スケジューラー統合前の、Backbone/local-latent の速度および任意のシーケンス/挿入 logits。
サンプラーを進行させずに、準備済みの分子テンソルから速度を予測します。
RFD4ProteinaTransformerOutput
クラス diffusers.models.transformers.transformer_rfd4_proteina.RFD4ProteinaTransformerOutput
< source >( backbone: Tensorlocal_latents: typing.Optional[torch.Tensor] = Nonesequence_logits: typing.Optional[torch.Tensor] = Noneinsertion_logits: typing.Optional[torch.Tensor] = None )
パラメータ
- backbone (
torch.Tensor) — 形状が(batch, tokens, 3)の Backbone の速度。単位はナノメートル/単位時間です。 - local_latents (
torch.Tensor) — 形状が(batch, tokens, latent_dim)の潜在速度。 - sequence_logits (
torch.Tensor, 任意) — 形が(batch, tokens, 147)の残基の logits。 - insertion_logits (
torch.Tensor, 任意) — 形が(batch, tokens, tokens)の、挿入割り当ての logits。
スケジューラー統合前のDenoiserの予測。