PXDesign / model /pxdesign /configs /configs_base.py

Upload folder using huggingface_hub (part 2)

4d20b62 verified about 2 months ago

6.51 kB

# Copyright 2025 ByteDance and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# pylint: disable=C0114,C0301
from protenix.config.extend_types import GlobalConfigValue, ValueMaybeNone

basic_configs = {
    "load_checkpoint_path": "",
    "load_strict": True,
    "deterministic": False,
    "model_name": "pxdesign_v0.1.0",  # train model name
}
model_configs = {
    # Model
    "c_s": 384,
    "c_z": 128,
    "c_s_inputs": 449,  # c_s_inputs == c_token + 32 + 32 + 1
    "c_atom": 128,
    "c_atompair": 16,
    "c_token": 384,
    "n_blocks": 48,
    "max_atoms_per_token": 24,  # DNA G max_atoms = 23
    "no_bins": 64,
    "sigma_data": 16.0,
    "diffusion_batch_size": 48,
    "diffusion_chunk_size": ValueMaybeNone(4),  # chunksize of diffusion_batch_size
    "blocks_per_ckpt": ValueMaybeNone(
        1
    ),  # NOTE: Number of blocks in each activation checkpoint, if None, no checkpointing is performed.
    # switch of kernels
    "use_memory_efficient_kernel": False,  # whether to use the torch.nn.functional.scaled_dot_product_attention, Defaults to False.
    "use_deepspeed_evo_attention": False,
    "use_flash": False,
    "use_lma": False,
    "use_xformer": False,
    "find_unused_parameters": False,
    "dtype": "bf16",  # default training dtype: bf16
    "loss_metrics_sparse_enable": True,  # the swicth for both sparse lddt metrics and sparse bond/smooth lddt loss
    "skip_amp": {
        "sample_diffusion": True,
        "confidence_head": True,
        "sample_diffusion_training": True,
        "loss": True,
    },
    "infer_setting": {
        "chunk_size": ValueMaybeNone(
            64
        ),  # should set to null for normal training and small dataset eval [for efficiency]
        "sample_diffusion_chunk_size": ValueMaybeNone(
            10
        ),  # should set to null for normal training and small dataset eval [for efficiency]
        "lddt_metrics_sparse_enable": GlobalConfigValue("loss_metrics_sparse_enable"),
        "lddt_metrics_chunk_size": ValueMaybeNone(
            1
        ),  # only works if loss_metrics_sparse_enable, can set as default 1
    },
    "inference_noise_scheduler": {
        "s_max": 160.0,
        "s_min": 4e-4,
        "rho": 7,
        "sigma_data": 16.0,  # NOTE: in EDM, this is 1.0
    },
    "sample_diffusion": {
        "gamma0": 1.0,
        "gamma_min": 0.01,
        "noise_scale_lambda": 1.003,
        "N_step": 400,
        "N_sample": 100,
        "eta_schedule": {"type": "piecewise_65", "min": 1.0, "max": 2.5},
    },
    "model": {
        "N_model_seed": 1,  # for inference
        "N_cycle": 4,
        "condition_embedding_drop_rate": 0.0,
        "confidence_embedding_drop_rate": 0.0,
        "input_embedder": {
            "c_atom": GlobalConfigValue("c_atom"),
            "c_atompair": GlobalConfigValue("c_atompair"),
            "c_token": GlobalConfigValue("c_token"),
        },
        "relative_position_encoding": {
            "r_max": 32,
            "s_max": 2,
            "c_z": GlobalConfigValue("c_z"),
        },
        "diffusion_module": {
            "use_fine_grained_checkpoint": True,
            "sigma_data": GlobalConfigValue("sigma_data"),
            "c_token": 768,
            "c_atom": GlobalConfigValue("c_atom"),
            "c_atompair": GlobalConfigValue("c_atompair"),
            "c_z": GlobalConfigValue("c_z"),
            "c_s": GlobalConfigValue("c_s"),
            "c_s_inputs": GlobalConfigValue("c_s_inputs"),
            "atom_encoder": {
                "n_blocks": 4,
                "n_heads": 4,
            },
            "transformer": {
                "n_blocks": 16,
                "n_heads": 16,
            },
            "atom_decoder": {
                "n_blocks": 4,
                "n_heads": 4,
            },
            "blocks_per_ckpt": GlobalConfigValue("blocks_per_ckpt"),
        },
    },
}

