Add files using upload-large-folder tool

config/callbacks/default.yaml

@@ -0,0 +1,5 @@
+defaults:
+  - model_checkpoint
+  - model_summary
+  - _self_
+

config/callbacks/model_summary.yaml

@@ -0,0 +1,5 @@
+# https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.RichModelSummary.html
+
+model_summary:
+  _target_: lightning.pytorch.callbacks.RichModelSummary
+  max_depth: -1 # the maximum depth of layer nesting that the summary will include

config/data/apo.yaml

@@ -0,0 +1,3 @@
+_target_: onescience.datapipes.simplefold.test_datamodule.SimpleFoldInferenceDataModule
+target_dir: data/apo/
+num_workers: 1

config/data/pdb_sp_afesm.yaml

@@ -0,0 +1,66 @@
+_target_: onescience.datapipes.simplefold.train_datamodule.SimpleFoldTrainingDataModule
+
+datasets:
+  - _target_: onescience.datapipes.simplefold.train_datamodule.DatasetConfig
+    data_name: rcsb_protein
+    tokenized_dir: data/rcsb_protein_tokenized
+    target_dir: data/rcsb_processed_targets
+    manifest_path: data/rcsb_protein_tokenized/manifest.json
+    cropper:
+      _target_: onescience.datapipes.boltz_data_pipeline.crop.boltz.BoltzCropper
+      min_neighborhood: 0
+      max_neighborhood: 40
+    filters:
+      - _target_: onescience.datapipes.boltz_data_pipeline.filter.dynamic.resolution.ResolutionFilter
+        resolution: 5.0
+      - _target_: onescience.datapipes.boltz_data_pipeline.filter.dynamic.date.DateFilter
+        date: "2020-05-01"
+        ref: released
+
+  - _target_: onescience.datapipes.simplefold.train_datamodule.DatasetConfig
+    data_name: swissprot
+    tokenized_dir: data/swissprot_tokenized
+    target_dir: data/swissprot_pdb_v4_boltz
+    manifest_path: data/swissprot_tokenized/manifest.json
+    record_list: data/swissprot_filtered.txt
+    cropper:
+      _target_: onescience.datapipes.boltz_data_pipeline.crop.boltz.BoltzCropper
+      min_neighborhood: 0
+      max_neighborhood: 40
+
+  - _target_: onescience.datapipes.simplefold.train_datamodule.DatasetConfig
+    data_name: processed_afesm
+    tokenized_dir: data/processed_afesm
+    target_dir: data/processed_afesm
+    manifest_path: data/processed_afesm/manifest.json
+    cropper:
+      _target_: onescience.datapipes.boltz_data_pipeline.crop.boltz.BoltzCropper
+      min_neighborhood: 0
+      max_neighborhood: 40
+
+filters:
+    - _target_: onescience.datapipes.boltz_data_pipeline.filter.dynamic.size.SizeFilter
+      min_chains: 1
+      max_chains: 300
+
+tokenizer:
+  _target_: onescience.datapipes.boltz_data_pipeline.tokenize.boltz_protein.BoltzTokenizer
+
+featurizer:
+  _target_: onescience.datapipes.boltz_data_pipeline.feature.featurizer.BoltzFeaturizer
+
+symmetries: data/symmetry.pkl
+max_tokens: 256
+max_atoms: 2304
+pad_to_max_tokens: False
+pad_to_max_atoms: False
+batch_size: 256
+num_workers: 16
+pin_memory: True
+return_train_symmetries: False
+min_dist: 2.0
+max_dist: 22.0
+num_bins: 64
+atoms_per_window_queries: 32
+rotation_augment_ref_pos: False
+rotation_augment_coords: True

config/data/pdb_sp_afesme.yaml

@@ -0,0 +1,67 @@
+_target_: onescience.datapipes.simplefold.train_datamodule.SimpleFoldTrainingDataModule
+
+datasets:
+  - _target_: onescience.datapipes.simplefold.train_datamodule.DatasetConfig
+    data_name: rcsb_protein
+    tokenized_dir: data/rcsb_protein_tokenized
+    target_dir: data/rcsb_processed_targets
+    manifest_path: data/rcsb_protein_tokenized/manifest.json
+    cropper:
+      _target_: onescience.datapipes.boltz_data_pipeline.crop.boltz.BoltzCropper
+      min_neighborhood: 0
+      max_neighborhood: 40
+    filters:
+      - _target_: onescience.datapipes.boltz_data_pipeline.filter.dynamic.resolution.ResolutionFilter
+        resolution: 5.0
+      - _target_: onescience.datapipes.boltz_data_pipeline.filter.dynamic.date.DateFilter
+        date: "2020-05-01"
+        ref: released
+
+  - _target_: onescience.datapipes.simplefold.train_datamodule.DatasetConfig
+    data_name: swissprot
+    tokenized_dir: data/swissprot_tokenized
+    target_dir: data/swissprot_pdb_v4_boltz
+    manifest_path: data/swissprot_tokenized/manifest.json
+    record_list: data/swissprot_filtered.txt
+    cropper:
+      _target_: onescience.datapipes.boltz_data_pipeline.crop.boltz.BoltzCropper
+      min_neighborhood: 0
+      max_neighborhood: 40
+
+  - _target_: onescience.datapipes.simplefold.train_datamodule.DatasetConfig
+    data_name: processed_afesm_cluster
+    tokenized_dir: data/tokenized_afesm_cluster
+    target_dir: data/processed_afesm_cluster
+    manifest_path: data/processed_afesm_cluster/manifest.json
+    cluster: data/filtered_afesm_cluster.json
+    cropper:
+      _target_: onescience.datapipes.boltz_data_pipeline.crop.boltz.BoltzCropper
+      min_neighborhood: 0
+      max_neighborhood: 40
+
+filters:
+    - _target_: onescience.datapipes.boltz_data_pipeline.filter.dynamic.size.SizeFilter
+      min_chains: 1
+      max_chains: 300
+
+tokenizer:
+  _target_: onescience.datapipes.boltz_data_pipeline.tokenize.boltz_protein.BoltzTokenizer
+
+featurizer:
+  _target_: onescience.datapipes.boltz_data_pipeline.feature.featurizer.BoltzFeaturizer
+
+symmetries: data/symmetry.pkl
+max_tokens: 256
+max_atoms: 2304
+pad_to_max_tokens: False
+pad_to_max_atoms: False
+batch_size: 256
+num_workers: 16
+pin_memory: True
+return_train_symmetries: False
+min_dist: 2.0
+max_dist: 22.0
+num_bins: 64
+atoms_per_window_queries: 32
+rotation_augment_ref_pos: False
+rotation_augment_coords: True

config/experiment/test.yaml

@@ -0,0 +1,30 @@
+# @package _global_
+defaults:
+  - override /data: cameo22
+  - override /model: simplefold
+  - override /model/architecture: foldingdit_100M
+  - override /model/sampler: euler_maruyama
+  - override /model/processor: protein_processor
+  - override /callbacks: default
+  - override /trainer: default
+
+load_ckpt_path: artifacts/simplefold_100M.ckpt
+seed: 12345
+
+data:
+  num_repeats: 1
+  filters:
+    - _target_: onescience.datapipes.boltz_data_pipeline.filter.dynamic.max_residues.MaxResiduesFilter
+      min_residues: 1
+      max_residues: 1000 # filter out T1044 in CASP14
+
+model:
+  sampler:
+    num_timesteps: 500
+    t_start: 1e-4
+    tau: 0.01
+    log_timesteps: True
+    w_cutoff: 0.99
+
+  processor:
+    inference_multiplicity: 1

config/experiment/train.yaml

@@ -0,0 +1,55 @@
+# @package _global_
+
+# All parameters below will be merged with parameters from default configurations set above 
+# this allows you to overwrite only specified parameters
+
+# Options for training:
+# - /data: [pdb, pdb_sp, pdb_sp_afesm, pdb_sp_afesme]
+# - /model/architecture: 
+#         [foldingdit_100M, foldingdit_360M, foldingdit_700M, 
+#          foldingdit_1.1B, foldingdit_1.6B, foldingdit_3B]
+# - /trainer: [default, fsdp]
+
+defaults:
+  - override /data: pdb
+  - override /model: simplefold
+  - override /model/architecture: foldingdit_100M
+  - override /model/sampler: euler_maruyama
+  - override /model/processor: protein_processor
+  - override /callbacks: default
+  - override /trainer: default
+  - override /logger: tensorboard
+
+# load_ckpt_path: [PATH_TO_YOUR_CKPT] # uncomment to load a checkpoint
+seed: 12345
+
+trainer:
+  max_steps: 300000
+  val_check_interval: 10000 # this controls BOTH checkpoint steps and validation
+  accumulate_grad_batches: 1 # gradient accumulation
+
+data:
+  batch_size: 1 # should be the same as the number of GPUs
+  num_workers: 16
+
+model:
+  _target_: models.simplefold.simplefold.SimpleFold
+  ema_decay: 0.999
+  clip_grad_norm_val: 2.0
+  esm_model: "esm2_3B"
+  lddt_weight_schedule: False # set to True in finetuning phase
+
+  scheduler:
+    _target_: onescience.utils.simplefold.lr_scheduler.LinearWarmup
+    _partial_: true
+    min_lr: 1e-6
+    max_lr: ${model.optimizer.lr}
+    warmup_steps: 5000
+
+  architecture:
+    esm_model: ${model.esm_model}
+
+  processor:
+    scale: 16.0
+    ref_scale: 5.0
+    multiplicity: 16 # number of copies per GPU

