RFD4-Proteína

RFD4-Proteína هو نموذج مطابقة تدفق من مختبر بيكر ونفيديا لتصميم الجزيئات الحيوية. يقوم بتوليد تسلسلات البروتينات والبنى الذرية الكاملة معًا، مشروطًا بالسياق الجزيئي للحفاظ على: هدف الارتباط، أو دافع الإنزيم، أو الرابطة، أو النقاط الساخنة، أو البنية الثانوية. يتم دمج محوّل مزدوج ومشفّر-فك تشفير بواسطة جدولين RFD4ProteinaScheduler، أحدهما للهيكل العظمي والآخر للمتغيرات المحلية.

يستهلك خط الأنابيب متجهات الميزات المُعدّة. وتحويل ملف CIF مُعلَّم إلى تلك المتجهات، وإعادة المتجهات المُولَّدة إلى ملف CIF، هي معالجة مسبقة خاصة بالمجال وتوجد في حزمة rfd4-proteina-prep.

مثال على الاستخدام

قم بتثبيت حزمة الإعداد عند مراجعة ثابتة. فهي تحدد متجهات الميزات التي يستقبلها النموذج، لذا فإن تثبيتًا غير مثبّت سيؤدي إلى تغيير الهياكل المولّدة عند تحديث الحزمة.

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")

نصائح

تحتوي مستودع النموذج على مدخلات مُعلَّمة وتصميم مُولَّد لكل مهمة تحت examples/: بروتين رابط، جسم مضاد، VHH، إنزيم، هياكل الدوافع، جزيء صغير رابط، غير مشروط، طي عكسي، وانتشار جزئي.
مرّر partial_t_bb_ca وpartial_t_local_latents مع أزمنة بداية ضمن [0, 1) لإعادة تصميم مدخل من بداية مُضجِبة، أو مرّر inverse_folding=True مع جعل كلا زمنَي البداية 0.0 للحفاظ على الهيكل الأساسي وتصميم التسلسل. يقوم كلاهما بتشفير المدخل، لذا يتطلبان إحداثيات coords_nm نظيفة ويرفضان الإحداثيات غير المحددة، والتي تكون عادةً بقايا غير محلولة.
عدد الخطوات هو طول إعداد خط الأنابيب recycle_depths، وهو 60 في الوصفة المنشورة (180 مرورًا بالجذع). يمكنك تغيير الجدول الزمني عن طريق تمرير قائمة recycle_depths مختلفة عند إنشاء خط الأنابيب.

لتوليد عدة تصاميم لكل مدخل، اضبط num_designs_per_prompt ومرر مولّدًا مختلفًا لكل تصميم. يحتوي المخرَج على batch * num_designs_per_prompt صفًا، والصف r هو التصميم r % num_designs_per_prompt من صف المدخل r // num_designs_per_prompt.

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 بقراءة خصائص الصف الذي يكتبه، لذا كررها بالطريقة نفسها، باستخدام repeat_interleave بدلاً من repeat:

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)

قم بتحميل torch.bfloat16 من أجل زيادة السرعة. تبقى حالة المُعيّن (sampler) وتضمينات الميزات (feature embeddings) في float32. قد يؤدي استخدام نوع بيانات (dtype) أو نواة (kernel) أو حجم دفعة (batch size) مختلف إلى تصميم مختلف لنفس البذرة (seed).
البنية المولّدة ليست تصميمًا مُتحقَّقًا منه. أعد طيّها وقِسها للمهمة المقصودة.

LoRA

النموذج المُطلَق هو نسخة LoRA مُكيَّفة بدقة بترتيب 32 (rank-32) من النموذج الأساسي المُدرَّب مسبقًا، وقد تم دمجها في الأوزان. يتوفر النموذج الأساسي والمُكيِّف (adapter) على مراجعة pretrain-base من مستودع النموذج، وتحمّل خط المعالجة المُكيِّف باستخدام load_lora_weights عبر RFD4ProteinaLoraLoaderMixin:

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

فئة diffusers.RFD4ProteinaPipeline

< source >

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

المعاملات

transformer (RFD4ProteinaTransformer) — هيكل أساسي مزدوج ومزيل ضجيج سرعة محلي الكامن.
vae (RFD4ProteinaAutoencoder) — مُشفِّر-المُفكِّك المصاحب لنقطة التحقق من نموذج إزالة الضوضاء.
bb_ca_scheduler (RFD4ProteinaScheduler) — يدمج إحداثيات الهيكل العظمي.
local_latents_scheduler (RFD4ProteinaScheduler) — يدمج الكامنات المحلية.
recycle_depths (list[int]، اختياري) — مرات إعادة تدوير الجذع لكل خطوة إزالة الضوضاء. طولها هو العدد الافتراضي لخطوات الاستدلال. الافتراضي هو الجدول المنشور المكوّن من 60 خطوة و180 مرة.

