RFD4-Proteína

RFD4-Proteína, biyomoleküler tasarım için Baker Lab ve NVIDIA tarafından geliştirilen bir akış eşleştirme modelidir. Bağlanma hedefi, enzim motifi, ligand, sıcak noktalar veya ikincil yapıyı korumak üzere moleküler bağlama koşullu olarak protein dizilerini ve tüm atom yapılarını birlikte üretir. Bir omurga ve bir yerel gizli değişken için olmak üzere iki RFD4ProteinaScheduler ile entegre edilmiş eşleşmiş bir transformer ve oto-encoder içerir.

Pipeline, hazırlanmış özellik tensörlerini tüketir. Annotasyonlu bir CIF dosyasını bu tensörlere dönüştürmek ve üretilen tensörleri tekrar bir CIF dosyasına çevirmek, rfd4-proteina-prep paketinde yer alan alanına özgü ön işlemedir.

Kullanım örneği

Hazırlık paketini sabit bir revizyonda yükleyin. Modelin gördüğü özellik tensörlerini belirlediğinden, sabitlenmemiş bir kurulum, paket güncellendiğinde üretilen yapıları değiştirir.

pip install diffusers git+https://huggingface.co/nvidia/RFD4-Proteina-prep@c9959efdebd75ac7fe9129fbbf0e291a6e0487d5

import torch
from huggingface_hub import hf_hub_download
from diffusers import RFD4ProteinaPipeline
from rfd4_proteina_prep import prepare_cif, write_cif

pipe = RFD4ProteinaPipeline.from_pretrained("nvidia/RFD4-Proteina-Diffusers", dtype=torch.bfloat16).to("cuda")

# An annotated CIF specifies what to hold fixed and what to design: here a PD-L1 target and a 50-120 residue binder.
input_cif = hf_hub_download("nvidia/RFD4-Proteina-Diffusers", "examples/protein_binder/input.cif")
features, atoms = prepare_cif(input_cif, seed=42, device=pipe.device)

# The masks and molecule tensors the pipeline always needs are arguments; the annotation tensors go in one mapping.
required = ["padding_mask", "is_in_context_conditioning", "supports_local_latent", "token_type", "atom37_mask"]
conditioning = {name: value for name, value in features.items() if name not in required}

output = pipe(
    padding_mask=features["padding_mask"],
    is_in_context_conditioning=features["is_in_context_conditioning"],
    supports_local_latent=features["supports_local_latent"],
    token_type=features["token_type"],
    atom37_mask=features["atom37_mask"],
    conditioning_features=conditioning,
    generator=torch.Generator("cuda").manual_seed(42),
)
write_cif(output, features, atoms, "design.cif")

İpuçları

Model deposu, examples/ altında her görev için etiketlenmiş bir girdi ve üretilmiş bir tasarım içerir: protein bağlayıcı, antikor, VHH, enzim, motif iskeletleme, küçük molekül bağlayıcı, koşulsuz, ters katlanma ve kısmi difüzyon.
Başlangıç zamanları [0, 1) aralığında olacak şekilde partial_t_bb_ca ve partial_t_local_latents parametrelerini geçirerek gürültülü bir başlangıçtan girdiyi yeniden tasarlayabilir veya her iki başlangıç zamanını da 0.0 olacak şekilde inverse_folding=True kullanarak iskeleti koruyup diziyi tasarlayabilirsiniz. Her iki yöntem de girdiyi kodladığı için temiz coords_nm değerleri gerektirir ve sonlu olmayan koordinatları reddeder; bunlar genellikle çözülmemiş kalıntıları temsil eder.
Adım sayısı, recycle_depths pipeline yapılandırmasının uzunluğudur; yayınlanan tarif için 60'tır (180 trunk geçişi). Pipeline oluşturulurken farklı bir recycle_depths listesi geçirerek zamanlamayı değiştirebilirsiniz.

Her giriş için birkaç tasarım oluşturmak için num_designs_per_prompt değerini ayarlayın ve her tasarım için ayrı bir üreteç geçin. Çıktı batch * num_designs_per_prompt satır içerir ve r. satır, r // num_designs_per_prompt. giriş satırının r % num_designs_per_prompt. tasarımını temsil eder.

generators = [torch.Generator("cuda").manual_seed(seed) for seed in range(8)]
output = pipe(
    padding_mask=features["padding_mask"],
    is_in_context_conditioning=features["is_in_context_conditioning"],
    supports_local_latent=features["supports_local_latent"],
    token_type=features["token_type"],
    atom37_mask=features["atom37_mask"],
    conditioning_features=conditioning,
    num_designs_per_prompt=8,
    generator=generators,
)