config/experiment/train_plddt.yaml

@@ -0,0 +1,85 @@
+# @package _global_
+
+# All parameters below will be merged with parameters from default configurations set above 
+# this allows you to overwrite only specified parameters
+
+# Options for training:
+# - /data: [pdb, pdb_sp, pdb_sp_afesm, pdb_sp_afesme]
+# - /model/architecture: 
+#         [foldingdit_100M, foldingdit_360M, foldingdit_700M, 
+#          foldingdit_1.1B, foldingdit_1.6B, foldingdit_3B]
+# - /trainer: [default, fsdp]
+
+defaults:
+  - override /data: pdb
+  - override /model: simplefold
+  - override /model/architecture: foldingdit_100M
+  - override /model/sampler: euler_maruyama
+  - override /model/processor: protein_processor
+  - override /callbacks: default
+  - override /trainer: default
+  - override /logger: tensorboard
+
+# load_ckpt_path: [PATH_TO_YOUR_CKPT] # uncomment to load a checkpoint
+seed: 12345
+
+trainer:
+  max_steps: 10000
+  val_check_interval: 1000 # this controls BOTH checkpoint steps and validation
+  accumulate_grad_batches: 1 # gradient accumulation
+
+data:
+  batch_size: 8 # should be the same as the number of GPUs
+  num_workers: 16
+
+model:
+  _target_: models.simplefold.simplefold.SimpleFold
+  ema_decay: 0.999
+  clip_grad_norm_val: 2.0
+  esm_model: "esm2_3B"
+  plddt_training: True
+
+  plddt_module:
+    _target_: models.simplefold.torch.confidence_module.ConfidenceModule
+    hidden_size: ${model.architecture.hidden_size}
+    num_plddt_bins: 50
+    transformer_blocks: 
+      _target_: models.simplefold.torch.blocks.HomogenTrunk
+      depth: 4
+      block:
+        _target_: models.simplefold.torch.blocks.TransformerBlock
+        _partial_: true
+        hidden_size: ${model.architecture.hidden_size}
+        mlp_ratio: 4.0
+        use_swiglu: True # SwiGLU FFN
+        self_attention_layer:
+          _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+          _partial_: True
+          hidden_size: ${model.architecture.hidden_size}
+          num_heads: ${model.architecture.num_heads}
+          qk_norm: True
+          pos_embedder:
+            _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+            in_dim: 4
+            embed_dim: ${model.architecture.hidden_size}
+            num_heads: ${model.architecture.num_heads}
+            base: 100.0
+
+  scheduler:
+    _target_: onescience.utils.simplefold.lr_scheduler.LinearWarmup
+    _partial_: true
+    min_lr: 1e-6
+    max_lr: ${model.optimizer.lr}
+    warmup_steps: 5000
+
+  architecture:
+    esm_model: ${model.esm_model}
+
+  sampler:
+    num_timesteps: 150
+    tau: 0.3
+
+  processor:
+    scale: 16.0
+    ref_scale: 5.0
+    multiplicity: 16 # number of copies per GPU

config/extras/default.yaml

@@ -0,0 +1,5 @@
+# disable python warnings if they annoy you
+ignore_warnings: False
+
+# pretty print config tree at the start of the run using Rich library
+print_config: True

config/model/architecture/default.yaml

@@ -0,0 +1,101 @@
+_target_: models.simplefold.torch.architecture.FoldingDiT
+
+hidden_size: 768
+num_heads: 12
+atom_num_heads: 4
+output_channels: 3
+use_atom_mask: False
+use_length_condition: True
+esm_dropout_prob: 0.0
+esm_model: esm2_3B
+
+time_embedder:
+  _target_: models.simplefold.torch.layers.TimestepEmbedder
+  hidden_size: 768
+
+aminoacid_pos_embedder:
+  _target_: models.simplefold.torch.pos_embed.AbsolutePositionEncoding
+  in_dim: 1
+  embed_dim: 768
+  include_input: True
+
+pos_embedder:
+  _target_: models.simplefold.torch.pos_embed.FourierPositionEncoding
+  in_dim: 3
+  include_input: True
+  min_freq_log2: 0
+  max_freq_log2: 12
+  num_freqs: 128
+  log_sampling: True
+
+trunk:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 8
+  block:
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True
+    hidden_size: 768
+    mlp_ratio: 4.0
+    use_swiglu: True
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 768
+      num_heads: 12
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 768
+        num_heads: 12
+        base: 100.0
+
+atom_hidden_size_enc: 256
+atom_n_queries_enc: 32
+atom_n_keys_enc: 128
+atom_encoder_transformer:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 1
+  block:
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True
+    hidden_size: 256
+    mlp_ratio: 4.0
+    use_swiglu: True
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 256
+      num_heads: 4
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 256
+        num_heads: 4
+        base: 100.0
+
+atom_hidden_size_dec: 256
+atom_n_queries_dec: 32
+atom_n_keys_dec: 128
+atom_decoder_transformer:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 1
+  block:
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True
+    hidden_size: 256
+    mlp_ratio: 4.0
+    use_swiglu: True
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 256
+      num_heads: 4
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 256
+        num_heads: 4
+        base: 100.0

config/model/architecture/foldingdit_1.6B.yaml

@@ -0,0 +1,99 @@
+_target_: models.simplefold.torch.architecture.FoldingDiT
+
+hidden_size: 1536
+num_heads: 24
+atom_num_heads: 8
+output_channels: 3
+use_atom_mask: False
+use_length_condition: True
+esm_dropout_prob: 0.0
+esm_model: esm2_3B
+
+time_embedder:
+  _target_: models.simplefold.torch.layers.TimestepEmbedder
+  hidden_size: 1536
+aminoacid_pos_embedder:
+  _target_: models.simplefold.torch.pos_embed.AbsolutePositionEncoding
+  in_dim: 1
+  embed_dim: 1536
+  include_input: True
+pos_embedder:
+  _target_: models.simplefold.torch.pos_embed.FourierPositionEncoding
+  in_dim: 3
+  include_input: True
+  min_freq_log2: 0
+  max_freq_log2: 12
+  num_freqs: 128
+  log_sampling: True
+
+trunk:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 36
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 1536
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 1536
+      num_heads: 24
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 1536
+        num_heads: 24
+        base: 100.0
+
+atom_hidden_size_enc: 512
+atom_n_queries_enc: 32
+atom_n_keys_enc: 128
+atom_encoder_transformer:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 3
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 512
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 512
+      num_heads: 8
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 512
+        num_heads: 8
+        base: 100.0
+
+atom_hidden_size_dec: 512
+atom_n_queries_dec: 32
+atom_n_keys_dec: 128
+atom_decoder_transformer:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 3
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 512
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 512
+      num_heads: 8
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 512
+        num_heads: 8
+        base: 100.0

config/model/architecture/plddt_module.yaml

@@ -0,0 +1,24 @@
+_target_: models.simplefold.torch.confidence_module.ConfidenceModule
+hidden_size: 1536
+num_plddt_bins: 50
+transformer_blocks: 
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 4
+  block:
+    _target_: models.simplefold.torch.blocks.TransformerBlock
+    _partial_: true
+    hidden_size: 1536
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 1536
+      num_heads: 24
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 1536
+        num_heads: 24
+        base: 100.0

config/model/default.yaml

@@ -0,0 +1,39 @@
+_target_: models.simplefold.simplefold.SimpleFold
+ema_decay: 0.999
+clip_grad_norm_val: 2.0
+use_rigid_align: True
+smooth_lddt_loss_weight: 1.0
+lddt_cutoff: 15.0
+esm_model: "esm2_3B"
+lddt_weight_schedule: False
+sample_dir: ${paths.sample_dir}
+
+architecture:
+  esm_model: ${model.esm_model}
+
+path:
+  _target_: models.simplefold.flow.LinearPath
+
+loss:
+  _target_: torch.nn.MSELoss
+  reduction: 'none'
+  reduce: False
+
+optimizer:
+  _target_: torch.optim.AdamW
+  _partial_: true
+  lr: 1e-4
+  weight_decay: 0.0
+  max_steps: ???
+
+scheduler:
+  _target_: onescience.utils.simplefold.lr_scheduler.LinearWarmup
+  _partial_: true
+  min_lr: 1e-6
+  max_lr: ${model.optimizer.lr}
+  warmup_steps: 10
+
+processor:
+  scale: 16.0
+  ref_scale: 5.0
+  multiplicity: 16

config/model/processor/default.yaml

@@ -0,0 +1,2 @@
+_target_: onescience.datapipes.simplefold.processor.protein_processor.ProteinDataProcessor
+_partial_: True