توليد الهياكل التسلسلات الجزيئية من الموترات الجزيئية المُعدّة.

يستهلك خط الأنابيب متجهات الميزات الناتجة عن تحضير مدخلات RFD4 (انظر حزمة rfd4-proteina-prep للحصول على ملفات CIF مُعلَّمة) ويعيد إحداثيات atom37 المُفكَّكة وهويات البقايا. الإحداثيات بالنانومتر.

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

المعاملات

padding_mask (torch.Tensor) — مصفوفة منطقية (batch, tokens)، تكون قيمتها صحيحة (true) حيث يكون الرمز (token) حقيقيًا وليس تعبئة (padding).
is_in_context_conditioning (torch.Tensor) — قيمة منطقية (batch, tokens)، تكون صحيحة (true) للرموز المقدمة كسياق وتُترك دون تغيير.
supports_local_latent (torch.Tensor) — قيمة منطقية (batch, tokens)، تكون صحيحة (true) للرموز التي تحمل كامنًا محليًا. تتشارك الأقنعة الثلاثة جميعها نفس الشكل والجهاز.
token_type (torch.Tensor، اختياري) — رموز أنواع الرموز (batch, tokens).
atom37_mask (torch.Tensor، اختياري) — (batch, tokens, 37)، تكون قيمتها صحيحة (true) حيث يُستخدم خانة ذرة.
coords_nm (torch.Tensor، اختياري) — إحداثيات (batch, tokens, 37, 3) نظيفة للذرات الـ37 بالنانومتر، مطلوبة في الانتشار الجزئي والطي العكسي.
conditioning_features (dict[str, Any]، اختياري) — المصفوفات المتبقية للهندسة البنيوية والشروط تحت المفاتيح التي يصدرها المعالج المسبق. يجب أن يحتوي على شرط التسلسل، condition_sequence_1_residue، وmolecule_iid. يتم تحديد أي من المتغيرات الأخرى تُقرأ من خلال إعداد الميزات الخاص بالنقطة الفاصلة. وهو عبارة عن خرائط واحدة لأن المعالج المسبق يصدر حوالي 30 مصفوفة اختيارية بأسماء ثابتة، وكان توقيعًا يحتوي على ذلك العدد من الحُجج سيربط خط الأنابيب بأسماء المعالج المسبق.
num_designs_per_prompt (int، القيمة الافتراضية 1) — عدد التصاميم التي سيتم توليدها لكل صف إدخال. تُكرَّر المدخلات على محور الدفعة (batch)، لذا يحتوي المخرج على batch * num_designs_per_prompt صفًا، ويكون الصف r هو التصميم r % num_designs_per_prompt من صف الإدخال r // num_designs_per_prompt.
generator (torch.Generator أو list[torch.Generator]، اختياري) — مولّد واحد مختلف لكل صف من صفوف المخرجات، من أجل تصاميم قابلة للتكرار. تتقدّم المولّدات وتُترك سلسلة PyTorch العشوائية العامة دون تغيير. احذفه للسحب من السلسلة العامة.
partial_t_bb_ca (float، اختياري) — وقت البدء في [0، 1) للعمود الفقري، للبدء من نسخة مُضجِمة من المدخل بدلاً من الضوضاء النقية. يُقدَّم مع partial_t_local_latents؛ ويتطلب coords_nm وatom37_mask.
partial_t_local_latents (float، اختياري) — وقت البدء في [0، 1) للـ latents المحلية؛ راجع partial_t_bb_ca.
inverse_folding (bool، القيمة الافتراضية False) — الحفاظ على الهيكل الأساسي المُدخل وتصميم التسلسل. يتطلب أن تكون قيمتا وقت البدء الجزئي صفرًا.
latents_bb_ca (torch.Tensor، اختياري) — حالة العمود الفقري الابتدائية المُولَّدة مسبقًا (batch * num_designs_per_prompt, tokens, 3)، لاستبدال الحالة التي كان من المفترض أن تُستخرج عشوائيًا في هذه الاستدعاء. يبقى الجدول الزمني دون تغيير. قدّمه معًا مع latents_local_latents.
latents_local_latents (torch.Tensor، اختياري) — محاولات محلية ابتدائية مُولّدة مسبقًا (batch * num_designs_per_prompt, tokens, latent_dim).
attention_kwargs (dict، اختياري) — يُمرَّر إلى المحوِّل؛ scale يضبط مقياس LoRA.
callback_on_step_end (Callable، اختياري) — يُستدعى على شكل callback_on_step_end(pipeline, step, timestep, callback_kwargs) بعد كل خطوة. يحتوي callback_kwargs على الموترات المسماة في callback_on_step_end_tensor_inputs؛ وقد يحل القاموس المُعاد محل "bb_ca" و"local_latents".
output_type (str)، القيمة الافتراضية "pt") — "pt" يعيد تينسورات torch بينما "np" يعيد مصفوفات numpy.
callback_on_step_end_tensor_inputs (list[str], defaults to ["bb_ca", "local_latents"]) — Tensors passed to the callback: "bb_ca", "local_latents", "last_clean_bb_ca", "last_clean_local_latents", "sequence_logits", "insertion_logits".
return_dict (bool، القيمة الافتراضية True) — إرجاع مخرجات من نوع محدد، أو مجموعة (tuple) من حقولها غير الفارغة (non-None).