write_cif, yazdığı satırın özelliklerini okuduğu için, bunları repeat yerine repeat_interleave ile aynı şekilde tekrarlayın:

repeated = {k: v.repeat_interleave(8, dim=0) if torch.is_tensor(v) else v for k, v in features.items()}
for row in range(8):
    write_cif(output, repeated, atoms, f"design_{row}.cif", row=row)

Hız için torch.bfloat16 olarak yükleyin. Örnekleme durumu ve özellik gömme değerleri float32'de kalır. Farklı bir veri tipi, çekirdek veya parti boyutu, aynı tohum için farklı bir tasarıma yol açabilir.
Oluşturulan bir yapı, doğrulanmış bir tasarım değildir. Hedeflenen görev için yeniden katlayın ve puanlayın.

LoRA

Yayınlanan model, önceden eğitilmiş tabanın rank-32 LoRA ince ayarı olup ağırlıklara birleştirilmiştir. Taban ve adaptör, model deposunun pretrain-base revizyonundadır ve pipeline, adaptörü RFD4ProteinaLoraLoaderMixin üzerinden load_lora_weights ile yükler:

pipe = RFD4ProteinaPipeline.from_pretrained(
    "nvidia/RFD4-Proteina-Diffusers", revision="pretrain-base", dtype=torch.bfloat16
).to("cuda")
pipe.load_lora_weights(
    "nvidia/RFD4-Proteina-Diffusers",
    revision="pretrain-base",
    subfolder="lora",
    weight_name="pytorch_lora_weights.safetensors",
)

RFD4ProteinaPipeline

sınıf diffusers.RFD4ProteinaPipeline

< source >

( transformer: RFD4ProteinaTransformervae: RFD4ProteinaAutoencoderbb_ca_scheduler: RFD4ProteinaSchedulerlocal_latents_scheduler: RFD4ProteinaSchedulerrecycle_depths: list[int] | None = None )

Parametreler

transformer (RFD4ProteinaTransformer) — Eşlenik omurga ve yerel-gizli hız gürültü temizleyici.
vae (RFD4ProteinaAutoencoder) — Gürültü azaltıcı kontrol noktasıyla eşleştirilen otoenkoder.
bb_ca_scheduler (RFD4ProteinaScheduler) — Omurga koordinatlarını entegre eder.
local_latents_scheduler (RFD4ProteinaScheduler) — Yerel latentleri entegre eder.
recycle_depths (list[int], isteğe bağlı) — Her gürültü azaltma adımı için gövde geri dönüşüm geçişleri. Uzunluğu, varsayılan çıkarım adım sayısıdır. Varsayılan olarak, 60 adımlık ve 180 geçişli yayınlanan zaman çizelgesidir.

Hazırlanmış moleküler tensörlerden moleküler yapılar ve diziler üretin.

Pipeline, RFD4 girdi hazırlığı tarafından üretilen özellik tensörlerini tüketir (işaretli CIF dosyaları için rfd4-proteina-prep paketine bakın) ve çözümlenmiş atom37 koordinatlarını ile kalıntı kimliklerini döndürür. Koordinatlar nanometre cinsindendir.

call

< source >

( padding_mask: Tensoris_in_context_conditioning: Tensorsupports_local_latent: Tensortoken_type: typing.Optional[torch.Tensor] = Noneatom37_mask: typing.Optional[torch.Tensor] = Nonecoords_nm: typing.Optional[torch.Tensor] = Noneconditioning_features: dict[str, typing.Any] | None = Nonenum_designs_per_prompt: int = 1generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonepartial_t_bb_ca: float | None = Nonepartial_t_local_latents: float | None = Noneinverse_folding: bool = Falselatents_bb_ca: typing.Optional[torch.Tensor] = Nonelatents_local_latents: typing.Optional[torch.Tensor] = Noneattention_kwargs: dict[str, typing.Any] | None = Nonecallback_on_step_end: collections.abc.Callable[[diffusers.pipelines.pipeline_utils.DiffusionPipeline, int, torch.Tensor, dict], dict | None] | None = Nonecallback_on_step_end_tensor_inputs: list = ['bb_ca', 'local_latents']output_type: str = 'pt'return_dict: bool = True ) → RFD4ProteinaPipelineOutput or tuple

Parametreler