config/model/processor/protein_processor.yaml

@@ -0,0 +1,2 @@
+_target_: onescience.datapipes.simplefold.processor.protein_processor.ProteinDataProcessor
+_partial_: True

config/model/simplefold.yaml

@@ -0,0 +1,39 @@
+_target_: models.simplefold.simplefold.SimpleFold
+ema_decay: 0.999
+clip_grad_norm_val: 2.0
+use_rigid_align: True
+smooth_lddt_loss_weight: 1.0
+lddt_cutoff: 15.0
+esm_model: "esm2_3B"
+lddt_weight_schedule: False
+sample_dir: ${paths.sample_dir}
+
+architecture:
+  esm_model: ${model.esm_model}
+
+path: 
+  _target_: models.simplefold.flow.LinearPath
+
+loss:
+  _target_: torch.nn.MSELoss
+  reduction: 'none'
+  reduce: False
+
+optimizer:
+  _target_: torch.optim.AdamW
+  _partial_: true
+  lr: 1e-4
+  weight_decay: 0.0
+  max_steps: ???
+
+scheduler:
+  _target_: onescience.utils.simplefold.lr_scheduler.LinearWarmup
+  _partial_: true
+  min_lr: 1e-6
+  max_lr: ${model.optimizer.lr}
+  warmup_steps: 10
+
+processor:
+  scale: 16.0
+  ref_scale: 5.0
+  multiplicity: 16

config/trainer/fsdp.yaml

@@ -0,0 +1,18 @@
+_target_: lightning.pytorch.trainer.Trainer
+
+default_root_dir: ${paths.output_dir}
+
+# set True to to ensure deterministic results
+# makes training slower but gives more reproducibility than just setting seeds
+deterministic: False
+
+accelerator: gpu
+devices: auto
+num_nodes: 1
+min_epochs: 1 # prevents early stopping
+max_steps: 100000
+precision: bf16-mixed
+val_check_interval: 5000
+check_val_every_n_epoch: null
+num_sanity_val_steps: 0
+limit_val_batches: 0.0

models/esm/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from onescience.datapipes import esm as data  # noqa
+from onescience.datapipes.esm import Alphabet, BatchConverter, FastaBatchedDataset  # noqa
+
+from .version import version as __version__  # noqa
+from .esm1 import ProteinBertModel  # noqa
+from .esm2 import ESM2  # noqa
+from .msa_transformer import MSATransformer  #noqa
+from . import pretrained  # noqa

models/esm/esm2.py

@@ -0,0 +1,152 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Union
+import torch
+import torch.nn as nn
+
+from onescience.datapipes.esm import Alphabet
+from onescience.modules.esm import (
+    ContactPredictionHead,
+    ESM1bLayerNorm,
+    RobertaLMHead,
+    TransformerLayer,
+)
+
+
+class ESM2(nn.Module):
+    def __init__(
+        self,
+        num_layers: int = 33,
+        embed_dim: int = 1280,
+        attention_heads: int = 20,
+        alphabet: Union[Alphabet, str] = "ESM-1b",
+        token_dropout: bool = True,
+    ):
+        super().__init__()
+        self.num_layers = num_layers
+        self.embed_dim = embed_dim
+        self.attention_heads = attention_heads
+        if not isinstance(alphabet, Alphabet):
+            alphabet = Alphabet.from_architecture(alphabet)
+        self.alphabet = alphabet
+        self.alphabet_size = len(alphabet)
+        self.padding_idx = alphabet.padding_idx
+        self.mask_idx = alphabet.mask_idx
+        self.cls_idx = alphabet.cls_idx
+        self.eos_idx = alphabet.eos_idx
+        self.prepend_bos = alphabet.prepend_bos
+        self.append_eos = alphabet.append_eos
+        self.token_dropout = token_dropout
+
+        self._init_submodules()
+
+    def _init_submodules(self):
+        self.embed_scale = 1
+        self.embed_tokens = nn.Embedding(
+            self.alphabet_size,
+            self.embed_dim,
+            padding_idx=self.padding_idx,
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                TransformerLayer(
+                    self.embed_dim,
+                    4 * self.embed_dim,
+                    self.attention_heads,
+                    add_bias_kv=False,
+                    use_esm1b_layer_norm=True,
+                    use_rotary_embeddings=True,
+                )
+                for _ in range(self.num_layers)
+            ]
+        )
+
+        self.contact_head = ContactPredictionHead(
+            self.num_layers * self.attention_heads,
+            self.prepend_bos,
+            self.append_eos,
+            eos_idx=self.eos_idx,
+        )
+        self.emb_layer_norm_after = ESM1bLayerNorm(self.embed_dim)
+
+        self.lm_head = RobertaLMHead(
+            embed_dim=self.embed_dim,
+            output_dim=self.alphabet_size,
+            weight=self.embed_tokens.weight,
+        )
+
+    def forward(self, tokens, repr_layers=[], need_head_weights=False, return_contacts=False):
+        if return_contacts:
+            need_head_weights = True
+
+        assert tokens.ndim == 2
+        padding_mask = tokens.eq(self.padding_idx)  # B, T
+
+        x = self.embed_scale * self.embed_tokens(tokens)
+
+        if self.token_dropout:
+            x.masked_fill_((tokens == self.mask_idx).unsqueeze(-1), 0.0)
+            # x: B x T x C
+            mask_ratio_train = 0.15 * 0.8
+            src_lengths = (~padding_mask).sum(-1)
+            mask_ratio_observed = (tokens == self.mask_idx).sum(-1).to(x.dtype) / src_lengths
+            x = x * (1 - mask_ratio_train) / (1 - mask_ratio_observed)[:, None, None]
+
+        if padding_mask is not None:
+            x = x * (1 - padding_mask.unsqueeze(-1).type_as(x))
+
+        repr_layers = set(repr_layers)
+        hidden_representations = {}
+        if 0 in repr_layers:
+            hidden_representations[0] = x
+
+        if need_head_weights:
+            attn_weights = []
+
+        # (B, T, E) => (T, B, E)
+        x = x.transpose(0, 1)
+
+        if not padding_mask.any():
+            padding_mask = None
+
+        for layer_idx, layer in enumerate(self.layers):
+            x, attn = layer(
+                x,
+                self_attn_padding_mask=padding_mask,
+                need_head_weights=need_head_weights,
+            )
+            if (layer_idx + 1) in repr_layers:
+                hidden_representations[layer_idx + 1] = x.transpose(0, 1)
+            if need_head_weights:
+                # (H, B, T, T) => (B, H, T, T)
+                attn_weights.append(attn.transpose(1, 0))
+
+        x = self.emb_layer_norm_after(x)
+        x = x.transpose(0, 1)  # (T, B, E) => (B, T, E)
+
+        # last hidden representation should have layer norm applied
+        if (layer_idx + 1) in repr_layers:
+            hidden_representations[layer_idx + 1] = x
+        x = self.lm_head(x)
+
+        result = {"logits": x, "representations": hidden_representations}
+        if need_head_weights:
+            # attentions: B x L x H x T x T
+            attentions = torch.stack(attn_weights, 1)
+            if padding_mask is not None:
+                attention_mask = 1 - padding_mask.type_as(attentions)
+                attention_mask = attention_mask.unsqueeze(1) * attention_mask.unsqueeze(2)
+                attentions = attentions * attention_mask[:, None, None, :, :]
+            result["attentions"] = attentions
+            if return_contacts:
+                contacts = self.contact_head(tokens, attentions)
+                result["contacts"] = contacts
+
+        return result
+
+    def predict_contacts(self, tokens):
+        return self(tokens, return_contacts=True)["contacts"]

models/esm/esmfold/v1/esmfold.py

