rc-foundry 0.1.1__py3-none-any.whl

rfd3/model/layers/pairformer_layers.py

@@ -0,0 +1,128 @@
+import torch
+from rfd3.model.layers.layer_utils import (
+    MultiDimLinear,
+    RMSNorm,
+    Transition,
+    linearNoBias,
+)
+from torch import nn
+
+from foundry.training.checkpoint import activation_checkpointing
+from foundry.utils.torch import device_of
+
+
+class AttentionPairBiasPairformerDeepspeed(nn.Module):
+    def __init__(self, c_a, c_s, c_pair, n_head, kq_norm=False):
+        super().__init__()
+        self.n_head = n_head
+        self.c_a = c_a
+        self.c_pair = c_pair
+        self.c = c_a // n_head
+
+        self.to_q = MultiDimLinear(c_a, (n_head, self.c))
+        self.to_k = MultiDimLinear(c_a, (n_head, self.c), bias=False, norm=kq_norm)
+        self.to_v = MultiDimLinear(c_a, (n_head, self.c), bias=False, norm=kq_norm)
+        self.to_b = linearNoBias(c_pair, n_head)
+        self.to_g = nn.Sequential(
+            MultiDimLinear(c_a, (n_head, self.c), bias=False),
+            nn.Sigmoid(),
+        )
+        self.to_a = linearNoBias(c_a, c_a)
+        # self.linear_output_project = nn.Sequential(
+        # LinearBiasInit(c_s, c_a, biasinit=-2.),
+        # nn.Sigmoid(),
+        # )
+        self.ln_0 = RMSNorm((c_pair,))
+        # self.ada_ln_1 = AdaLN(c_a=c_a, c_s=c_s)
+        self.ln_1 = RMSNorm((c_a,))
+        self.use_deepspeed_evo = False
+        self.force_bfloat16 = True
+
+    def forward(
+        self,
+        A_I,  # [I, C_a]
+        S_I,  # [I, C_a] | None
+        Z_II,  # [I, I, C_z]
+        Beta_II=None,  # [I, I]
+    ):
+        # Input projections
+        assert S_I is None
+        A_I = self.ln_1(A_I)
+
+        if self.use_deepspeed_evo or self.force_bfloat16:
+            A_I = A_I.to(torch.bfloat16)
+
+        Q_IH = self.to_q(A_I)  # / np.sqrt(self.c)
+        K_IH = self.to_k(A_I)
+        V_IH = self.to_v(A_I)
+        B_IIH = self.to_b(self.ln_0(Z_II)) + Beta_II[..., None]
+        G_IH = self.to_g(A_I)
+
+        B, L = B_IIH.shape[:2]
+
+        if not self.use_deepspeed_evo or L <= 24:
+            Q_IH = Q_IH / torch.sqrt(
+                torch.tensor(self.c).to(Q_IH.device, torch.bfloat16)
+            )
+            # Attention
+            A_IIH = torch.softmax(
+                torch.einsum("...ihd,...jhd->...ijh", Q_IH, K_IH) + B_IIH, dim=-2
+            )  # softmax over j
+            ## G_IH: [I, H, C]
+            ## A_IIH: [I, I, H]
+            ## V_IH: [I, H, C]
+            A_I = torch.einsum("...ijh,...jhc->...ihc", A_IIH, V_IH)
+            A_I = G_IH * A_I  # [B, I, H, C]
+            A_I = A_I.flatten(start_dim=-2)  # [B, I, Ca]
+        else:
+            raise NotImplementedError
+
+        A_I = self.to_a(A_I)
+
+        return A_I
+
+
+class PairformerBlock(nn.Module):
+    """
+    Attempt to replicate AF3 architecture from scratch.
+    """
+
+    def __init__(
+        self,
+        c_s,
+        c_z,
+        attention_pair_bias,
+        p_drop=0.1,
+        triangle_multiplication=None,
+        triangle_attention=None,
+        n_transition=4,
+        use_deepspeed_evo=True,
+        use_triangle_mult=False,
+        use_triangle_attn=False,
+    ):
+        super().__init__()
+
+        # self.drop_row = Dropout(broadcast_dim=-2, p_drop=p_drop)
+        # self.drop_col = Dropout(broadcast_dim=-3, p_drop=p_drop)
+
+        self.z_transition = Transition(c=c_z, n=n_transition)
+
+        if c_s > 0:
+            self.s_transition = Transition(c=c_s, n=n_transition)
+
+            self.attention_pair_bias = AttentionPairBiasPairformerDeepspeed(
+                c_a=c_s, c_s=0, c_pair=c_z, **attention_pair_bias
+            )
+
+    @activation_checkpointing
+    def forward(self, S_I, Z_II):
+        with torch.amp.autocast(
+            device_type=device_of(self).type, enabled=True, dtype=torch.bfloat16
+        ):
+            Z_II = Z_II + self.z_transition(Z_II)
+            if S_I is not None:
+                S_I = S_I + self.attention_pair_bias(
+                    S_I, None, Z_II, Beta_II=torch.tensor([0.0], device=Z_II.device)
+                )
+                S_I = S_I + self.s_transition(S_I)
+        return S_I, Z_II