padding_mask (torch.Tensor) — Gerçek bir token olan yerlerde true, dolgu (padding) olan yerlerde false olan Boole (batch, tokens).
is_in_context_conditioning (torch.Tensor) — Bağlam olarak verilen ve değiştirilmeyen tokenlar için doğru olan Boolean (batch, tokens).
supports_local_latent (torch.Tensor) — Boolean (batch, tokens), yerel latent taşıyan tokenlar için true. Üç maskenin de aynı şekli ve cihazı vardır.
token_type (torch.Tensor, isteğe bağlı) — (batch, tokens) token-türü kodları.
atom37_mask (torch.Tensor, isteğe bağlı) — (batch, tokens, 37), bir atom yuvasının kullanıldığı yerlerde true.
coords_nm (torch.Tensor, isteğe bağlı) — Kısmi difüzyon ve ters katlanma için gereken, nanometre cinsinden temiz (batch, tokens, 37, 3) atom37 koordinatları.
conditioning_features (dict[str, Any], isteğe bağlı) — Önişleyicinin (preprocessor) yayınladığı anahtarlar altındaki kalan topoloji ve koşullama tensörleri. Dizi koşulunu, condition_sequence_1_residue ve molecule_iid değerlerini içermelidir. Geri kalanların hangilerinin okunacağı, kontrol noktasının (checkpoint) özellik yapılandırması tarafından belirlenir. Önişleyici, sabit adlar altında yaklaşık 30 isteğe bağlı tensör yayınladığı için ve bu kadar çok argümana sahip bir imza, hattı (pipeline) önişleyicinin adlandırma şemasına bağlayacağından, bu veriler tek bir eşleme (mapping) olarak tutulur.
num_designs_per_prompt (int, varsayılan 1) — Her giriş satırı için üretilecek tasarım sayısı. Girdiler, batch ekseni boyunca tekrarlandığından çıktı batch * num_designs_per_prompt satır içerir ve r. satır, r // num_designs_per_prompt. giriş satırının r % num_designs_per_prompt. tasarımına karşılık gelir.
generator (torch.Generator veya list[torch.Generator], isteğe bağlı) — Tekrarlanabilir tasarımlar için her çıktı satırı için ayrı bir üreteç. Üreteçler ilerler ve küresel PyTorch rastgele akışı dokunulmadan bırakılır. Küresel akımdan çekim yapmak için bunu atlayın.
partial_t_bb_ca (float, isteğe bağlı) — Girişin saf gürültü yerine gürültülü bir versiyonundan başlamak için omurganın [0, 1) aralığındaki başlangıç zamanı. partial_t_local_latents ile birlikte verilmelidir; coords_nm ve atom37_mask gerektirir.
partial_t_local_latents (float, isteğe bağlı) — Yerel latentler için [0, 1) aralığında başlangıç zamanı; bkz. partial_t_bb_ca.
inverse_folding (bool, varsayılan False) — Girdi iskeletini koruyup diziyi tasarlar. Kısmi başlangıç sürelerinin her ikisi için de sıfır gerektirir.
latents_bb_ca (torch.Tensor, isteğe bağlı) — Bu çağrının aksi halde örnekleyeceği durumu değiştirmek için önceden oluşturulmuş başlangıç omurga durumu (batch * num_designs_per_prompt, tokens, 3). Zamanlama değiştirilmez. latents_local_latents ile birlikte verin.
latents_local_latents (torch.Tensor, isteğe bağlı) — Önceden oluşturulmuş başlangıç yerel latensleri (batch * num_designs_per_prompt, tokens, latent_dim).
attention_kwargs (dict, isteğe bağlı) — Transformer'a iletilir; scale LoRA ölçeğini ayarlar.
callback_on_step_end (Callable, isteğe bağlı) — Her adımın ardından callback_on_step_end(pipeline, step, timestep, callback_kwargs) olarak çağrılır. callback_kwargs, callback_on_step_end_tensor_inputs içinde adları belirtilen tensörleri içerir; döndürülen bir sözlük, "bb_ca" ve "local_latents" öğelerini değiştirebilir.
output_type (str, varsayılan "pt") — "pt" torch tensörlerini, "np" ise numpy dizilerini döndürür.
callback_on_step_end_tensor_inputs (list[str], varsayılan ["bb_ca", "local_latents"]) — Geri çağrıya (callback) iletilen tensörler: "bb_ca", "local_latents", "last_clean_bb_ca", "last_clean_local_latents", "sequence_logits", "insertion_logits".
return_dict (bool, varsayılan True) — Tipi belirlenmiş bir çıktı veya None olmayan alanlarının bir demetini döndürür.

Dönüş değeri

RFD4ProteinaPipelineOutput veya tuple