@@ -0,0 +1,364 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import typing as T
+from dataclasses import dataclass, field
+from functools import partial
+
+import torch
+import torch.nn as nn
+from torch import nn
+from torch.nn import LayerNorm
+
+import onescience.models.esm as esm
+from onescience.models.esm import Alphabet
+from onescience.models.esm.esmfold.v1.categorical_mixture import categorical_lddt
+from onescience.models.esm.esmfold.v1.misc import (
+    batch_encode_sequences,
+    collate_dense_tensors,
+    output_to_pdb,
+)
+from onescience.models.esm.esmfold.v1.trunk import FoldingTrunk, FoldingTrunkConfig
+from onescience.datapipes.openfold.data_transforms import make_atom14_masks
+from onescience.utils.openfold.np import residue_constants
+from onescience.utils.openfold.loss import compute_predicted_aligned_error, compute_tm
+
+
+@dataclass
+class ESMFoldConfig:
+    trunk: T.Any = field(default_factory=FoldingTrunkConfig)
+    lddt_head_hid_dim: int = 128
+
+
+load_fn = esm.pretrained.load_model_and_alphabet
+esm_registry = {
+    "esm2_8M": partial(load_fn, "esm2_t6_8M_UR50D_500K"),
+    "esm2_8M_270K": esm.pretrained.esm2_t6_8M_UR50D,
+    "esm2_35M": partial(load_fn, "esm2_t12_35M_UR50D_500K"),
+    "esm2_35M_270K": esm.pretrained.esm2_t12_35M_UR50D,
+    "esm2_150M": partial(load_fn, "esm2_t30_150M_UR50D_500K"),
+    "esm2_150M_270K": partial(load_fn, "esm2_t30_150M_UR50D_270K"),
+    "esm2_650M": esm.pretrained.esm2_t33_650M_UR50D,
+    "esm2_650M_270K": partial(load_fn, "esm2_t33_650M_270K_UR50D"),
+    "esm2_3B": esm.pretrained.esm2_t36_3B_UR50D,
+    "esm2_3B_270K": partial(load_fn, "esm2_t36_3B_UR50D_500K"),
+    "esm2_15B": esm.pretrained.esm2_t48_15B_UR50D,
+}
+
+
+class ESMFold(nn.Module):
+    def __init__(self, esmfold_config=None, **kwargs):
+        super().__init__()
+
+        self.cfg = esmfold_config if esmfold_config else ESMFoldConfig(**kwargs)
+        cfg = self.cfg
+
+        self.distogram_bins = 64
+
+        self.esm, self.esm_dict = esm_registry.get(cfg.esm_type)()
+
+        self.esm.requires_grad_(False)
+        self.esm.half()
+
+        self.esm_feats = self.esm.embed_dim
+        self.esm_attns = self.esm.num_layers * self.esm.attention_heads
+        self.register_buffer("af2_to_esm", ESMFold._af2_to_esm(self.esm_dict))
+        self.esm_s_combine = nn.Parameter(torch.zeros(self.esm.num_layers + 1))
+
+        c_s = cfg.trunk.sequence_state_dim
+        c_z = cfg.trunk.pairwise_state_dim
+
+        self.esm_s_mlp = nn.Sequential(
+            LayerNorm(self.esm_feats),
+            nn.Linear(self.esm_feats, c_s),
+            nn.ReLU(),
+            nn.Linear(c_s, c_s),
+        )
+        if cfg.use_esm_attn_map:
+            self.esm_z_mlp = nn.Sequential(
+                LayerNorm(self.esm_attns),
+                nn.Linear(self.esm_attns, c_z),
+                nn.ReLU(),
+                nn.Linear(c_z, c_z),
+            )
+
+        # 0 is padding, N is unknown residues, N + 1 is mask.
+        self.n_tokens_embed = residue_constants.restype_num + 3
+        self.pad_idx = 0
+        self.unk_idx = self.n_tokens_embed - 2
+        self.mask_idx = self.n_tokens_embed - 1
+        self.embedding = nn.Embedding(self.n_tokens_embed, c_s, padding_idx=0)
+
+        self.trunk = FoldingTrunk(**cfg.trunk)
+
+        self.distogram_head = nn.Linear(c_z, self.distogram_bins)
+        self.ptm_head = nn.Linear(c_z, self.distogram_bins)
+        self.lm_head = nn.Linear(c_s, self.n_tokens_embed)
+        self.lddt_bins = 50
+        self.lddt_head = nn.Sequential(
+            nn.LayerNorm(cfg.trunk.structure_module.c_s),
+            nn.Linear(cfg.trunk.structure_module.c_s, cfg.lddt_head_hid_dim),
+            nn.Linear(cfg.lddt_head_hid_dim, cfg.lddt_head_hid_dim),
+            nn.Linear(cfg.lddt_head_hid_dim, 37 * self.lddt_bins),
+        )
+
+    @staticmethod
+    def _af2_to_esm(d: Alphabet):
+        # Remember that t is shifted from residue_constants by 1 (0 is padding).
+        esm_reorder = [d.padding_idx] + [
+            d.get_idx(v) for v in residue_constants.restypes_with_x
+        ]
+        return torch.tensor(esm_reorder)
+
+    def _af2_idx_to_esm_idx(self, aa, mask):
+        aa = (aa + 1).masked_fill(mask != 1, 0)
+        return self.af2_to_esm[aa]
+
+    def _compute_language_model_representations(
+        self, esmaa: torch.Tensor
+    ) -> torch.Tensor:
+        """Adds bos/eos tokens for the language model, since the structure module doesn't use these."""
+        batch_size = esmaa.size(0)
+
+        bosi, eosi = self.esm_dict.cls_idx, self.esm_dict.eos_idx
+        bos = esmaa.new_full((batch_size, 1), bosi)
+        eos = esmaa.new_full((batch_size, 1), self.esm_dict.padding_idx)
+        esmaa = torch.cat([bos, esmaa, eos], dim=1)
+        # Use the first padding index as eos during inference.
+        esmaa[range(batch_size), (esmaa != 1).sum(1)] = eosi
+
+        res = self.esm(
+            esmaa,
+            repr_layers=range(self.esm.num_layers + 1),
+            need_head_weights=self.cfg.use_esm_attn_map,
+        )
+        esm_s = torch.stack(
+            [v for _, v in sorted(res["representations"].items())], dim=2
+        )
+        esm_s = esm_s[:, 1:-1]  # B, L, nLayers, C
+        esm_z = (
+            res["attentions"].permute(0, 4, 3, 1, 2).flatten(3, 4)[:, 1:-1, 1:-1, :]
+            if self.cfg.use_esm_attn_map
+            else None
+        )
+        return esm_s, esm_z
+
+    def _mask_inputs_to_esm(self, esmaa, pattern):
+        new_esmaa = esmaa.clone()
+        new_esmaa[pattern == 1] = self.esm_dict.mask_idx
+        return new_esmaa
+
+    def forward(
+        self,
+        aa: torch.Tensor,
+        mask: T.Optional[torch.Tensor] = None,
+        residx: T.Optional[torch.Tensor] = None,
+        masking_pattern: T.Optional[torch.Tensor] = None,
+        num_recycles: T.Optional[int] = None,
+    ):
+        """Runs a forward pass given input tokens. Use `model.infer` to
+        run inference from a sequence.
+
+        Args:
+            aa (torch.Tensor): Tensor containing indices corresponding to amino acids. Indices match
+                modules.utils.openfold.np.residue_constants.restype_order_with_x.
+            mask (torch.Tensor): Binary tensor with 1 meaning position is unmasked and 0 meaning position is masked.
+            residx (torch.Tensor): Residue indices of amino acids. Will assume contiguous if not provided.
+            masking_pattern (torch.Tensor): Optional masking to pass to the input. Binary tensor of the same size
+                as `aa`. Positions with 1 will be masked. ESMFold sometimes produces different samples when
+                different masks are provided.
+            num_recycles (int): How many recycle iterations to perform. If None, defaults to training max
+                recycles, which is 3.
+        """
+
+        if mask is None:
+            mask = torch.ones_like(aa)
+
+        B = aa.shape[0]
+        L = aa.shape[1]
+        device = aa.device
+
+        if residx is None:
+            residx = torch.arange(L, device=device).expand_as(aa)
+
+        # === ESM ===
+        esmaa = self._af2_idx_to_esm_idx(aa, mask)
+
+        if masking_pattern is not None:
+            esmaa = self._mask_inputs_to_esm(esmaa, masking_pattern)
+
+        esm_s, esm_z = self._compute_language_model_representations(esmaa)
+
+        # Convert esm_s to the precision used by the trunk and
+        # the structure module. These tensors may be a lower precision if, for example,
+        # we're running the language model in fp16 precision.
+        esm_s = esm_s.to(self.esm_s_combine.dtype)
+        esm_s = esm_s.detach()
+
+        # === preprocessing ===
+        esm_s = (self.esm_s_combine.softmax(0).unsqueeze(0) @ esm_s).squeeze(2)
+
+        s_s_0 = self.esm_s_mlp(esm_s)
+        if self.cfg.use_esm_attn_map:
+            esm_z = esm_z.to(self.esm_s_combine.dtype)
+            esm_z = esm_z.detach()
+            s_z_0 = self.esm_z_mlp(esm_z)
+        else:
+            s_z_0 = s_s_0.new_zeros(B, L, L, self.cfg.trunk.pairwise_state_dim)
+
+        s_s_0 += self.embedding(aa)
+
+        structure: dict = self.trunk(
+            s_s_0, s_z_0, aa, residx, mask, no_recycles=num_recycles
+        )
+        # Documenting what we expect:
+        structure = {
+            k: v
+            for k, v in structure.items()
+            if k
+            in [
+                "s_z",
+                "s_s",
+                "frames",
+                "sidechain_frames",
+                "unnormalized_angles",
+                "angles",
+                "positions",
+                "states",
+            ]
+        }
+
+        disto_logits = self.distogram_head(structure["s_z"])
+        disto_logits = (disto_logits + disto_logits.transpose(1, 2)) / 2
+        structure["distogram_logits"] = disto_logits
+
+        lm_logits = self.lm_head(structure["s_s"])
+        structure["lm_logits"] = lm_logits
+
+        structure["aatype"] = aa
+        make_atom14_masks(structure)
+
+        for k in [
+            "atom14_atom_exists",
+            "atom37_atom_exists",
+        ]:
+            structure[k] *= mask.unsqueeze(-1)
+        structure["residue_index"] = residx
+
+        lddt_head = self.lddt_head(structure["states"]).reshape(
+            structure["states"].shape[0], B, L, -1, self.lddt_bins
+        )
+        structure["lddt_head"] = lddt_head
+        plddt = categorical_lddt(lddt_head[-1], bins=self.lddt_bins)
+        structure["plddt"] = (
+            100 * plddt
+        )  # we predict plDDT between 0 and 1, scale to be between 0 and 100.
+
+        ptm_logits = self.ptm_head(structure["s_z"])
+
+        seqlen = mask.type(torch.int64).sum(1)
+        structure["ptm_logits"] = ptm_logits
+        structure["ptm"] = torch.stack(
+            [
+                compute_tm(
+                    batch_ptm_logits[None, :sl, :sl],
+                    max_bins=31,
+                    no_bins=self.distogram_bins,
+                )
+                for batch_ptm_logits, sl in zip(ptm_logits, seqlen)
+            ]
+        )
+        structure.update(
+            compute_predicted_aligned_error(
+                ptm_logits, max_bin=31, no_bins=self.distogram_bins
+            )
+        )
+
+        return structure
+
+    @torch.no_grad()
+    def infer(
+        self,
+        sequences: T.Union[str, T.List[str]],
+        residx=None,
+        masking_pattern: T.Optional[torch.Tensor] = None,
+        num_recycles: T.Optional[int] = None,
+        residue_index_offset: T.Optional[int] = 512,
+        chain_linker: T.Optional[str] = "G" * 25,
+    ):
+        """Runs a forward pass given input sequences.
+
+        Args:
+            sequences (Union[str, List[str]]): A list of sequences to make predictions for. Multimers can also be passed in,
+                each chain should be separated by a ':' token (e.g. "<chain1>:<chain2>:<chain3>").
+            residx (torch.Tensor): Residue indices of amino acids. Will assume contiguous if not provided.
+            masking_pattern (torch.Tensor): Optional masking to pass to the input. Binary tensor of the same size
+                as `aa`. Positions with 1 will be masked. ESMFold sometimes produces different samples when
+                different masks are provided.
+            num_recycles (int): How many recycle iterations to perform. If None, defaults to training max
+                recycles (cfg.trunk.max_recycles), which is 4.
+            residue_index_offset (int): Residue index separation between chains if predicting a multimer. Has no effect on
+                single chain predictions. Default: 512.
+            chain_linker (str): Linker to use between chains if predicting a multimer. Has no effect on single chain
+                predictions. Default: length-25 poly-G ("G" * 25).
+        """
+        if isinstance(sequences, str):
+            sequences = [sequences]
+
+        aatype, mask, _residx, linker_mask, chain_index = batch_encode_sequences(
+            sequences, residue_index_offset, chain_linker
+        )
+
+        if residx is None:
+            residx = _residx
+        elif not isinstance(residx, torch.Tensor):
+            residx = collate_dense_tensors(residx)
+
+        aatype, mask, residx, linker_mask = map(
+            lambda x: x.to(self.device), (aatype, mask, residx, linker_mask)
+        )
+
+        output = self.forward(
+            aatype,
+            mask=mask,
+            residx=residx,
+            masking_pattern=masking_pattern,
+            num_recycles=num_recycles,
+        )
+
+        output["atom37_atom_exists"] = output[
+            "atom37_atom_exists"
+        ] * linker_mask.unsqueeze(2)
+
+        output["mean_plddt"] = (output["plddt"] * output["atom37_atom_exists"]).sum(
+            dim=(1, 2)
+        ) / output["atom37_atom_exists"].sum(dim=(1, 2))
+        output["chain_index"] = chain_index
+
+        return output
+
+    def output_to_pdb(self, output: T.Dict) -> T.List[str]:
+        """Returns the pbd (file) string from the model given the model output."""
+        return output_to_pdb(output)
+
+    def infer_pdbs(self, seqs: T.List[str], *args, **kwargs) -> T.List[str]:
+        """Returns list of pdb (files) strings from the model given a list of input sequences."""
+        output = self.infer(seqs, *args, **kwargs)
+        return self.output_to_pdb(output)
+
+    def infer_pdb(self, sequence: str, *args, **kwargs) -> str:
+        """Returns the pdb (file) string from the model given an input sequence."""
+        return self.infer_pdbs([sequence], *args, **kwargs)[0]
+
+    def set_chunk_size(self, chunk_size: T.Optional[int]):
+        # This parameter means the axial attention will be computed
+        # in a chunked manner. This should make the memory used more or less O(L) instead of O(L^2).
+        # It's equivalent to running a for loop over chunks of the dimension we're iterative over,
+        # where the chunk_size is the size of the chunks, so 128 would mean to parse 128-lengthed chunks.
+        # Setting the value to None will return to default behavior, disable chunking.
+        self.trunk.set_chunk_size(chunk_size)
+
+    @property
+    def device(self):
+        return self.esm_s_combine.device