design_model_configs = {
    "condition_embedder": {
        "input_embedder": {
            "c_atom": GlobalConfigValue("c_atom"),
            "c_atompair": GlobalConfigValue("c_atompair"),
            "c_token": GlobalConfigValue("c_token"),
            "c_s_inputs": GlobalConfigValue("c_s_inputs"),
        },
        "template_embedder": {
            "c_templ_in": 64 + 1,
            "c_z": GlobalConfigValue("c_z"),
        },
    },
    "fuse_condition": {
        "c_s": GlobalConfigValue("c_s"),
        "c_z": GlobalConfigValue("c_z"),
        "c_s_inputs": GlobalConfigValue("c_s_inputs"),
        "fuse_s": False,
        "fuse_z": True,
        "fuse_s_inputs": False,
        "s_clamp": 30000.0,
    },
    "sequence_head": {
        "pairformer": {
            "n_blocks": 4,
            "c_z": GlobalConfigValue("c_z"),
            "c_s": GlobalConfigValue("c_s"),
            "n_heads": 8,
            "dropout": 0.20,
            "blocks_per_ckpt": GlobalConfigValue("blocks_per_ckpt"),
        },
        "seq_encoder": {
            "min_bin": 2.0,
            "max_bin": 12.0,
            "no_bins": 20,
            "eps": 1e-6,
            "c_z": GlobalConfigValue("c_z"),
            "c_s": GlobalConfigValue("c_s"),
            "restype_num": 32 + 4,
        },
        "seq_single_decoder": {
            "c_s": GlobalConfigValue("c_s"),
            "c_s_out": 21,
            "dropout_rate": 0.2,
        },
        "seq_pair_decoder": {
            "c_z": GlobalConfigValue("c_z"),
            "num_token": 21,
            "c_z_out": 441,
            "dropout_rate": 0.2,
        },
        "relative_position_encoding": {
            "r_max": 32,
            "s_max": 2,
            "c_z": GlobalConfigValue("c_z"),
        },
    },
    "design_diffusion_distogram": {
        "c_z": 768,
        "no_bins": GlobalConfigValue("no_bins"),
    },
    "design_distogram_head": {
        "c_z": GlobalConfigValue("c_z"),
        "no_bins": GlobalConfigValue("no_bins"),
    },
}
model_configs["model"].update(design_model_configs)

