Upload folder using huggingface_hub

.hydra/config.yaml

@@ -0,0 +1,417 @@
+task_name: train
+experiment_group: openfold_fs50_ur90_memmap
+tags:
+- ${experiment_group}
+train: true
+test: false
+ckpt_path: null
+seed: 12345
+float32_matmul_precision: high
+model:
+  _target_: src.models.llama.LlamaLitModule
+  scheduler_name: constant_with_warmup
+  num_warmup_steps: 200
+  num_training_steps: 1000000
+  lr: 0.001
+  embed_coords: false
+  embed_sequence_index: false
+  embed_residue_index: false
+  use_kv_cache_for_scoring: true
+  max_seq_pos_in_doc: 1024
+  max_res_pos_in_seq: 4096
+  pass_res_pos_in_doc_as_position_ids: true
+  optimizer: adamw
+  config:
+    _target_: transformers.LlamaConfig
+    vocab_size: ${constants.vocab_size}
+    hidden_size: 1024
+    intermediate_size: 4096
+    num_attention_heads: 16
+    num_hidden_layers: 16
+    num_key_value_heads: 8
+    rope_theta: 500000
+    max_position_embeddings: 131072
+    scoring_max_tokens: 10240
+    attn_implementation: flash_attention_2
+    attention_bias: false
+    attention_dropout: 0.0
+    rms_norm_eps: 1.0e-05
+    hidden_act: silu
+    torch_dtype: bfloat16
+    use_cache: true
+    pretraining_tp: 1
+    rope_scaling:
+      factor: 32.0
+      high_freq_factor: 4.0
+      low_freq_factor: 1.0
+      original_max_position_embeddings: 32000
+      rope_type: llama3
+callbacks:
+  throughput:
+    _target_: src.utils.callbacks.TokenThroughputMonitor
+  model_checkpoint:
+    _target_: lightning.pytorch.callbacks.ModelCheckpoint
+    dirpath: ${paths.output_dir}/checkpoints
+    filename: epoch_{epoch:03d}
+    monitor: val/loss
+    verbose: false
+    save_last: true
+    save_top_k: 1
+    mode: min
+    auto_insert_metric_name: false
+    save_weights_only: false
+    every_n_train_steps: null
+    train_time_interval: null
+    every_n_epochs: null
+    save_on_train_epoch_end: null
+  model_summary:
+    _target_: lightning.pytorch.callbacks.RichModelSummary
+    max_depth: -1
+  rich_progress_bar:
+    _target_: lightning.pytorch.callbacks.RichProgressBar
+  timer:
+    _target_: src.utils.callbacks.EpochTimerCallback
+  print:
+    _target_: src.utils.callbacks.PrintCallback
+  sample_counter:
+    _target_: src.utils.callbacks.SampleCounter
+logger:
+  wandb:
+    _target_: src.utils.loggers.WandbLogger
+    save_dir: ${paths.output_dir}
+    offline: false
+    id: null
+    anonymous: null
+    project: profam
+    log_model: false
+    prefix: ''
+    entity: ProFam
+    group: ''
+    name: null
+    tags: ${tags}
+    job_type: ''
+    log_hydra_config_file: true
+    log_git_hash: true
+trainer:
+  _target_: src.utils.trainer.ProFamTrainer
+  default_root_dir: ${paths.output_dir}
+  max_epochs: 10000
+  max_steps: -1
+  accelerator: gpu
+  devices: auto
+  check_val_every_n_epoch: 1
+  val_check_interval: 50000
+  target_tokens_per_batch: null
+  tokens_per_document: 30000
+  batch_size: ${data.batch_size}
+  deterministic: false
+  log_every_n_steps: 10
+  timeout: 120
+  profiler:
+    name: null
+    log_tensorboard: false
+    simple:
+      _target_: SimpleProfiler
+    advance:
+      _target_: AdvancedProfiler
+      filename: advanced_perf_logs
+      dirpath: ./profiler_logs
+    pytorch:
+      _target_: PyTorchProfiler
+      filename: pytorch_perf_logs
+      dirpath: ./profiler_logs
+      record_shapes: true
+      profile_memory: true
+      with_stack: true
+      with_flops: false
+      with_modules: false
+      acc_events: false
+  strategy: ddp
+  num_nodes: 1
+  sync_batchnorm: true
+  precision: bf16-true
+  min_epochs: 1000
+  accumulate_grad_batches: 2
+  use_distributed_sampler: false
+paths:
+  root_dir: ${oc.env:PROJECT_ROOT}
+  data_dir: /home/jovyan/shared/judewells/profam/data
+  log_dir: ${paths.root_dir}/logs/${experiment_group}
+  output_dir: ${hydra:runtime.output_dir}
+  work_dir: ${hydra:runtime.cwd}
+extras:
+  ignore_warnings: false
+  enforce_tags: true
+  print_config: true
+tokenizer:
+  _target_: src.data.tokenizers.ProFamTokenizer
+  tokenizer_file: data/profam_tokenizer.json
+  unk_token: '[UNK]'
+  pad_token: '[PAD]'
+  bos_token: '[start-of-document]'
+  sep_token: '[SEP]'
+  mask_token: '?'
+  seq_struct_sep_token: '|'
+  add_final_sep: true
+  add_bos_token: true
+  add_document_token: true
+  mask_below_plddt: 80
+  max_res_pos_in_seq: ${model.max_res_pos_in_seq}
+  embed_residue_index: ${model.embed_residue_index}
+constants:
+  vocab_size: 68
+  gym_val_assay_list:
+  - BLAT_ECOLX_Jacquier_2013
+  - CALM1_HUMAN_Weile_2017
+  - DYR_ECOLI_Thompson_2019
+  - DLG4_RAT_McLaughlin_2012
+  - REV_HV1H2_Fernandes_2016
+  - TAT_HV1BR_Fernandes_2016
+  - RL40A_YEAST_Roscoe_2013
+  - P53_HUMAN_Giacomelli_2018_WT_Nutlin
+  sequence_features:
+  - ds_name
+  - identifier
+  - input_ids
+  - attention_mask
+  - original_size
+  - residue_index
+  - batch_size
+  structure_features:
+  - ds_name
+  - identifier
+  - input_ids
+  - attention_mask
+  - original_size
+  - residue_index
+  - coords
+  - coords_mask
+  - interleaved_coords_mask
+  - aa_mask
+  - plddts
+  - structure_mask
+extra_callbacks:
+  unconditional_sampling_callback:
+    prompt_builder:
+      preprocessor:
+        cfg:
+          _target_: src.data.processors.PreprocessingConfig
+          document_token: '[RAW]'
+          drop_first_protein: false
+          keep_first_protein: false
+          allow_unk: false
+          max_tokens_per_example: 8192
+          shuffle_proteins_in_document: true
+          padding: do_not_pad
+        transform_fns:
+        - _target_: src.data.processors.transforms.replace_nans_in_coords
+          _partial_: true
+          fill_value: 0.0
+        _target_: src.data.processors.ProteinDocumentPreprocessor
+      _target_: src.models.inference.PromptBuilder
+    _target_: src.pipelines.callback.SamplingEvaluationPipelineCallback
+    pipeline:
+      _target_: src.pipelines.unconditional_sequence.UnconditionalSequenceEvaluationPipeline
+      num_generations: 5
+      max_tokens: 20000
+      max_generated_length: 300
+      pipeline_id: unconditional_sampling
+      save_results_to_file: false
+    evaluators:
+      _target_: src.evaluators.esmfold.ESMFoldSamplingEvaluator
+      name: esmfold_example
+data:
+  _target_: src.data.datamodule.ProteinDataMixture
+  dataset_builders:
+    openfold_train:
+      _target_: src.data.builders.family_text_memmap_datasets.ProteinFamilyMemmapDatasetBuilder
+      name: openfold_train
+      dataset_root: ${paths.data_dir}/openfold/uniclust30_clustered_shuffled_final_text/train_test_split_v2/train_filtered
+      tokenizer: ${tokenizer}
+      preprocessor:
+        _target_: src.data.processors.ProteinDocumentPreprocessor
+        cfg:
+          _target_: src.data.processors.AlignedProteinPreprocessingConfig
+          document_token: '[RAW]'
+          drop_first_protein: false
+          keep_first_protein: false
+          allow_unk: false
+          max_tokens_per_example: 8192
+          shuffle_proteins_in_document: true
+          padding: do_not_pad
+          keep_gaps: false
+          keep_insertions: true
+          to_upper: true
+          use_msa_pos: false
+        transform_fns:
+        - _target_: src.data.processors.transforms.replace_nans_in_coords
+          _partial_: true
+          fill_value: 0.0
+    proteingym:
+      _target_: src.data.builders.proteingym.ProteinGymDataset
+      name: proteingym
+      dms_ids: ${constants.gym_val_assay_list}
+      seed: 42
+      max_mutated_sequences: null
+      mutant_bos_token: sep
+      keep_gaps: false
+      use_filtered_msa: true
+      extra_tokens_per_document: 2
+      use_msa_pos: false
+      num_proc: null
+      max_tokens_per_example: 7500
+      max_context_seqs: null
+      keep_wt: false
+      drop_wt: true
+    foldseek_s50_train:
+      _target_: src.data.builders.family_text_memmap_datasets.ProteinFamilyMemmapDatasetBuilder
+      name: foldseek_s50_train
+      dataset_root: ${paths.data_dir}/foldseek/foldseek_s50_seq_only_text/train_test_split_v2/train_filtered
+      tokenizer: ${tokenizer}
+      seed: ${seed}
+      preprocessor:
+        _target_: src.data.processors.ProteinDocumentPreprocessor
+        cfg:
+          _target_: src.data.processors.PreprocessingConfig
+          document_token: '[RAW]'
+          drop_first_protein: false
+          keep_first_protein: false
+          allow_unk: false
+          max_tokens_per_example: 8192
+          shuffle_proteins_in_document: true
+          padding: do_not_pad
+        transform_fns:
+        - _target_: src.data.processors.transforms.replace_nans_in_coords
+          _partial_: true
+          fill_value: 0.0
+    uniref90_train:
+      _target_: src.data.builders.family_text_memmap_datasets.ProteinFamilyMemmapDatasetBuilder
+      name: uniref90_train
+      dataset_root: ${paths.data_dir}/uniref/uniref90_text_shuffled/train_test_split_v2/train_filtered
+      tokenizer: ${tokenizer}
+      preprocessor:
+        _target_: src.data.processors.ProteinDocumentPreprocessor
+        cfg:
+          _target_: src.data.processors.PreprocessingConfig
+          document_token: '[RAW]'
+          drop_first_protein: false
+          keep_first_protein: false
+          allow_unk: false
+          max_tokens_per_example: 320
+          shuffle_proteins_in_document: true
+          padding: do_not_pad
+        transform_fns:
+        - _target_: src.data.processors.transforms.replace_nans_in_coords
+          _partial_: true
+          fill_value: 0.0
+    uniref90_val:
+      _target_: src.data.builders.family_text_memmap_datasets.ProteinFamilyMemmapDatasetBuilder
+      name: uniref90_val
+      dataset_root: ${paths.data_dir}/uniref/uniref90_text_shuffled/train_test_split_v2/val_filtered
+      tokenizer: ${tokenizer}
+      preprocessor:
+        _target_: src.data.processors.ProteinDocumentPreprocessor
+        cfg:
+          _target_: src.data.processors.PreprocessingConfig
+          document_token: '[RAW]'
+          drop_first_protein: false
+          keep_first_protein: false
+          allow_unk: false
+          max_tokens_per_example: 320
+          shuffle_proteins_in_document: true
+          padding: do_not_pad
+        transform_fns:
+        - _target_: src.data.processors.transforms.replace_nans_in_coords
+          _partial_: true
+          fill_value: 0.0
+    funfams_s50_train:
+      _target_: src.data.builders.family_text_memmap_datasets.ProteinFamilyMemmapDatasetBuilder
+      name: funfams_s50_train
+      dataset_root: ${paths.data_dir}/funfams/s50_text/train_test_split_v2/train_filtered
+      tokenizer: ${tokenizer}
+      seed: ${seed}
+      preprocessor:
+        _target_: src.data.processors.ProteinDocumentPreprocessor
+        cfg:
+          _target_: src.data.processors.AlignedProteinPreprocessingConfig
+          document_token: '[RAW]'
+          drop_first_protein: false
+          keep_first_protein: false
+          allow_unk: false
+          max_tokens_per_example: 8192
+          shuffle_proteins_in_document: true
+          padding: do_not_pad
+          keep_gaps: false
+          keep_insertions: true
+          to_upper: true
+          use_msa_pos: false
+        transform_fns:
+        - _target_: src.data.processors.transforms.replace_nans_in_coords
+          _partial_: true
+          fill_value: 0.0
+    funfams_s50_val:
+      _target_: src.data.builders.family_text_memmap_datasets.ProteinFamilyMemmapDatasetBuilder
+      name: funfams_s50_val
+      dataset_root: ${paths.data_dir}/funfams/s50_text/train_test_split_v2/val_filtered
+      tokenizer: ${tokenizer}
+      preprocessor:
+        _target_: src.data.processors.ProteinDocumentPreprocessor
+        cfg:
+          _target_: src.data.processors.AlignedProteinPreprocessingConfig
+          document_token: '[RAW]'
+          drop_first_protein: false
+          keep_first_protein: false
+          allow_unk: false
+          max_tokens_per_example: 8192
+          shuffle_proteins_in_document: true
+          padding: do_not_pad
+          keep_gaps: false
+          keep_insertions: true
+          to_upper: true
+          use_msa_pos: false
+        transform_fns:
+        - _target_: src.data.processors.transforms.replace_nans_in_coords
+          _partial_: true
+          fill_value: 0.0
+    foldseek_s50_val:
+      _target_: src.data.builders.family_text_memmap_datasets.ProteinFamilyMemmapDatasetBuilder
+      name: foldseek_s50_val
+      dataset_root: ${paths.data_dir}/foldseek/foldseek_s50_seq_only_text/train_test_split_v2/val_filtered
+      tokenizer: ${tokenizer}
+      seed: ${seed}
+      preprocessor:
+        _target_: src.data.processors.ProteinDocumentPreprocessor
+        cfg:
+          _target_: src.data.processors.PreprocessingConfig
+          document_token: '[RAW]'
+          drop_first_protein: false
+          keep_first_protein: false
+          allow_unk: false
+          max_tokens_per_example: 8192
+          shuffle_proteins_in_document: true
+          padding: do_not_pad
+        transform_fns:
+        - _target_: src.data.processors.transforms.replace_nans_in_coords
+          _partial_: true
+          fill_value: 0.0
+  data_weights:
+    foldseek_s50_train: 1
+    uniref90_train: 1
+    openfold_train: 1
+    funfams_s50_train: 0.03
+  val_dataset_batch_sizes:
+    funfams_s50_val: 1
+    proteingym: 1
+    foldseek_s50_val: 1
+    uniref90_val: 1
+  batch_size: 100
+  data_dir: ${paths.data_dir}
+  num_workers: 32
+  ignore_gaps: true
+  feature_names: ${constants.sequence_features}
+  pack_to_max_tokens: 52000
+  prefetch_factor: 4
+  shuffle: true
+  interleaved: true
+  interleaved_block_size: 1000
+  total_num_train_samples: null