models/esm/msa_transformer.py

@@ -0,0 +1,238 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+
+from onescience.modules.esm import (
+    AxialTransformerLayer,
+    LearnedPositionalEmbedding,
+    RobertaLMHead,
+    ESM1bLayerNorm,
+    ContactPredictionHead,
+)
+
+from onescience.modules.attention import RowSelfAttention, ColumnSelfAttention
+
+
+
+class MSATransformer(nn.Module):
+    @classmethod
+    def add_args(cls, parser):
+        # fmt: off
+        parser.add_argument(
+            "--num_layers",
+            default=12,
+            type=int,
+            metavar="N",
+            help="number of layers"
+        )
+        parser.add_argument(
+            "--embed_dim",
+            default=768,
+            type=int,
+            metavar="N",
+            help="embedding dimension"
+        )
+        parser.add_argument(
+            "--logit_bias",
+            action="store_true",
+            help="whether to apply bias to logits"
+        )
+        parser.add_argument(
+            "--ffn_embed_dim",
+            default=3072,
+            type=int,
+            metavar="N",
+            help="embedding dimension for FFN",
+        )
+        parser.add_argument(
+            "--attention_heads",
+            default=12,
+            type=int,
+            metavar="N",
+            help="number of attention heads",
+        )
+        parser.add_argument(
+            "--dropout",
+            default=0.1,
+            type=float,
+            help="Dropout to apply."
+        )
+        parser.add_argument(
+            "--attention_dropout",
+            default=0.1,
+            type=float,
+            help="Dropout to apply."
+        )
+        parser.add_argument(
+            "--activation_dropout",
+            default=0.1,
+            type=float,
+            help="Dropout to apply."
+        )
+        parser.add_argument(
+            "--max_tokens_per_msa",
+            default=2 ** 14,
+            type=int,
+            help=(
+                "Used during inference to batch attention computations in a single "
+                "forward pass. This allows increased input sizes with less memory."
+            ),
+        )
+        # fmt: on
+
+    def __init__(self, args, alphabet):
+        super().__init__()
+        self.args = args
+        self.alphabet_size = len(alphabet)
+        self.padding_idx = alphabet.padding_idx
+        self.mask_idx = alphabet.mask_idx
+        self.cls_idx = alphabet.cls_idx
+        self.eos_idx = alphabet.eos_idx
+        self.prepend_bos = alphabet.prepend_bos
+        self.append_eos = alphabet.append_eos
+
+        self.embed_tokens = nn.Embedding(
+            self.alphabet_size, self.args.embed_dim, padding_idx=self.padding_idx
+        )
+
+        if getattr(self.args, "embed_positions_msa", False):
+            emb_dim = getattr(self.args, "embed_positions_msa_dim", self.args.embed_dim)
+            self.msa_position_embedding = nn.Parameter(
+                0.01 * torch.randn(1, 1024, 1, emb_dim),
+                requires_grad=True,
+            )
+        else:
+            self.register_parameter("msa_position_embedding", None)
+
+        self.dropout_module = nn.Dropout(self.args.dropout)
+        self.layers = nn.ModuleList(
+            [
+                AxialTransformerLayer(
+                    self.args.embed_dim,
+                    self.args.ffn_embed_dim,
+                    self.args.attention_heads,
+                    self.args.dropout,
+                    self.args.attention_dropout,
+                    self.args.activation_dropout,
+                    getattr(self.args, "max_tokens_per_msa", self.args.max_tokens),
+                )
+                for _ in range(self.args.layers)
+            ]
+        )
+
+        self.contact_head = ContactPredictionHead(
+            self.args.layers * self.args.attention_heads,
+            self.prepend_bos,
+            self.append_eos,
+            eos_idx=self.eos_idx,
+        )
+        self.embed_positions = LearnedPositionalEmbedding(
+            self.args.max_positions,
+            self.args.embed_dim,
+            self.padding_idx,
+        )
+        self.emb_layer_norm_before = ESM1bLayerNorm(self.args.embed_dim)
+        self.emb_layer_norm_after = ESM1bLayerNorm(self.args.embed_dim)
+        self.lm_head = RobertaLMHead(
+            embed_dim=self.args.embed_dim,
+            output_dim=self.alphabet_size,
+            weight=self.embed_tokens.weight,
+        )
+
+    def forward(self, tokens, repr_layers=[], need_head_weights=False, return_contacts=False):
+        if return_contacts:
+            need_head_weights = True
+
+        assert tokens.ndim == 3
+        batch_size, num_alignments, seqlen = tokens.size()
+        padding_mask = tokens.eq(self.padding_idx)  # B, R, C
+        if not padding_mask.any():
+            padding_mask = None
+
+        x = self.embed_tokens(tokens)
+        x += self.embed_positions(tokens.view(batch_size * num_alignments, seqlen)).view(x.size())
+        if self.msa_position_embedding is not None:
+            if x.size(1) > 1024:
+                raise RuntimeError(
+                    "Using model with MSA position embedding trained on maximum MSA "
+                    f"depth of 1024, but received {x.size(1)} alignments."
+                )
+            x += self.msa_position_embedding[:, :num_alignments]
+
+        x = self.emb_layer_norm_before(x)
+
+        x = self.dropout_module(x)
+
+        if padding_mask is not None:
+            x = x * (1 - padding_mask.unsqueeze(-1).type_as(x))
+
+        repr_layers = set(repr_layers)
+        hidden_representations = {}
+        if 0 in repr_layers:
+            hidden_representations[0] = x
+
+        if need_head_weights:
+            row_attn_weights = []
+            col_attn_weights = []
+
+        # B x R x C x D -> R x C x B x D
+        x = x.permute(1, 2, 0, 3)
+
+        for layer_idx, layer in enumerate(self.layers):
+            x = layer(
+                x,
+                self_attn_padding_mask=padding_mask,
+                need_head_weights=need_head_weights,
+            )
+            if need_head_weights:
+                x, col_attn, row_attn = x
+                # H x C x B x R x R -> B x H x C x R x R
+                col_attn_weights.append(col_attn.permute(2, 0, 1, 3, 4))
+                # H x B x C x C -> B x H x C x C
+                row_attn_weights.append(row_attn.permute(1, 0, 2, 3))
+            if (layer_idx + 1) in repr_layers:
+                hidden_representations[layer_idx + 1] = x.permute(2, 0, 1, 3)
+
+        x = self.emb_layer_norm_after(x)
+        x = x.permute(2, 0, 1, 3)  # R x C x B x D -> B x R x C x D
+
+        # last hidden representation should have layer norm applied
+        if (layer_idx + 1) in repr_layers:
+            hidden_representations[layer_idx + 1] = x
+        x = self.lm_head(x)
+
+        result = {"logits": x, "representations": hidden_representations}
+        if need_head_weights:
+            # col_attentions: B x L x H x C x R x R
+            col_attentions = torch.stack(col_attn_weights, 1)
+            # row_attentions: B x L x H x C x C
+            row_attentions = torch.stack(row_attn_weights, 1)
+            result["col_attentions"] = col_attentions
+            result["row_attentions"] = row_attentions
+            if return_contacts:
+                contacts = self.contact_head(tokens, row_attentions)
+                result["contacts"] = contacts
+
+        return result
+
+    def predict_contacts(self, tokens):
+        return self(tokens, return_contacts=True)["contacts"]
+
+    @property
+    def num_layers(self):
+        return self.args.layers
+
+    def max_tokens_per_msa_(self, value: int) -> None:
+        """The MSA Transformer automatically batches attention computations when
+        gradients are disabled to allow you to pass in larger MSAs at test time than
+        you can fit in GPU memory. By default this occurs when more than 2^14 tokens
+        are passed in the input MSA. You can set this value to infinity to disable
+        this behavior.
+        """
+        for module in self.modules():
+            if isinstance(module, (RowSelfAttention, ColumnSelfAttention)):
+                module.max_tokens_per_msa = value