يعيد

RFD4ProteinaPipelineOutput أو tuple

متجهات atom37 المُفكّكة (coords_nm بالنانومتر، residue_type، emit_mask، atom_mask)، والحالة الضجيجية النهائية، وآخر تقدير نظيف، وأخيراً قيم logits للتسلسل والإدراج.

توليد الهياكل والتسلسلات من الموترات الجزيئية المُعدّة.

أمثلة:

>>> 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 )

تحقّق من أقنّة الجزيئات، ومولّدات التصميم لكل تصميم، والانتشار الجزئي، والنداءات العكسية.

noise_backbone

< source >

( target_backbonebb_maskconditioninggeneratorsindex )

أضف ضوضاء إلى الهيكل المحفوظ حتى الوقت المجدول للخطوة index، من أجل الطي العكسي.

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 )

ارسم الحالة الابتدائية: ضوضاء نقية، أو الإدخال النظيف المُضاف إليه ضوضاء حتى start_times (الهيكل الأساسي، الكمونيات). يمكن إهمال clean_latents عندما يكون زمن البدء الكموني صفراً.

تتبع عمليات السحب ترتيب المُعيّن الأصلي في كل مولّد صف: ففي التوليد الكامل، تأتي ضوضاء الهيكل الأساسي وإزاحاتها العشوائية قبل ضوضاء الكمون، بينما في الانتشار الجزئي، تُسحب كلتا الضوضاء قبل الإزاحات.

RFD4ProteinaPipelineOutput

فئة 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 )

المعاملات

coords_nm (torch.Tensor) — إحداثيات الذرّات 37 المُفكّكة (B, N, 37, 3).
residue_type (torch.Tensor) — فهارس مفردات البقايا (B, N)؛ وقيمة -1 حيث لا تُصدر أي بقايا.
emit_mask (torch.Tensor) — قناع منطقي (B, N) للرموز التي تم تصميمها.
atom_mask (torch.Tensor) — قناع منطقي (B, N, 37) للذرات الصالحة.
bb_ca (torch.Tensor) — الحالة النهائية للعمود الفقري (B, N, 3) بالنانومتر.
local_latents (torch.Tensor) — الحالة النهائية المحلية الكامنة (B, N, D).
clean_bb_ca (torch.Tensor) — تنبؤ الهيكل الأساسي بالحالة النظيفة في خطوة إزالة الضوضاء الأخيرة.
clean_local_latents (torch.Tensor) — تنبؤ الحالة النظيفة المحلية في خطوة إزالة الضوضاء الأخيرة.
sequence_logits (torch.Tensor) — قيم اللوجيت النهائية لكل رمز (B, N, 147).
insertion_logits (torch.Tensor) — توقعات الإدخال النهائية (B, N, N) للأنماط غير المؤرَّخة.

نتائج الاستدلال الجزيئي. الإحداثيات بالنانومتر.

Update on GitHub

Diffusers

RFD4-Proteína

مثال على الاستخدام

نصائح

LoRA

RFD4ProteinaPipeline

فئة diffusers.RFD4ProteinaPipeline

__call__

check_inputs

noise_backbone

prepare_latents

RFD4ProteinaPipelineOutput

فئة diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput

call