rfd3/run_inference.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env -S /bin/sh -c '"$(dirname "$0")/../../../../.ipd/shebang/rf3_exec.sh" "$0" "$@"'
+
+import os
+
+import hydra
+import rootutils
+from dotenv import load_dotenv
+from omegaconf import DictConfig, OmegaConf
+
+from rfd3.engine import RFD3InferenceConfig, RFD3InferenceEngine
+
+# Setup root dir and environment variables (more info: https://github.com/ashleve/rootutils)
+# NOTE: Sets the `PROJECT_ROOT` environment variable to the root directory of the project (where `.project-root` is located)
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+load_dotenv(override=True)
+
+# If the user has set `PROJECT_PATH`, use it to build the config path; otherwise, fall back to `PROJECT_ROOT`
+_config_path = os.path.join(os.environ["PROJECT_ROOT"], "models/rfd3/configs")
+
+
+@hydra.main(
+    config_path=_config_path,
+    config_name="inference",
+    version_base="1.3",
+)
+def run_inference(cfg: DictConfig) -> None:
+    """Execute the specified inference pipeline"""
+
+    run_params_set = {"inputs", "n_batches", "out_dir"}
+    run_params = {k: v for k, v in cfg.items() if k in run_params_set}
+
+    # Create __init__ args by filtering for all configs not in run_params
+    cfg_dict = OmegaConf.to_container(cfg, resolve=True)
+    init_cfg_dict = {
+        k: v for k, v in cfg_dict.items() if k not in run_params_set | {"_target_"}
+    }
+
+    # Run
+    engine = RFD3InferenceEngine(**RFD3InferenceConfig(**init_cfg_dict))
+    engine.run(**run_params)
+
+
+if __name__ == "__main__":
+    run_inference()

rfd3/testing/debug.py

@@ -0,0 +1,139 @@
+#!/usr/bin/env -S /bin/sh -c '"$(dirname "$0")/../../../scripts/shebang/modelhub_exec.sh" "$0" "$@"'
+# JBs debugging file, please create your own and go crazy!
+import logging
+import os
+import sys
+import time
+
+import hydra
+import ipdb  # noqa: F401
+import numpy as np
+import rootutils
+import torch
+import tree
+from atomworks.ml.utils.token import (
+    get_token_starts,
+)
+from rfd3.testing.testing_utils import (
+    TEST_CFG_TRAIN,
+    TEST_JSON_DATA,
+    build_pipelines,
+    instantiate_example,
+    load_train_or_val_cfg,
+)
+
+from foundry.utils.ddp import set_accelerator_based_on_availability
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Same as train.py
+rootutils.setup_root(__file__ + "/../..", indicator=".project-root", pythonpath=True)
+_config_path = os.path.join(
+    os.environ.get("PROJECT_PATH", os.environ.get("PROJECT_ROOT", "../..")), "configs"
+)
+print(f"Config path: {_config_path}")
+print(f"Project root: {os.environ.get('PROJECT_ROOT', '../..')}")
+
+
+is_inference = True
+args = TEST_JSON_DATA["1qys-1-refactored"]
+input = instantiate_example(args, is_inference=is_inference)
+
+
+TEST_CFG_TRAIN = (
+    load_train_or_val_cfg(name=sys.argv[1].split("=")[-1])
+    if len(sys.argv) > 1
+    else TEST_CFG_TRAIN
+)
+
+
+def forward(example, trainer, model, is_inference=is_inference):
+    network_input = trainer._assemble_network_inputs(example)
+
+    # Forward pass
+    device = "cuda:0"
+
+    def _inmap(path, x):
+        if hasattr(x, "cpu") and path != ("f", "msa_stack"):
+            return x.to(device)
+        else:
+            return x
+
+    network_input = tree.map_structure_with_path(_inmap, network_input)
+    model.eval() if is_inference else model.train()
+    if not is_inference:
+        network_output = model.forward(
+            input=network_input,
+            n_cycle=1,
+            coord_atom_lvl_to_be_noised=example["coord_atom_lvl_to_be_noised"].to(
+                device
+            ),
+        )
+    else:
+        with torch.no_grad():
+            network_output = model.forward(
+                input=network_input,
+                n_cycle=1,
+                coord_atom_lvl_to_be_noised=example["coord_atom_lvl_to_be_noised"].to(
+                    device
+                ),
+            )
+    return network_output
+
+
+def prep_forward(cfg):
+    trainer = hydra.utils.instantiate(
+        cfg.trainer,
+        loggers=None,
+        callbacks=None,
+        _convert_="partial",
+        _recursive_=False,
+    )
+    set_accelerator_based_on_availability(cfg)
+    trainer.initialize_or_update_trainer_state({"train_cfg": cfg})
+    cfg.trainer.devices_per_node = 1
+    cfg.trainer.num_nodes = 1
+    try:
+        trainer.fabric.launch()
+    except Exception as e:
+        print(f"Error: {e}")
+        print("Switching port")
+        os.environ["MASTER_PORT"] = str(1024 + np.random.randint(64512))
+        trainer.fabric.launch()
+    trainer.construct_model()
+    model = trainer.state["model"]
+    return model, trainer
+
+
+def test_conditional_forward():
+    unindexed_cfg = load_train_or_val_cfg("test-unindexed")
+    unindexed_cfg.datasets.global_transform_args.train_conditions.island.frequency = (
+        1e10
+    )
+    unindexed_cfg.datasets.global_transform_args.train_conditions.island.p_unindex_motif_tokens = 1.0
+    unindexed_pipes = build_pipelines(composed_config=unindexed_cfg)
+
+    t0 = time.time()
+    example = unindexed_pipes[is_inference](input)
+    example["example_id"] = "debug_example"
+    print(f"Time taken to process example: {time.time() - t0}")
+
+    aa = example["atom_array"]
+    t_aa = aa[get_token_starts(aa)]  # noqa: F841
+
+    from rfd3.testing.debug_utils import pipe_out_to_file
+
+    pipe_out_to_file(example, save=True)
+
+    print("Preparing model")
+    model, trainer = prep_forward(TEST_CFG_TRAIN)
+    if is_inference:
+        model.eval()
+        trainer.state["model"].eval()
+    network_output = forward(example, trainer, model, is_inference=is_inference)  # noqa: F841
+
+
+if __name__ == "__main__":
+    test_conditional_forward()
+    print("Finished main")