models/esm/pretrained.py

@@ -0,0 +1,545 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import re
+import urllib
+import warnings
+from argparse import Namespace
+from pathlib import Path
+
+import torch
+
+import onescience.models.esm as esm
+from onescience.datapipes.esm import Alphabet
+from models.esm.esm2 import ESM2
+
+
+def _has_regression_weights(model_name):
+    """Return whether we expect / require regression weights;
+    Right now that is all models except ESM-1v, ESM-IF, and partially trained ESM2 models"""
+    return not ("esm1v" in model_name or "esm_if" in model_name or "270K" in model_name or "500K" in model_name)
+
+
+def load_model_and_alphabet(model_name):
+    if model_name.endswith(".pt"):  # treat as filepath
+        return load_model_and_alphabet_local(model_name)
+    else:
+        return load_model_and_alphabet_hub(model_name)
+
+
+def load_hub_workaround(url):
+    raise FileNotFoundError(
+        "Remote ESM download is disabled in the packaged SimpleFold runtime. "
+        "Use load_model_and_alphabet_local() with a file under weight/esm_models."
+    )
+
+
+def load_regression_hub(model_name):
+    raise FileNotFoundError(
+        f"Remote ESM regression download is disabled: {model_name}. "
+        "Place the contact-regression file under weight/esm_models."
+    )
+
+
+def _download_model_and_regression_data(model_name):
+    raise FileNotFoundError(
+        f"Remote ESM download is disabled: {model_name}. "
+        "Place the required ESM files under weight/esm_models."
+    )
+
+
+def load_model_and_alphabet_hub(model_name):
+    raise FileNotFoundError(
+        f"Remote ESM download is disabled in the packaged SimpleFold runtime: {model_name}. "
+        "Place the required ESM .pt file and contact-regression file under weight/esm_models."
+    )
+
+
+def load_model_and_alphabet_local(model_location):
+    """Load from local path. The regression weights need to be co-located"""
+    model_location = Path(model_location)
+    model_data = torch.load(str(model_location), map_location="cpu")
+    model_name = model_location.stem
+    if _has_regression_weights(model_name):
+        regression_location = str(model_location.with_suffix("")) + "-contact-regression.pt"
+        regression_data = torch.load(regression_location, map_location="cpu")
+    else:
+        regression_data = None
+    return load_model_and_alphabet_core(model_name, model_data, regression_data)
+
+
+def has_emb_layer_norm_before(model_state):
+    """Determine whether layer norm needs to be applied before the encoder"""
+    return any(k.startswith("emb_layer_norm_before") for k, param in model_state.items())
+
+
+def _load_model_and_alphabet_core_v1(model_data):
+    import models.esm as esm  # since esm.inverse_folding is imported below, you actually have to re-import esm here
+
+    alphabet = esm.Alphabet.from_architecture(model_data["args"].arch)
+
+    if model_data["args"].arch == "roberta_large":
+        # upgrade state dict
+        pra = lambda s: "".join(s.split("encoder_")[1:] if "encoder" in s else s)
+        prs1 = lambda s: "".join(s.split("encoder.")[1:] if "encoder" in s else s)
+        prs2 = lambda s: "".join(
+            s.split("sentence_encoder.")[1:] if "sentence_encoder" in s else s
+        )
+        model_args = {pra(arg[0]): arg[1] for arg in vars(model_data["args"]).items()}
+        model_state = {prs1(prs2(arg[0])): arg[1] for arg in model_data["model"].items()}
+        model_state["embed_tokens.weight"][alphabet.mask_idx].zero_()  # For token drop
+        model_args["emb_layer_norm_before"] = has_emb_layer_norm_before(model_state)
+        model_type = esm.ProteinBertModel
+
+    elif model_data["args"].arch == "protein_bert_base":
+
+        # upgrade state dict
+        pra = lambda s: "".join(s.split("decoder_")[1:] if "decoder" in s else s)
+        prs = lambda s: "".join(s.split("decoder.")[1:] if "decoder" in s else s)
+        model_args = {pra(arg[0]): arg[1] for arg in vars(model_data["args"]).items()}
+        model_state = {prs(arg[0]): arg[1] for arg in model_data["model"].items()}
+        model_type = esm.ProteinBertModel
+    elif model_data["args"].arch == "msa_transformer":
+
+        # upgrade state dict
+        pra = lambda s: "".join(s.split("encoder_")[1:] if "encoder" in s else s)
+        prs1 = lambda s: "".join(s.split("encoder.")[1:] if "encoder" in s else s)
+        prs2 = lambda s: "".join(
+            s.split("sentence_encoder.")[1:] if "sentence_encoder" in s else s
+        )
+        prs3 = lambda s: s.replace("row", "column") if "row" in s else s.replace("column", "row")
+        model_args = {pra(arg[0]): arg[1] for arg in vars(model_data["args"]).items()}
+        model_state = {prs1(prs2(prs3(arg[0]))): arg[1] for arg in model_data["model"].items()}
+        if model_args.get("embed_positions_msa", False):
+            emb_dim = model_state["msa_position_embedding"].size(-1)
+            model_args["embed_positions_msa_dim"] = emb_dim  # initial release, bug: emb_dim==1
+
+        model_type = esm.MSATransformer
+
+    elif "invariant_gvp" in model_data["args"].arch:
+        import models.esm.inverse_folding
+
+        model_type = esm.inverse_folding.gvp_transformer.GVPTransformerModel
+        model_args = vars(model_data["args"])  # convert Namespace -> dict
+
+        def update_name(s):
+            # Map the module names in checkpoints trained with internal code to
+            # the updated module names in open source code
+            s = s.replace("W_v", "embed_graph.embed_node")
+            s = s.replace("W_e", "embed_graph.embed_edge")
+            s = s.replace("embed_scores.0", "embed_confidence")
+            s = s.replace("embed_score.", "embed_graph.embed_confidence.")
+            s = s.replace("seq_logits_projection.", "")
+            s = s.replace("embed_ingraham_features", "embed_dihedrals")
+            s = s.replace("embed_gvp_in_local_frame.0", "embed_gvp_output")
+            s = s.replace("embed_features_in_local_frame.0", "embed_gvp_input_features")
+            return s
+
+        model_state = {
+            update_name(sname): svalue
+            for sname, svalue in model_data["model"].items()
+            if "version" not in sname
+        }
+
+    else:
+        raise ValueError("Unknown architecture selected")
+
+    model = model_type(
+        Namespace(**model_args),
+        alphabet,
+    )
+
+    return model, alphabet, model_state
+
+
+def _load_model_and_alphabet_core_v2(model_data):
+    def upgrade_state_dict(state_dict):
+        """Removes prefixes 'model.encoder.sentence_encoder.' and 'model.encoder.'."""
+        prefixes = ["encoder.sentence_encoder.", "encoder."]
+        pattern = re.compile("^" + "|".join(prefixes))
+        state_dict = {pattern.sub("", name): param for name, param in state_dict.items()}
+        return state_dict
+
+    cfg = model_data["cfg"]["model"]
+    state_dict = model_data["model"]
+    state_dict = upgrade_state_dict(state_dict)
+    alphabet = Alphabet.from_architecture("ESM-1b")
+    model = ESM2(
+        num_layers=cfg.encoder_layers,
+        embed_dim=cfg.encoder_embed_dim,
+        attention_heads=cfg.encoder_attention_heads,
+        alphabet=alphabet,
+        token_dropout=cfg.token_dropout,
+    )
+    return model, alphabet, state_dict
+
+
+def load_model_and_alphabet_core(model_name, model_data, regression_data=None):
+    if regression_data is not None:
+        model_data["model"].update(regression_data["model"])
+
+    if model_name.startswith("esm2"):
+        model, alphabet, model_state = _load_model_and_alphabet_core_v2(model_data)
+    else:
+        model, alphabet, model_state = _load_model_and_alphabet_core_v1(model_data)
+
+    expected_keys = set(model.state_dict().keys())
+    found_keys = set(model_state.keys())
+
+    if regression_data is None:
+        expected_missing = {"contact_head.regression.weight", "contact_head.regression.bias"}
+        error_msgs = []
+        missing = (expected_keys - found_keys) - expected_missing
+        if missing:
+            error_msgs.append(f"Missing key(s) in state_dict: {missing}.")
+        unexpected = found_keys - expected_keys
+        if unexpected:
+            error_msgs.append(f"Unexpected key(s) in state_dict: {unexpected}.")
+
+        if error_msgs:
+            raise RuntimeError(
+                "Error(s) in loading state_dict for {}:\n\t{}".format(
+                    model.__class__.__name__, "\n\t".join(error_msgs)
+                )
+            )
+        if expected_missing - found_keys:
+            warnings.warn(
+                "Regression weights not found, predicting contacts will not produce correct results."
+            )
+
+    model.load_state_dict(model_state, strict=regression_data is not None)
+
+    return model, alphabet
+
+
+def esm1_t34_670M_UR50S():
+    """34 layer transformer model with 670M params, trained on Uniref50 Sparse.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1_t34_670M_UR50S")
+
+
+def esm1_t34_670M_UR50D():
+    """34 layer transformer model with 670M params, trained on Uniref50 Dense.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1_t34_670M_UR50D")
+
+
+def esm1_t34_670M_UR100():
+    """34 layer transformer model with 670M params, trained on Uniref100.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1_t34_670M_UR100")
+
+
+def esm1_t12_85M_UR50S():
+    """12 layer transformer model with 85M params, trained on Uniref50 Sparse.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1_t12_85M_UR50S")
+
+
+def esm1_t6_43M_UR50S():
+    """6 layer transformer model with 43M params, trained on Uniref50 Sparse.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1_t6_43M_UR50S")
+
+
+def esm1b_t33_650M_UR50S():
+    """33 layer transformer model with 650M params, trained on Uniref50 Sparse.
+    This is our best performing model, which will be described in a future publication.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1b_t33_650M_UR50S")
+
+
+def esm_msa1_t12_100M_UR50S():
+    warnings.warn(
+        "This model had a minor bug in the positional embeddings, "
+        "please use ESM-MSA-1b: esm.pretrained.esm_msa1b_t12_100M_UR50S()",
+    )
+    return load_model_and_alphabet_hub("esm_msa1_t12_100M_UR50S")
+
+
+def esm_msa1b_t12_100M_UR50S():
+    return load_model_and_alphabet_hub("esm_msa1b_t12_100M_UR50S")
+
+
+def esm1v_t33_650M_UR90S():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 1 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_1")
+
+
+def esm1v_t33_650M_UR90S_1():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 1 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_1")
+
+
+def esm1v_t33_650M_UR90S_2():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 2 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_2")
+
+
+def esm1v_t33_650M_UR90S_3():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 3 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_3")
+
+
+def esm1v_t33_650M_UR90S_4():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 4 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_4")
+
+
+def esm1v_t33_650M_UR90S_5():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 5 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_5")
+
+
+def esm_if1_gvp4_t16_142M_UR50():
+    """Inverse folding model with 142M params, with 4 GVP-GNN layers, 8
+    Transformer encoder layers, and 8 Transformer decoder layers, trained on
+    CATH structures and 12 million alphafold2 predicted structures from UniRef50
+    sequences.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm_if1_gvp4_t16_142M_UR50")
+
+
+def esm2_t6_8M_UR50D():
+    """6 layer ESM-2 model with 8M params, trained on UniRef50.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t6_8M_UR50D")
+
+
+def esm2_t12_35M_UR50D():
+    """12 layer ESM-2 model with 35M params, trained on UniRef50.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t12_35M_UR50D")
+
+
+def esm2_t30_150M_UR50D():
+    """30 layer ESM-2 model with 150M params, trained on UniRef50.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t30_150M_UR50D")
+
+
+def esm2_t33_650M_UR50D():
+    """33 layer ESM-2 model with 650M params, trained on UniRef50.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t33_650M_UR50D")
+
+
+def esm2_t36_3B_UR50D():
+    """36 layer ESM-2 model with 3B params, trained on UniRef50.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t36_3B_UR50D")
+
+
+def esm2_t48_15B_UR50D():
+    """48 layer ESM-2 model with 15B params, trained on UniRef50.
+    If you have OOM while loading this model, please refer to README
+    on how to employ FSDP and ZeRO CPU offloading
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t48_15B_UR50D")
+
+
+def esmfold_v0():
+    """
+    ESMFold v0 model with 3B ESM-2, 48 folding blocks.
+    This version was used for the paper (Lin et al, 2022). It was trained 
+    on all PDB chains until 2020-05, to ensure temporal holdout with CASP14
+    and the CAMEO validation and test set reported there.
+    """
+    import models.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_v0()
+
+
+def esmfold_v1():
+    """
+    ESMFold v1 model using 3B ESM-2, 48 folding blocks.
+    ESMFold provides fast high accuracy atomic level structure prediction
+    directly from the individual sequence of a protein. ESMFold uses the ESM2
+    protein language model to extract meaningful representations from the
+    protein sequence.
+    """
+    import models.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_v1()
+
+def esmfold_structure_module_only_8M():
+    """
+    ESMFold baseline model using 8M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import models.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_8M()
+
+
+def esmfold_structure_module_only_8M_270K():
+    """
+    ESMFold baseline model using 8M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import models.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_8M_270K()
+
+
+def esmfold_structure_module_only_35M():
+    """
+    ESMFold baseline model using 35M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import models.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_35M()
+
+
+def esmfold_structure_module_only_35M_270K():
+    """
+    ESMFold baseline model using 35M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import models.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_35M_270K()
+
+
+def esmfold_structure_module_only_150M():
+    """
+    ESMFold baseline model using 150M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import models.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_150M()
+
+
+def esmfold_structure_module_only_150M_270K():
+    """
+    ESMFold baseline model using 150M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import models.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_150M_270K()
+
+
+def esmfold_structure_module_only_650M():
+    """
+    ESMFold baseline model using 650M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import models.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_650M()
+
+
+def esmfold_structure_module_only_650M_270K():
+    """
+    ESMFold baseline model using 650M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import models.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_650M_270K()
+
+
+def esmfold_structure_module_only_3B():
+    """
+    ESMFold baseline model using 3B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import models.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_3B()
+
+
+def esmfold_structure_module_only_3B_270K():
+    """
+    ESMFold baseline model using 3B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import models.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_3B_270K()
+
+
+def esmfold_structure_module_only_15B():
+    """
+    ESMFold baseline model using 15B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    The 15B parameter ESM-2 was not trained out to 500K updates
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import models.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_15B()