.hydra/gym_config.yaml

@@ -0,0 +1,242 @@
+task_name: train
+experiment_group: openfold_fs50_ur90_memmap_GYM_ONLY
+tags:
+- ${experiment_group}
+train: true
+test: false
+ckpt_path: null
+seed: 12345
+float32_matmul_precision: high
+model:
+  _target_: src.models.llama.LlamaLitModule
+  scheduler_name: constant_with_warmup
+  num_warmup_steps: 200
+  num_training_steps: 1000000
+  lr: 0.001
+  embed_coords: false
+  embed_sequence_index: false
+  embed_residue_index: false
+  use_kv_cache_for_scoring: true
+  max_seq_pos_in_doc: 1024
+  max_res_pos_in_seq: 4096
+  pass_res_pos_in_doc_as_position_ids: true
+  optimizer: adamw
+  scoring_max_tokens: 10240
+  config:
+    _target_: transformers.LlamaConfig
+    vocab_size: ${constants.vocab_size}
+    hidden_size: 1024
+    intermediate_size: 4096
+    num_attention_heads: 16
+    num_hidden_layers: 16
+    num_key_value_heads: 8
+    rope_theta: 500000
+    max_position_embeddings: 131072
+    scoring_max_tokens: 10240
+    attn_implementation: flash_attention_2
+    attention_bias: false
+    attention_dropout: 0.0
+    rms_norm_eps: 1.0e-05
+    hidden_act: silu
+    torch_dtype: bfloat16
+    use_cache: true
+    pretraining_tp: 1
+    rope_scaling:
+      factor: 32.0
+      high_freq_factor: 4.0
+      low_freq_factor: 1.0
+      original_max_position_embeddings: 32000
+      rope_type: llama3
+callbacks:
+  throughput:
+    _target_: src.utils.callbacks.TokenThroughputMonitor
+  model_checkpoint:
+    _target_: lightning.pytorch.callbacks.ModelCheckpoint
+    dirpath: ${paths.output_dir}/checkpoints
+    filename: epoch_{epoch:03d}
+    monitor: val/loss
+    verbose: false
+    save_last: true
+    save_top_k: 1
+    mode: min
+    auto_insert_metric_name: false
+    save_weights_only: false
+    every_n_train_steps: null
+    train_time_interval: null
+    every_n_epochs: null
+    save_on_train_epoch_end: null
+  model_summary:
+    _target_: lightning.pytorch.callbacks.RichModelSummary
+    max_depth: -1
+  rich_progress_bar:
+    _target_: lightning.pytorch.callbacks.RichProgressBar
+  timer:
+    _target_: src.utils.callbacks.EpochTimerCallback
+  print:
+    _target_: src.utils.callbacks.PrintCallback
+  sample_counter:
+    _target_: src.utils.callbacks.SampleCounter
+logger: null
+
+trainer:
+  _target_: src.utils.trainer.ProFamTrainer
+  default_root_dir: ${paths.output_dir}
+  max_epochs: 10000
+  max_steps: -1
+  accelerator: gpu
+  devices: auto
+  check_val_every_n_epoch: 1
+  val_check_interval: 50000
+  target_tokens_per_batch: null
+  tokens_per_document: 30000
+  batch_size: ${data.batch_size}
+  deterministic: false
+  log_every_n_steps: 10
+  timeout: 120
+  profiler:
+    name: null
+    log_tensorboard: false
+    simple:
+      _target_: SimpleProfiler
+    advance:
+      _target_: AdvancedProfiler
+      filename: advanced_perf_logs
+      dirpath: ./profiler_logs
+    pytorch:
+      _target_: PyTorchProfiler
+      filename: pytorch_perf_logs
+      dirpath: ./profiler_logs
+      record_shapes: true
+      profile_memory: true
+      with_stack: true
+      with_flops: false
+      with_modules: false
+      acc_events: false
+  strategy: ddp
+  num_nodes: 1
+  sync_batchnorm: true
+  precision: bf16-true
+  min_epochs: 1000
+  accumulate_grad_batches: 2
+  use_distributed_sampler: false
+paths:
+  root_dir: ${oc.env:PROJECT_ROOT}
+  data_dir: "../data"
+  log_dir: ${paths.root_dir}/logs/${experiment_group}
+  output_dir: ${hydra:runtime.output_dir}
+  work_dir: ${hydra:runtime.cwd}
+extras:
+  ignore_warnings: false
+  enforce_tags: true
+  print_config: true
+tokenizer:
+  _target_: src.data.tokenizers.ProFamTokenizer
+  tokenizer_file: data/profam_tokenizer.json
+  unk_token: '[UNK]'
+  pad_token: '[PAD]'
+  bos_token: '[start-of-document]'
+  sep_token: '[SEP]'
+  mask_token: '?'
+  seq_struct_sep_token: '|'
+  add_final_sep: true
+  add_bos_token: true
+  add_document_token: true
+  mask_below_plddt: 80
+  max_res_pos_in_seq: ${model.max_res_pos_in_seq}
+  embed_residue_index: ${model.embed_residue_index}
+constants:
+  vocab_size: 68
+  gym_val_assay_list:
+  - BLAT_ECOLX_Jacquier_2013
+  - CALM1_HUMAN_Weile_2017
+  - DYR_ECOLI_Thompson_2019
+  - DLG4_RAT_McLaughlin_2012
+  - REV_HV1H2_Fernandes_2016
+  - TAT_HV1BR_Fernandes_2016
+  - RL40A_YEAST_Roscoe_2013
+  - P53_HUMAN_Giacomelli_2018_WT_Nutlin
+  sequence_features:
+  - ds_name
+  - identifier
+  - input_ids
+  - attention_mask
+  - original_size
+  - residue_index
+  - batch_size
+  structure_features:
+  - ds_name
+  - identifier
+  - input_ids
+  - attention_mask
+  - original_size
+  - residue_index
+  - coords
+  - coords_mask
+  - interleaved_coords_mask
+  - aa_mask
+  - plddts
+  - structure_mask
+extra_callbacks:
+  unconditional_sampling_callback:
+    prompt_builder:
+      preprocessor:
+        cfg:
+          _target_: src.data.processors.PreprocessingConfig
+          document_token: '[RAW]'
+          drop_first_protein: false
+          keep_first_protein: false
+          allow_unk: false
+          max_tokens_per_example: 8192
+          shuffle_proteins_in_document: true
+          padding: do_not_pad
+        transform_fns:
+        - _target_: src.data.processors.transforms.replace_nans_in_coords
+          _partial_: true
+          fill_value: 0.0
+        _target_: src.data.processors.ProteinDocumentPreprocessor
+      _target_: src.models.inference.PromptBuilder
+    _target_: src.pipelines.callback.SamplingEvaluationPipelineCallback
+    pipeline:
+      _target_: src.pipelines.unconditional_sequence.UnconditionalSequenceEvaluationPipeline
+      num_generations: 5
+      max_tokens: 20000
+      max_generated_length: 300
+      pipeline_id: unconditional_sampling
+      save_results_to_file: false
+    evaluators:
+      _target_: src.evaluators.esmfold.ESMFoldSamplingEvaluator
+      name: esmfold_example
+data:
+  _target_: src.data.datamodule.ProteinDataMixture
+  dataset_builders:
+    proteingym:
+      _target_: src.data.builders.proteingym.ProteinGymDataset
+      name: proteingym
+      dms_ids: ${constants.gym_val_assay_list}
+      seed: 42
+      max_mutated_sequences: null
+      mutant_bos_token: sep
+      keep_gaps: false
+      use_filtered_msa: true
+      extra_tokens_per_document: 2
+      use_msa_pos: false
+      num_proc: null
+      max_tokens_per_example: 7500
+      max_context_seqs: null
+      keep_wt: false
+      drop_wt: true
+  data_weights:
+    foldseek_s50_train: 1
+  val_dataset_batch_sizes:
+    proteingym: 1
+  batch_size: 100
+  data_dir: ${paths.data_dir}
+  num_workers: 32
+  ignore_gaps: true
+  feature_names: ${constants.sequence_features}
+  pack_to_max_tokens: 52000
+  prefetch_factor: 4
+  shuffle: true
+  interleaved: true
+  interleaved_block_size: 1000
+  total_num_train_samples: null

checkpoints/last.ckpt

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