Çözümlenmiş atom37 tensörleri (coords_nm nanometre cinsinden, residue_type, emit_mask, atom_mask), son gürültülü durum, son temiz tahmin ve son dizi ile ekleme logitleri.

Hazırlanmış moleküler tensörlerden yapılar ve diziler oluşturun.

Örnekler:

>>> import torch
>>> from diffusers import RFD4ProteinaPipeline
>>> from rfd4_proteina_prep import prepare_cif, write_cif

>>> pipe = RFD4ProteinaPipeline.from_pretrained("nvidia/RFD4-Proteina-Diffusers", dtype=torch.bfloat16)
>>> pipe.to("cuda")

>>> features, atoms = prepare_cif("input.cif", seed=42, device=pipe.device)
>>> required = ["padding_mask", "is_in_context_conditioning", "supports_local_latent", "token_type"]
>>> required += ["atom37_mask"]
>>> output = pipe(
...     padding_mask=features["padding_mask"],
...     is_in_context_conditioning=features["is_in_context_conditioning"],
...     supports_local_latent=features["supports_local_latent"],
...     token_type=features["token_type"],
...     atom37_mask=features["atom37_mask"],
...     conditioning_features={name: value for name, value in features.items() if name not in required},
...     generator=torch.Generator("cuda").manual_seed(42),
... )
>>> write_cif(output, features, atoms, "design.cif")

check_inputs

< source >

( padding_maskis_in_context_conditioningsupports_local_latentatom37_maskcoords_nmconditioning_featuresnum_designs_per_promptgeneratorpartial_tinverse_foldinglatentscallback_on_step_end_tensor_inputsoutput_type )

Moleküler maskeleri, tasarım başına üreteçleri, kısmi difüzyonu ve geri çağrıları doğrulayın.

noise_backbone

< source >

( target_backbonebb_maskconditioninggeneratorsindex )

Ters katlanma için korunan omurgayı, index adımının planlanan zamanına kadar gürültüleyin.

prepare_latents

< source >

( bb_mask: Tensorlatent_mask: Tensorconditioning: dictgenerators: list[torch.Generator] | Noneclean_bb: typing.Optional[torch.Tensor] = Noneclean_latents: typing.Optional[torch.Tensor] = Nonestart_times: tuple[float, float] | None = None )

Başlangıç durumunu çizin: saf gürültü veya start_times değerine kadar gürültülenmiş temiz giriş (backbone, latents). Latent başlangıç zamanı sıfır olduğunda clean_latents parametresi atlanabilir.

Çekimler, her satır üreteci üzerinde yerel örnekleyicinin sırasını izler: tam bir üretimde omurga gürültüsü ve onun stokastik ofsetleri, gizli gürültüsünden önce gelir; kısmi difüzyonda ise her iki gürültü de ofsetlerden önce çekilir.

RFD4ProteinaPipelineOutput

sınıf diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput

< source >

( coords_nm: Tensorresidue_type: Tensoremit_mask: Tensoratom_mask: Tensorbb_ca: Tensorlocal_latents: Tensorclean_bb_ca: Tensorclean_local_latents: Tensorsequence_logits: Tensorinsertion_logits: Tensor )

Parametreler

coords_nm (torch.Tensor) — Çözümlenmiş atom37 koordinatları (B, N, 37, 3).
residue_type (torch.Tensor) — Kalanlık sözlük dizinleri (B, N); bir kalanlık üretilmediği yerlerde -1.
emit_mask (torch.Tensor) — Tasarlanan jetonların Boole maskesi (B, N).
atom_mask (torch.Tensor) — Geçerli atomların Boole maskesi (B, N, 37).
bb_ca (torch.Tensor) — Nanometre cinsinden son backbone durumu (B, N, 3).
local_latents (torch.Tensor) — Son yerel latent durum (B, N, D).
clean_bb_ca (torch.Tensor) — Son gürültü azaltma adımında temiz durumun omurga tahmini.
clean_local_latents (torch.Tensor) — Son gürültü azaltma adımında temiz durumun yerel-latent tahmini.
sequence_logits (torch.Tensor) — Son token başına rezidü logitleri (B, N, 147).
insertion_logits (torch.Tensor) — Dizinlenmemiş motifler için nihai ekleme tahminleri (B, N, N).

Moleküler çıkarım sonuçları. Koordinatlar nanometre cinsindendir.

Update on GitHub

Diffusers

RFD4-Proteína

Kullanım örneği

İpuçları

LoRA

RFD4ProteinaPipeline

sınıf diffusers.RFD4ProteinaPipeline

__call__

check_inputs

noise_backbone

prepare_latents

RFD4ProteinaPipelineOutput

sınıf diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput

call