models/esm/version.py

@@ -0,0 +1,6 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+version = "2.0.1"

models/simplefold/torch/blocks.py

@@ -0,0 +1,146 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import torch
+from torch import nn
+from .layers import modulate, SwiGLUFeedForward
+
+try:
+    from timm.models.vision_transformer import Mlp
+except ModuleNotFoundError:
+    class Mlp(nn.Module):
+        def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0):
+            super().__init__()
+            out_features = out_features or in_features
+            hidden_features = hidden_features or in_features
+            self.fc1 = nn.Linear(in_features, hidden_features)
+            self.act = act_layer()
+            self.drop1 = nn.Dropout(drop)
+            self.fc2 = nn.Linear(hidden_features, out_features)
+            self.drop2 = nn.Dropout(drop)
+
+        def forward(self, x):
+            x = self.fc1(x)
+            x = self.act(x)
+            x = self.drop1(x)
+            x = self.fc2(x)
+            x = self.drop2(x)
+            return x
+
+
+class DiTBlock(nn.Module):
+    """
+    A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+    """
+
+    def __init__(
+        self,
+        self_attention_layer,
+        hidden_size,
+        mlp_ratio=4.0,
+        use_swiglu=True,
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = self_attention_layer()
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        if use_swiglu:
+            self.mlp = SwiGLUFeedForward(hidden_size, mlp_hidden_dim)
+        else:
+            approx_gelu = lambda: nn.GELU(approximate="tanh")
+            self.mlp = Mlp(
+                in_features=hidden_size,
+                hidden_features=mlp_hidden_dim,
+                act_layer=approx_gelu,
+                drop=0,
+            )
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+        self.apply(_basic_init)
+
+        # Zero-out adaLN modulation layers in DiT encoder blocks:
+        nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
+
+    def forward(
+        self,
+        latents,
+        c,
+        **kwargs,
+    ):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            self.adaLN_modulation(c).chunk(6, dim=1)
+        )
+        _latents = self.attn(
+            modulate(self.norm1(latents), shift_msa, scale_msa), **kwargs
+        )
+        latents = latents + gate_msa.unsqueeze(1) * _latents
+        latents = latents + gate_mlp.unsqueeze(1) * self.mlp(
+            modulate(self.norm2(latents), shift_mlp, scale_mlp)
+        )
+        return latents
+
+
+class TransformerBlock(nn.Module):
+    """
+    A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+    """
+
+    def __init__(
+        self,
+        self_attention_layer,
+        hidden_size,
+        mlp_ratio=4.0,
+        use_swiglu=False,
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = self_attention_layer()
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        if use_swiglu:
+            self.mlp = SwiGLUFeedForward(hidden_size, mlp_hidden_dim)
+        else:
+            self.mlp = Mlp(
+                in_features=hidden_size,
+                hidden_features=mlp_hidden_dim,
+                act_layer=approx_gelu,
+                drop=0,
+            )
+
+    def forward(
+        self,
+        latents,
+        **kwargs,
+    ):
+        _latents = self.attn(self.norm1(latents), **kwargs)
+        latents = latents + _latents
+        latents = latents + self.mlp(self.norm2(latents))
+        return latents
+
+
+class HomogenTrunk(nn.Module):
+    def __init__(self, block, depth):
+        super().__init__()
+        self.blocks = nn.ModuleList([block() for _ in range(depth)])
+
+    def forward(self, latents, c, **kwargs):
+        for i, block in enumerate(self.blocks):
+            kwargs["layer_idx"] = i
+            latents = block(latents=latents, c=c, **kwargs)
+        return latents

models/simplefold/torch/confidence_module.py

@@ -0,0 +1,85 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import torch
+from torch import nn
+
+
+def compute_aggregated_metric(logits, end=1.0):
+    """Compute the metric from the logits.
+
+    Parameters
+    ----------
+    logits : torch.Tensor
+        The logits of the metric
+    end : float
+        Max value of the metric, by default 1.0
+
+    Returns
+    -------
+    Tensor
+        The metric value
+
+    """
+    num_bins = logits.shape[-1]
+    bin_width = end / num_bins
+    bounds = torch.arange(
+        start=0.5 * bin_width, end=end, step=bin_width, device=logits.device
+    )
+    probs = nn.functional.softmax(logits, dim=-1)
+    plddt = torch.sum(
+        probs * bounds.view(*((1,) * len(probs.shape[:-1])), *bounds.shape),
+        dim=-1,
+    )
+    return plddt
+
+
+class ConfidenceModule(nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        transformer_blocks,
+        num_plddt_bins=50,
+    ):
+        super().__init__()
+        self.transformer_blocks = transformer_blocks
+        self.to_plddt_logits = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.LayerNorm(hidden_size),
+            nn.SiLU(),
+            nn.Linear(hidden_size, num_plddt_bins),
+        )
+
+    def forward(
+        self,
+        latent,
+        feats,
+    ):
+        token_pe_pos = torch.cat(
+            [
+                feats["residue_index"].unsqueeze(-1).float(),        # (B, M, 1)
+                feats["entity_id"].unsqueeze(-1).float(),            # (B, M, 1)
+                feats["asym_id"].unsqueeze(-1).float(),              # (B, M, 1)
+                feats["sym_id"].unsqueeze(-1).float(),               # (B, M, 1)
+            ],
+            dim=-1,
+        )  
+
+        latent = self.transformer_blocks(
+            latents=latent, 
+            c=None,
+            pos=token_pe_pos,
+        )
+
+        # Compute the pLDDT
+        plddt_logits = self.to_plddt_logits(latent)
+
+        # Compute the aggregated pLDDT
+        plddt = compute_aggregated_metric(plddt_logits)
+
+        return dict(
+            plddt=plddt,
+            plddt_logits=plddt_logits,
+        )