rfd3/testing/debug_utils.py

@@ -0,0 +1,73 @@
+import numpy as np
+from atomworks.common import sum_string_arrays
+from atomworks.io.utils.io_utils import to_cif_file
+from atomworks.ml.transforms.center_random_augmentation import CenterRandomAugmentation
+from biotite.structure import AtomArrayStack
+from rfd3.trainer.rfd3 import _reassign_unindexed_token_chains
+from rfd3.transforms.design_transforms import (
+    MotifCenterRandomAugmentation,
+)
+
+
+def pipe_out_to_file(output, save=True):
+    atom_array = output["atom_array"]
+
+    xyz = output["coord_atom_lvl_to_be_noised"]
+    idxs = np.argsort(output["t"].numpy())
+    eps = output["noise"].numpy()[idxs]
+    eps[0] = eps[0] * 0
+    x = AtomArrayStack(xyz.shape[0], xyz.shape[1])
+    x.coord = xyz[idxs].numpy() + eps
+
+    x.set_annotation("chain_id", ["A"] * xyz.shape[1])
+    x.set_annotation("atom_name", [f"C{i}" for i in range(x.shape[-1])])
+    x.set_annotation("res_id", output["feats"]["atom_to_token_map"])
+    x.set_annotation("element", ["C"] * x.shape[-1])
+    x.set_annotation("res_name", [atom_array.res_name[i] for i in range(x.shape[-1])])
+
+    if save:
+        f = f"{output.get('example_id', 'example')}_debug_out.cif"
+        to_cif_file(
+            x,
+            f,
+            id="x",
+        )
+        print("Saved cif file to:", f)
+    else:
+        return x
+
+
+def save_pipe_out(atom_array):
+    atom_array = _reassign_unindexed_token_chains(atom_array)
+
+    f = "debug_out.cif"
+    to_cif_file(
+        atom_array,
+        f,
+        id="x",
+    )
+    print("Saved cif file to:", f)
+
+
+def to_debug_pipe(pipe):
+    pipe.transforms = [
+        t
+        for t in pipe.transforms
+        if not isinstance(t, (CenterRandomAugmentation, MotifCenterRandomAugmentation))
+    ]
+    return pipe
+
+
+# Allows to use atom-array whenever debugging by removing friction in atoms having the same identifiers
+def save_debug_cif(atom_array, filepath, name="debug_out.cif"):
+    dummy_array = atom_array.copy()
+    dummy_array.chain_id = sum_string_arrays(
+        dummy_array.chain_id, "-", dummy_array.transformation_id.astype(str)
+    )
+
+    f = filepath + name
+    to_cif_file(
+        dummy_array,
+        f,
+    )
+    print("Saved cif file to:", f)