configs = {**basic_configs, **model_configs}

	# Copyright 2025 ByteDance and/or its affiliates.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.

	# pylint: disable=C0114,C0301
	from protenix.config.extend_types import GlobalConfigValue, ValueMaybeNone

	basic_configs = {
	"load_checkpoint_path": "",
	"load_strict": True,
	"deterministic": False,
	"model_name": "pxdesign_v0.1.0", # train model name
	}
	model_configs = {
	# Model
	"c_s": 384,
	"c_z": 128,
	"c_s_inputs": 449, # c_s_inputs == c_token + 32 + 32 + 1
	"c_atom": 128,
	"c_atompair": 16,
	"c_token": 384,
	"n_blocks": 48,
	"max_atoms_per_token": 24, # DNA G max_atoms = 23
	"no_bins": 64,
	"sigma_data": 16.0,
	"diffusion_batch_size": 48,
	"diffusion_chunk_size": ValueMaybeNone(4), # chunksize of diffusion_batch_size
	"blocks_per_ckpt": ValueMaybeNone(
	1
	), # NOTE: Number of blocks in each activation checkpoint, if None, no checkpointing is performed.
	# switch of kernels
	"use_memory_efficient_kernel": False, # whether to use the torch.nn.functional.scaled_dot_product_attention, Defaults to False.
	"use_deepspeed_evo_attention": False,
	"use_flash": False,
	"use_lma": False,
	"use_xformer": False,
	"find_unused_parameters": False,
	"dtype": "bf16", # default training dtype: bf16
	"loss_metrics_sparse_enable": True, # the swicth for both sparse lddt metrics and sparse bond/smooth lddt loss
	"skip_amp": {
	"sample_diffusion": True,
	"confidence_head": True,
	"sample_diffusion_training": True,
	"loss": True,
	},
	"infer_setting": {
	"chunk_size": ValueMaybeNone(
	64
	), # should set to null for normal training and small dataset eval [for efficiency]
	"sample_diffusion_chunk_size": ValueMaybeNone(
	10
	), # should set to null for normal training and small dataset eval [for efficiency]
	"lddt_metrics_sparse_enable": GlobalConfigValue("loss_metrics_sparse_enable"),
	"lddt_metrics_chunk_size": ValueMaybeNone(
	1
	), # only works if loss_metrics_sparse_enable, can set as default 1
	},
	"inference_noise_scheduler": {
	"s_max": 160.0,
	"s_min": 4e-4,
	"rho": 7,
	"sigma_data": 16.0, # NOTE: in EDM, this is 1.0
	},
	"sample_diffusion": {
	"gamma0": 1.0,
	"gamma_min": 0.01,
	"noise_scale_lambda": 1.003,
	"N_step": 400,
	"N_sample": 100,
	"eta_schedule": {"type": "piecewise_65", "min": 1.0, "max": 2.5},
	},
	"model": {
	"N_model_seed": 1, # for inference
	"N_cycle": 4,
	"condition_embedding_drop_rate": 0.0,
	"confidence_embedding_drop_rate": 0.0,
	"input_embedder": {
	"c_atom": GlobalConfigValue("c_atom"),
	"c_atompair": GlobalConfigValue("c_atompair"),
	"c_token": GlobalConfigValue("c_token"),
	},
	"relative_position_encoding": {
	"r_max": 32,
	"s_max": 2,
	"c_z": GlobalConfigValue("c_z"),
	},
	"diffusion_module": {
	"use_fine_grained_checkpoint": True,
	"sigma_data": GlobalConfigValue("sigma_data"),
	"c_token": 768,
	"c_atom": GlobalConfigValue("c_atom"),
	"c_atompair": GlobalConfigValue("c_atompair"),
	"c_z": GlobalConfigValue("c_z"),
	"c_s": GlobalConfigValue("c_s"),
	"c_s_inputs": GlobalConfigValue("c_s_inputs"),
	"atom_encoder": {
	"n_blocks": 4,
	"n_heads": 4,
	},
	"transformer": {
	"n_blocks": 16,
	"n_heads": 16,
	},
	"atom_decoder": {
	"n_blocks": 4,
	"n_heads": 4,
	},
	"blocks_per_ckpt": GlobalConfigValue("blocks_per_ckpt"),
	},
	},
	}

	design_model_configs = {
	"condition_embedder": {
	"input_embedder": {
	"c_atom": GlobalConfigValue("c_atom"),
	"c_atompair": GlobalConfigValue("c_atompair"),
	"c_token": GlobalConfigValue("c_token"),
	"c_s_inputs": GlobalConfigValue("c_s_inputs"),
	},
	"template_embedder": {
	"c_templ_in": 64 + 1,
	"c_z": GlobalConfigValue("c_z"),
	},
	},
	"fuse_condition": {
	"c_s": GlobalConfigValue("c_s"),
	"c_z": GlobalConfigValue("c_z"),
	"c_s_inputs": GlobalConfigValue("c_s_inputs"),
	"fuse_s": False,
	"fuse_z": True,
	"fuse_s_inputs": False,
	"s_clamp": 30000.0,
	},
	"sequence_head": {
	"pairformer": {
	"n_blocks": 4,
	"c_z": GlobalConfigValue("c_z"),
	"c_s": GlobalConfigValue("c_s"),
	"n_heads": 8,
	"dropout": 0.20,
	"blocks_per_ckpt": GlobalConfigValue("blocks_per_ckpt"),
	},
	"seq_encoder": {
	"min_bin": 2.0,
	"max_bin": 12.0,
	"no_bins": 20,
	"eps": 1e-6,
	"c_z": GlobalConfigValue("c_z"),
	"c_s": GlobalConfigValue("c_s"),
	"restype_num": 32 + 4,
	},
	"seq_single_decoder": {
	"c_s": GlobalConfigValue("c_s"),
	"c_s_out": 21,
	"dropout_rate": 0.2,
	},
	"seq_pair_decoder": {
	"c_z": GlobalConfigValue("c_z"),
	"num_token": 21,
	"c_z_out": 441,
	"dropout_rate": 0.2,
	},
	"relative_position_encoding": {
	"r_max": 32,
	"s_max": 2,
	"c_z": GlobalConfigValue("c_z"),
	},
	},
	"design_diffusion_distogram": {
	"c_z": 768,
	"no_bins": GlobalConfigValue("no_bins"),
	},
	"design_distogram_head": {
	"c_z": GlobalConfigValue("c_z"),
	"no_bins": GlobalConfigValue("no_bins"),
	},
	}
	model_configs["model"].update(design_model_configs)

	configs = {basic_configs, model_configs}