models/simplefold/torch/sampler.py

@@ -0,0 +1,108 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import torch
+from tqdm import tqdm
+from einops import repeat
+from onescience.utils.simplefold.boltz_utils import center_random_augmentation
+
+
+class EMSampler():
+    """
+    A Euler-Maruyama solver for SDEs.
+    """
+    def __init__(
+        self,
+        num_timesteps=500,
+        t_start=1e-4,
+        tau=0.3,
+        log_timesteps=False,
+        w_cutoff=0.99,
+    ):
+        self.num_timesteps = num_timesteps
+        self.log_timesteps = log_timesteps
+        self.t_start = t_start
+        self.tau = tau
+        self.w_cutoff = w_cutoff
+
+        if self.log_timesteps:
+            t = 1.0 - torch.logspace(-2, 0, self.num_timesteps + 1).flip(0)
+            t = t - torch.min(t)
+            t = t / torch.max(t)
+            self.steps = t.clamp(min=self.t_start, max=1.0)
+        else:
+            self.steps = torch.linspace(
+                self.t_start, 1.0, steps=self.num_timesteps + 1
+            )
+
+    def diffusion_coefficient(self, t, eps=0.01):
+        # determine diffusion coefficient
+        w = (1.0 - t) / (t + eps)
+        if t >= self.w_cutoff:
+            w = 0.0
+        return w
+
+    @torch.no_grad()
+    def euler_maruyama_step(
+        self,
+        model_fn,
+        flow,
+        y, 
+        t, 
+        t_next, 
+        batch, 
+    ):
+        dt = t_next - t
+        eps = torch.randn_like(y).to(y)
+
+        y = center_random_augmentation(
+            y,
+            batch["atom_pad_mask"],
+            augmentation=False,
+            centering=True,
+        )
+
+        batched_t = repeat(t, " -> b", b=y.shape[0])
+        velocity = model_fn(
+            noised_pos=y,
+            t=batched_t,
+            feats=batch,
+        )['predict_velocity']
+        score = flow.compute_score_from_velocity(velocity, y, t)
+
+        diff_coeff = self.diffusion_coefficient(t)
+        drift = velocity + diff_coeff * score
+        mean_y = y + drift * dt
+        y_sample = mean_y + torch.sqrt(2.0 * dt * diff_coeff * self.tau) * eps
+
+        return y_sample
+
+    @torch.no_grad()
+    def sample(self, model_fn, flow, noise, batch):
+        sampling_timesteps = self.num_timesteps
+        steps = self.steps.to(noise.device)
+        y_sampled = noise
+        feats = batch
+
+        for i in tqdm(
+            range(sampling_timesteps),
+            desc="Sampling",
+            total=sampling_timesteps,
+        ):
+            t = steps[i]
+            t_next = steps[i + 1]
+
+            y_sampled = self.euler_maruyama_step(
+                model_fn,
+                flow,
+                y_sampled,
+                t,
+                t_next,
+                feats,
+            )
+
+        return {
+            "denoised_coords": y_sampled
+        }

weight/ccd.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2d3b2f03a3c5665944adba51e33263511e51b21c9cd05d902f9c4b7c1e58d2f4
+size 345859128