Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,48 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -textmodel/multievolve_workflow.png filter=lfs diff=lfs merge=lfs -text
+model/streamlit_1.png filter=lfs diff=lfs merge=lfs -text

.streamlit/config.toml

@@ -0,0 +1,6 @@
+[theme]
+primaryColor = "#1E88E5"
+backgroundColor = "#FFFFFF"
+secondaryBackgroundColor = "#E3F2FD"
+textColor = "#262730"
+font = "sans serif"

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

MANIFEST.in

@@ -0,0 +1 @@
+include model/predictors/sweep_configs/*.yaml

README.md

@@ -0,0 +1,500 @@
+---
+license: apache-2.0
+language:
+- en
+- zh
+tags:
+- OneScience
+- life-science
+- protein
+- directed-evolution
+- multi-mutant
+- protein-language-model
+- MULTI-evolve
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">MULTI-evolve</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+MULTI-evolve (model-guided, universal, targeted installation of multi-mutants) is an end-to-end framework for protein directed evolution. It trains sequence-to-fitness prediction models, proposes combinatorial multi-mutants, generates MULTI-assembly site-directed mutagenesis oligonucleotides, and supports screening single-mutant candidates through a protein language model zero-shot ensemble method.
+
+Paper:
+
+> **Rapid directed evolution guided by protein language models and epistatic interactions**  
+> Science, 2026  
+> https://doi.org/10.1126/science.aea1820
+
+# Model Description
+
+The core workflow of MULTI-evolve includes:
+
+1. Train fully connected neural networks using experimental sequence-to-fitness data.
+2. Compare different data splits, sequence representations, and machine learning models.
+3. Select the best-performing prediction model to score combinatorial mutants and propose candidates.
+4. Generate MULTI-assembly site-directed mutagenesis oligonucleotides from the selected multi-mutants.
+5. In selected iterations, use a protein language model zero-shot ensemble method to screen single-mutant candidates.
+
+# Use Cases
+
+| Use case | Description |
+| --- | --- |
+| Protein directed evolution | Train fitness prediction models from experimental data and screen candidate mutations |
+| Multi-mutant design | Predict combinatorial mutations and screen multi-mutants with high predicted fitness |
+| Protein complex optimization | Support mutation formats and inputs for multichain proteins |
+| Zero-shot mutation screening | Use a protein language model ensemble method to screen candidate single mutations |
+
+# Usage
+
+## 1. Using OneCode
+
+Experience intelligent one-click AI4S programming in the OneCode online environment:
+
+[Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- Supervised model training and standard combinatorial mutation prediction in MULTI-evolve can run on a CPU or GPU/DCU.
+- Protein language model zero-shot prediction uses models such as ESM and ESM-IF; a GPU/DCU is recommended.
+
+### Set Up the Runtime Environment
+
+#### DCU Environment
+
+```bash
+# Activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+
+pip install onescience[bio] \
+  -i http://mirrors.onescience.ai:3141/pypi/simple/ \
+  --trusted-host mirrors.onescience.ai
+```
+
+#### Environment Notes
+
+- If you encounter missing dependencies or version incompatibilities during execution, refer to the dependency versions specified in `env.yml` and install or adjust the relevant dependencies as needed.
+
+### Prepare Models and Data
+
+The different MULTI-evolve functions have different model and data requirements. Prepare the resources for standard supervised training, combinatorial multi-mutant recommendation, MULTI-assembly design, protein language model zero-shot screening, and benchmark reproduction separately.
+
+#### 1) Supervised Learning Input Data
+
+To train your own protein fitness prediction model with MULTI-evolve, prepare:
+
+```text
+Wild-type protein FASTA
++
+Experimental training data CSV
+```
+
+The training data CSV must contain at least:
+
+```text
+mutation
+property_value
+```
+
+For example, the mutation format for a single-chain protein is:
+
+```text
+A40P/E61Y
+```
+
+For multichain proteins, use `:` to separate different chains:
+
+```text
+A40P/E61Y:WT
+```
+
+Where:
+
+```text
+/   Separates multiple mutations on the same chain
+:   Separates different protein chains
+WT  Indicates that the corresponding chain remains wild type
+```
+
+The official repository provides example data:
+
+```text
+data/
+├── example_protein/
+└── example_multichain_protein/
+```
+
+Therefore, no additional training data download is required to run the official basic examples.
+
+#### 2) Combinatorial Mutation Candidate Pool
+
+When running combinatorial multi-mutant recommendation, in addition to the wild-type FASTA and training data, provide a mutation pool: a list of candidate single mutations eligible for combinatorial design.
+
+Example:
+
+```text
+data/example_protein/combo_muts.csv
+```
+
+Pass this file as the:
+
+```text
+--mutation-pool
+```
+
+parameter, for example:
+
+```bash
+p2_propose.py \
+  --experiment-name multievolve_example \
+  --protein-name example_protein \
+  --wt-files apex.fasta \
+  --training-dataset example_dataset.csv \
+  --mutation-pool combo_muts.csv \
+  --top-muts-per-load 3 \
+  --export-name multievolve_proposals
+```
+
+#### 3) Protein Language Model Zero-Shot Mode
+
+The MULTI-evolve protein language model zero-shot ensemble workflow requires:
+
+```text
+Wild-type FASTA
++
+PDB/CIF protein structure
+```
+
+The current official code uses the following models:
+
+```text
+ESM-1v:
+esm1v_t33_650M_UR90S_1
+esm1v_t33_650M_UR90S_2
+esm1v_t33_650M_UR90S_3
+esm1v_t33_650M_UR90S_4
+esm1v_t33_650M_UR90S_5
+
+ESM-2:
+esm2_t36_3B_UR50D
+
+ESM-IF1:
+esm_if1_gvp4_t16_142M_UR50
+```
+
+MULTI-evolve calls these models through `fair-esm`. On the first run, if the corresponding weights are not available locally, `fair-esm` automatically downloads the models and caches them in the PyTorch Hub checkpoint directory.
+
+The default cache location is:
+
+```text
+~/.cache/torch/hub/checkpoints/
+```
+
+ESM-2 also uses the corresponding contact regression weights:
+
+```text
+esm2_t36_3B_UR50D-contact-regression.pt
+```
+- The current repository already includes `esm2_t36_3B_UR50D-contact-regression.pt` under `hub/checkpoints/`.
+
+For network-restricted or offline environments, download the weights in advance:
+
+```bash
+mkdir -p ~/.cache/torch/hub/checkpoints
+cd ~/.cache/torch/hub/checkpoints
+
+wget https://dl.fbaipublicfiles.com/fair-esm/models/esm1v_t33_650M_UR90S_1.pt
+wget https://dl.fbaipublicfiles.com/fair-esm/models/esm1v_t33_650M_UR90S_2.pt
+wget https://dl.fbaipublicfiles.com/fair-esm/models/esm1v_t33_650M_UR90S_3.pt
+wget https://dl.fbaipublicfiles.com/fair-esm/models/esm1v_t33_650M_UR90S_4.pt
+wget https://dl.fbaipublicfiles.com/fair-esm/models/esm1v_t33_650M_UR90S_5.pt
+
+wget https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t36_3B_UR50D.pt
+wget https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t36_3B_UR50D-contact-regression.pt
+
+wget https://dl.fbaipublicfiles.com/fair-esm/models/esm_if1_gvp4_t16_142M_UR50.pt
+```
+
+To save the models in the current project or another location, set the PyTorch Hub cache root with `TORCH_HOME`. For example, to use a directory in the current project:
+
+```bash
+cd /path/to/MULTI-evolve
+export TORCH_HOME=$PWD
+mkdir -p ${TORCH_HOME}/hub/checkpoints
+```
+
+Then save or symlink the weights above to:
+
+```text
+/path/to/MULTI-evolve/hub/checkpoints/
+```
+
+This avoids modifying the MULTI-evolve source code.
+
+#### 4) Benchmark DMS Data
+
+To run the official benchmark and reproduce the performance comparison across different:
+
+```text
+data splitting methods
+sequence representation methods
+machine learning models
+```
+
+you must additionally prepare the official benchmark DMS data by downloading it separately from Zenodo.
+
+Download:
+
+```text
+DOI: 10.5281/zenodo.17620759
+https://zenodo.org/records/17620759
+```
+
+After downloading, place the DMS CSV files directly in the following directory. If the directory does not exist, create it manually first:
+
+```text
+data/benchmark/datasets/
+```
+
+The benchmark script entry point in the current repository is:
+
+```text
+scripts/notebooks/benchmark/multievolve_hyperparameter_tuning.py
+```
+
+## 3. Quick Start
+
+### Download the Model Package
+
+```bash
+hf download OneScience-Group/MULTI-evolve \
+  --local-dir ./MULTI-evolve
+
+cd MULTI-evolve
+```
+
+- Standard supervised training and combinatorial mutation recommendation in MULTI-evolve do not require additional large fixed datasets; you can use the example data in the repository or your own experimental data.
+
+- The protein language model zero-shot mode may require additional ESM/ESM-IF model caches; prepare them in advance in offline environments.
+
+### Quick Verification
+
+Install the current repository:
+
+```bash
+python -m pip install -e . --no-deps
+```
+
+Check the commands:
+
+```bash
+p1_train.py --help
+p2_propose.py --help
+p3_assembly_design.py --help
+plm_zeroshot_ensemble.py --help
+```
+
+
+# Example Data
+
+The official repository provides:
+
+```text
+data/
+├── example_protein/
+├── example_multichain_protein/
+└── benchmark/
+```
+
+The official command-line examples primarily use:
+
+```bash
+cd data/example_protein
+```
+
+Typical inputs include:
+
+```text
+apex.fasta
+example_dataset.csv
+combo_muts.csv
+APEX_33overhang.fasta
+apex.cif
+```
+
+These files are used for:
+
+| File | Purpose |
+| --- | --- |
+| `apex.fasta` | Wild-type protein amino acid sequence |
+| `example_dataset.csv` | Training data |
+| `combo_muts.csv` | Combinatorial mutation candidate pool |
+| `APEX_33overhang.fasta` | DNA input required for MULTI-assembly oligonucleotide design |
+| `apex.cif` | Structure-conditioned scoring for protein language models |
+
+# Inference and Training Examples
+
+## Step 1: Train the Neural Network Model
+
+```bash
+# If the runtime environment is not active, activate the conda environment in use, such as onescience311
+conda activate onescience311
+cd data/example_protein
+
+p1_train.py \
+  --experiment-name multievolve_example \
+  --protein-name example_protein \
+  --wt-files apex.fasta \
+  --training-dataset-fname example_dataset.csv \
+  --wandb-key dummy \
+  --mode test
+```
+
+Key parameters:
+
+| Parameter | Description |
+| --- | --- |
+| `--experiment-name` | Current experiment name; keep it consistent in subsequent steps |
+| `--protein-name` | Protein name |
+| `--wt-files` | Wild-type FASTA; use commas to separate multiple FASTA files for multichain proteins |
+| `--training-dataset-fname` | Training data CSV |
+| `--mode` | `test` or `standard` |
+
+## Step 2: Propose Combinatorial Multi-Mutants
+
+```bash
+p2_propose.py \
+  --experiment-name multievolve_example \
+  --protein-name example_protein \
+  --wt-files apex.fasta \
+  --training-dataset example_dataset.csv \
+  --mutation-pool combo_muts.csv \
+  --top-muts-per-load 3 \
+  --export-name multievolve_proposals
+```
+
+The script loads the trained model saved to the local cache in Step 1 and scores the combinatorial mutation candidates.
+
+Typical output:
+
+```text
+multievolve_proposals.csv
+```
+
+For protein complexes, candidate files are also generated separately for each chain.
+
+## Step 3: Design MULTI-assembly Oligonucleotides
+
+```bash
+p3_assembly_design.py \
+  --mutations-file multievolve_proposals.csv \
+  --wt-fasta APEX_33overhang.fasta \
+  --overhang 33 \
+  --species human \
+  --oligo-direction top \
+  --tm 80 \
+  --output design
+```
+
+Where:
+
+| Parameter | Description |
+| --- | --- |
+| `--mutations-file` | Candidate mutation CSV generated in Step 2 |
+| `--wt-fasta` | Wild-type DNA FASTA containing overhangs at both ends |
+| `--overhang` | Overhang length |
+| `--species` | `human`, `ecoli`, or `yeast` |
+| `--oligo-direction` | `top` or `bottom` |
+| `--tm` | Target oligonucleotide melting temperature; the official recommendation is 80 °C |
+| `--output` | `design` or `update` |
+
+Outputs:
+
+```text
+cloning_sheet.csv
+oligos.csv
+```
+
+## Protein Language Model Zero-Shot Ensemble
+
+```bash
+plm_zeroshot_ensemble.py \
+  --wt-file apex.fasta \
+  --pdb-files apex.cif \
+  --chain-id A \
+  --variants 24 \
+  --excluded-positions 1,14,41,112 \
+  --normalizing-method aa_substitution_type
+```
+
+Where:
+
+| Parameter | Description |
+| --- | --- |
+| `--wt-file` | Wild-type protein FASTA |
+| `--pdb-files` | PDB/CIF structure files; use commas to separate multiple structures |
+| `--chain-id` | Chain ID of the target protein in the structure file |
+| `--variants` | Number of mutations nominated by each method |
+| `--excluded-positions` | Positions excluded from mutation |
+| `--normalizing-method` | `aa_substitution_type` or `aa_mutation` |
+
+This workflow ensembles four methods and produces:
+
+```text
+plm_zeroshot_ensemble_nominated_mutations.csv
+```
+
+# Output Description
+
+MULTI-evolve generates model caches, evaluation results, and candidate sequences at different stages.
+
+The official repository automatically creates the following after execution:
+
+```text
+proteins/
+└── <protein_name>/
+    ├── feature_cache/
+    ├── model_cache/
+    │   └── <dataset>/
+    │       ├── objects/
+    │       └── results/
+    ├── proposers/
+    │   └── results/
+    └── split_cache/
+        └── <dataset>/
+```
+
+The main outputs include:
+
+| Output | Description |
+| --- | --- |
+| `model_cache/` | Trained models and comparison results |
+| `feature_cache/` | Cached sequence representations |
+| `multievolve_proposals.csv` | Recommended multi-mutant candidates |
+| `cloning_sheet.csv` | MULTI-assembly cloning design sheet |
+| `oligos.csv` | Site-directed mutagenesis oligonucleotide sequences |
+| `plm_zeroshot_ensemble_nominated_mutations.csv` | Protein language model zero-shot ensemble recommendations |
+
+
+# Official OneScience Information
+
+| Platform | Main OneScience repository | Skills repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+
+# Citation and License
+
+- Original MULTI-evolve paper: [Rapid directed evolution guided by protein language models and epistatic interactions](https://doi.org/10.1126/science.aea1820).
+- The `LICENSE` in the repository root is currently **Apache License 2.0**. This license permits use, modification, distribution, and commercial use, but redistribution requires retaining the license, copyright, and attribution notices, and clearly indicating modified files.
+- Apache-2.0 also includes a patent license and explicitly does not grant rights to use the project's trademarks.
+- `setup.py` still contains an `MIT License` classifier, which is inconsistent with the actual `LICENSE` file in the repository root. For SCNet/ModelScope redistribution, use the Apache-2.0 `LICENSE` in the repository root as the authoritative license, and retain the original license file.
+- This repository is a DCU-adapted version of MULTI-evolve, with some environment configurations, dependencies, and execution procedures adjusted. Use of the repository code, model weights, and related data remains subject to the licenses and terms of use of their respective original projects.

config.json

@@ -0,0 +1,25 @@
+{
+  "framework": "Pytorch",
+  "task": "protein-engineering",
+  "subtasks": [
+    "supervised-protein-fitness-training",
+    "combinatorial-mutant-proposal",
+    "multi-assembly-oligo-design",
+    "protein-language-model-zeroshot-screening",
+    "benchmark-evaluation"
+  ],
+  "entrypoints": {
+    "train": "scripts/p1_train.py",
+    "propose": "scripts/p2_propose.py",
+    "assembly_design": "scripts/p3_assembly_design.py",
+    "zeroshot": "scripts/plm_zeroshot_ensemble.py",
+    "streamlit_app": "scripts/app.py",
+    "benchmark": "scripts/notebooks/benchmark/multievolve_hyperparameter_tuning.py"
+  },
+  "package": "model",
+  "example_data": [
+    "data/example_protein",
+    "data/example_multichain_protein"
+  ],
+  "model_cache": "hub/checkpoints"
+}

data/README.md

@@ -0,0 +1,13 @@
+# Example data
+
+## Training dataset
+
+CSV file with columns ```mutation``` and ```property_value```. Variants should be formatted as ```A40P/E61Y```, or for protein complexes as ```A40P/E61Y:WT```, where ```:``` separates the individual chains (e.g. chain 1 mutations:chain 2 mutations), ```/``` separates the individual mutations, and ```WT``` indicates the wildtype sequence.
+
+## Multi-chain protein
+
+The variable heavy domain (vh) and variable light domain (vl) chains of HuABC2, anti-CD122 antibody. 
+
+## Single chain protein
+
+Engineered ascorbate peroxidase (APEX)

data/benchmark/README.md

@@ -0,0 +1,10 @@
+# Benchmarking
+
+The multievolve_hyperparameter_tuning.py script in notebooks/benchmark can be used to train the models for benchmarking MULTI-evolve. 
+
+"dataset_summary.csv" provides details on the individual datasets used for benchmarking. The table was derived and modified from [ProteinGym](https://proteingym.org/).
+
+# Data
+
+Large datasets are not included in this repository due to size constraints. Please download DMS dataset files from Zenodo (10.5281/zenodo.17620759) and place in:
+- `data/benchmark/datasets/`

data/benchmark/dataset_summary.csv

@@ -0,0 +1,13 @@
+DMS_id,DMS_filename,UniProt_ID,taxon,source_organism,target_seq,seq_len,includes_multiple_mutants,DMS_total_number_mutants,DMS_number_multiple_mutants,DMS_binarization_cutoff,DMS_binarization_method,first_author,title,year,jo,molecule_name,selection_assay,selection_type,ProteinGym_version,coarse_selection_type
+CAPSD_AAV2S_Sinai_2021,CAPSD_AAV2S_Sinai_2021.csv,CAPSD_AAV2S,Virus,Adeno-associated virus 2,MAADGYLPDWLEDTLSEGIRQWWKLKPGPPPPKPAERHKDDSRGLVLPGYKYLGPFNGLDKGEPVNEADAAALEHDKAYDRQLDSGDNPYLKYNHADAEFQERLKEDTSFGGNLGRAVFQAKKRVLEPLGLVEEPVKTAPGKKRPVEHSPVEPDSSSGTGKAGQQPARKRLNFGQTGDADSVPDPQPLGQPPAAPSGLGTNTMATGSGAPMADNNEGADGVGNSSGNWHCDSTWMGDRVITTSTRTWALPTYNNHLYKQISSQSGASNDNHYFGYSTPWGYFDFNRFHCHFSPRDWQRLINNNWGFRPKRLNFKLFNIQVKEVTQNDGTTTIANNLTSTVQVFTDSEYQLPYVLGSAHQGCLPPFPADVFMVPQYGYLTLNNGSQAVGRSSFYCLEYFPSQMLRTGNNFTFSYTFEDVPFHSSYAHSQSLDRLMNPLIDQYLYYLSRTNTPSGTTTQSRLQFSQAGASDIRDQSRNWLPGPCYRQQRVSKTSADNNNSEYSWTGATKYHLNGRDSLVNPGPAMASHKDDEEKFFPQSGVLIFGKQGSEKTNVDIEKVMITDEEEIRTTNPVATEQYGSVSTNLQRGNRQAATADVNTQGVLPGMVWQDRDVYLQGPIWAKIPHTDGHFHPSPLMGGFGLKHPPPQILIKNTPVPANPSTTFSAAKFASFITQYSTGQVSVEIEWELQKENSKRWNPEIQYTSNYNKSVNVDFTVDTNGVYSEPRPIGTRYLTRNL,735,TRUE,42328,41796,-1.2,manual,Sinai,Generative AAV capsid diversification by latent interpolation,2021,10.1101/2021.04.16.440236,AAV,viability for AAV capsid production,,0.1,OrganismalFitness
+D7PM05_CLYGR_Somermeyer_2022,D7PM05_CLYGR_Somermeyer_2022.csv,D7PM05_CLYGR,Eukaryote,Clytia gregaria,MTALTEGAKLFEKEIPYITELEGDVEGMKFIIKGEGTGDATTGTIKAKYICTTGDLPVPWATILSSLSYGVFCFAKYPRHIADFFKSTQPDGYSQDRIISFDNDGQYDVKAKVTYENGTLYNRVTVKGTGFKSNGNILGMRVLYHSPPHAVYILPDRKNGGMKIEYNKAFDVMGGGHQMARHAQFNKPLGAWEEDYPLYHHLTVWTSFGKDPDDDETDHLTIVEVIKAVDLETYR,235,TRUE,24515,23346,12500,manual,Somermeyer,Heterogeneity of the GFP fitness landscape and data-driven protein design,2022,10.7554/eLife.75842,Green fluorescent protein cgreGFP,Fluorescence,FACS,1,Activity
+F7YBW8_MESOW_Aakre_2015,F7YBW8_MESOW_Aakre_2015.csv,F7YBW8_MESOW,Prokaryote,Mesorhizobium opportunistum (strain LMG 24607 / HAMBI 3007 / WSM2075),MANVEKMSVAVTPQQAAVMREAVEAGEYATASEIVREAVRDWLAKRELRHDDIRRLRQLWDEGKASGRPEPVDFDALRKEARQKLTEVPPNGR,93,TRUE,9192,9155,-0.001724,median,Aakre,Evolving New Protein-Protein Interaction Specificity through Promiscuous Intermediates,2015,10.1016/j.cell.2015.09.055,Antitoxin ParD3,fitness,Growth (antitoxin neutralization of ParE3),0.1,OrganismalFitness
+GFP_AEQVI_Sarkisyan_2016,GFP_AEQVI_Sarkisyan_2016.csv,GFP_AEQVI,Eukaryote,Aequorea victoria,MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTLSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK,238,TRUE,51714,50630,2.5,manual,Sarkisyan,Local fitness landscape of the green fluorescent protein,2016,10.1038/nature17995,GFP,Fluorescence,FACS,0.1,Activity
+HIS7_YEAST_Pokusaeva_2019,HIS7_YEAST_Pokusaeva_2019.csv,HIS7_YEAST,Eukaryote,Saccharomyces cerevisiae,MTEQKALVKRITNETKIQIAISLKGGPLAIEHSIFPEKEAEAVAEQATQSQVINVHTGIGFLDHMIHALAKHSGWSLIVECIGDLHIDDHHTTEDCGIALGQAFKEALGAVRGVKRFGSGFAPLDEALSRAVVDLSNRPYAVVELGLQREKVGDLSCEMIPHFLESFAEASRITLHVDCLRGKNDHHRSESAFKALAVAIREATSPNGTNDVPSTKGVLM,220,TRUE,496137,495969,0.3,manual,Pokusaeva,An experimental assay of the interactions of amino acids from orthologous sequences shaping a complex fitness landscape,2019,10.1371/journal.pgen.1008079,IGP dehydratase (HIS3),Growth,Growth,0.1,OrganismalFitness
+PHOT_CHLRE_Chen_2023,PHOT_CHLRE_Chen_2023.csv,PHOT_CHLRE,Eukaryote,Chlamydomonas reinhardtii,AGLRHTFVVADATLPDCPLVYASEGFYAMTGYGPDEVLGHNARFLQGEGTDPKEVQKIRDAIKKGEACSVRLLNYRKDGTPFWNLLTVTPIKTPDGRVSKFVGVQVDVTSKTEGKALA,118,TRUE,167529,165407,0.631701888,median,Chen,Deep Mutational Scanning of an Oxygen-Independent Fluorescent Protein CreiLOV for Comprehensive Profiling of Mutational and Epistatic Effects,2023,10.1021/acssynbio.2c00662,Phototropin,Fluorescence,FACS,1,Activity
+Q6WV13_9MAXI_Somermeyer_2022,Q6WV13_9MAXI_Somermeyer_2022.csv,Q6WV12_9MAXI,Eukaryote,Pontellina plumata,MPAMKIECRITGTLNGVEFELVGGGEGTPEQGRMTNKMKSTKGALTFSPYLLSHVMGYGFYHFGTYPSGYENPFLHAINNGGYTNTRIEKYEDGGVLHVSFSYRYEAGRVIGDFKVVGTGFPEDSVIFTDKIIRSNATVEHLHPMGDNVLVGSFARTFSLRDGGYYSFVVDSHMHFKSAIHPSILQNGGPMFAFRRVEELHSNTELGIVEYQHAFKTPIAFA,222,TRUE,31401,30260,15721.24977,median,Somermeyer,Heterogeneity of the GFP fitness landscape and data-driven protein design,2022,10.7554/eLife.75842,Green fluorescent protein ppluGFP2,Fluorescence,FACS,1,Activity
+Q8WTC7_9CNID_Somermeyer_2022,Q8WTC7_9CNID_Somermeyer_2022.csv,Q8WTC7_9CNID,Eukaryote,Aequorea macrodactyla,MSKGEELFTGIVPVLIELDGDVHGHKFSVRGEGEGDADYGKLEIKFICTTGKLPVPWPTLVTTLSYGILCFARYPEHMKMNDFFKSAMPEGYIQERTIFFQDDGKYKTRGEVKFEGDTLVNRIELKGMDFKEDGNILGHKLEYNFNSHNVYIMPDKANNGLKVNFKIRHNIEGGGVQLADHYQTNVPLGDGPVLIPINHYLSCQTAISKDRNETRDHMVFLEFFSACGHTHGMDELYK,238,TRUE,33510,32309,5000,manual,Somermeyer,Heterogeneity of the GFP fitness landscape and data-driven protein design,2022,10.7554/eLife.75842,Green fluorescent protein amacGFP,Fluorescence,FACS,1,Activity
+SPG1_STRSG_Wu_2016,SPG1_STRSG_Wu_2016.csv,SPG1_STRSG,Prokaryote,Streptococcus sp. group G,MEKEKKVKYFLRKSAFGLASVSAAFLVGSTVFAVDSPIEDTPIIRNGGELTNLLGNSETTLALRNEESATADLTAAAVADTVAAAAAENAGAAAWEAAAAADALAKAKADALKEFNKYGVSDYYKNLINNAKTVEGIKDLQAQVVESAKKARISEATDGLSDFLKSQTPAEDTVKSIELAEAKVLANRELDKYGVSDYHKNLINNAKTVEGVKELIDEILAALPKTDQYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDDATKTFTVTEKPEVIDASELTPAVTTYKLVINGKTLKGETTTKAVDAETAEKAFKQYANDNGVDGVWTYDDATKTFTVTEMVTEVPGDAPTEPEKPEASIPLVPLTPATPIAKDDAKKDDTKKEDAKKPEAKKDDAKKAETLPTTGEGSNPFFTAAALAVMAGAGALAVASKRKED,448,TRUE,149360,149284,0.122438875,median,Wu,Adaptation in protein fitness landscapes is facilitated by indirect paths,2016,10.7554/eLife.16965,GB1,Binding (IgG),binding,1,Binding
+Q65J43_BACLD_g4_Thomas_2025,Q65J43_BACLD_g4_Thomas_2025.csv,Q65J43_BACLD,Prokaryote,Bacillus licheniformis,MIKKWAVHLLFSALVLLGLSGGAAYSPQHAEGAARYDDVLYFPASRYPETGAHISDAIKAGHADVCTIERSGADKRRQESLKGIPTKPGFDRDEWPMAMCEEGGKGASVRYVSSSDNRGAGSWVGNRLNGYADGTRILFIVQ,142,TRUE,15404,15183,1.33474,WT,Thomas,Engineering highly active nuclease enzymes with machine learning and high-throughput screening,2025,10.1016/j.cels.2025.101236,NucB,Fluorescence,FACS,,Activity
+TRPB1_THEMA_Tm9D8_Johnston_2024,TRPB1_THEMA_Tm9D8_Johnston_2024.csv,TRPB1_THEMA,Prokaryote,Escherichia coli,MKGYFGPYGGQYVPEILMGALEELEAAYEGIMKDESFWKEFNDLLRDYAGRPTPLYFARRLSEKYGARVYLKREDLLHTGAHKINNAIGQVLLAKLMGKTRIIAETGAGQHGVATATAAALFGMECVIYMGEEDTIRQKLNVERMKLLGAKVVPVKSGSRTLKDAIDEALRDWITNLQTTYYVFGSVVGPHPYPIIVRNFQKVIGEETKKQIPEKEGRLPDYIVACVSGGSNAAGIFYPFIDSGVKLIGVEAGGEGLETGKHAASLLKGKIGYLHGSKTFVLQDDWGQVQVSHSVSAGLDYSGVGPEHAYWRETGKVLYDAVTDEEALDAFIELSRLEGIIPALESSHALAYLKKINIKGKVVVVNLSGRGDKDLESVLNHPYVRERIR,389,TRUE,159129,159052,0.40807,WT,Johnston,A combinatorially complete epistatic fitness landscape in an enzyme active site,2024,10.1073/pnas.2400439121,TrpB,Growth,Growth,,Stability
+VH14_Li_2023,VH14_Li_2023.csv,,Human,Homo sapiens,EVQLVETGGGLVQPGGSLRLSCAASGFTLNSYGISWVRQAPGKGPEWVSVIYSDGRRTFYGDSVKGRFTISRDTSTNTVYLQMNSLRVEDTAVYYCAKGRAAGTFDSWGQGTLVTVSSGGGGSGGGGSGGGGSDVVMTQSPESLAVSLGERATISCKSSQSVLYESRNKNSVAWYQQKAGQPPKLLIYWASTRESGVPDRFSGSGSGTDFTLTISSLQAEDAAVYYCQQYHRLPLSFGGGTKVEIK,246,TRUE,35889,35452,-0.787659364,WT,Li,Machine learning optimization of candidate antibody yields highly diverse sub-nanomolar affinity antibody libraries,2023,10.1038/s41467-023-39022-2,scFv,Binding,,,Binding

data/example_multichain_protein/combo_muts.csv

@@ -0,0 +1,10 @@
+A40P
+S51I
+T116L
+L11Q
+Q13P
+S63A
+S65H
+F203G
+I226L
+I226P

data/example_multichain_protein/example_dataset.csv

@@ -0,0 +1,254 @@
+mutation,property_value
+WT:WT,1
+A40P:WT,1.006
+A50R:WT,1.034
+D106W:WT,0.923
+D76N:WT,0.977
+E35H:WT,0.867
+E35N:WT,0.93
+E61Y:WT,0.748
+E91S:WT,0.871
+F27W:WT,0.651
+F32Y:WT,0.802
+G9V:WT,0.238
+I48V:WT,0.339
+I71T:WT,0.783
+M108F:WT,0.853
+M34H:WT,0.498
+Q113D:WT,0.818
+Q113P:WT,0.858
+Q84D:WT,0.897
+R100I:WT,1.015
+R100L:WT,0.86
+R104G:WT,0.924
+R19T:WT,1.297
+S121P:WT,1.312
+S51I:WT,1.077
+S80T:WT,0.99
+T116L:WT,1.369
+T28D:WT,1.171
+T28P:WT,0.901
+T90P:WT,1.18
+V72I:WT,1.179
+Y103N:WT,1.179
+Y105G:WT,1.098
+Y110V:WT,1.437
+Y33W:WT,1.409
+A55P:WT,1.219
+A64D:WT,1.077
+D106Y:WT,1.036
+D57G:WT,0.763
+E6A:WT,0.929
+E6G:WT,0.435
+E91R:WT,0.659
+F70W:WT,0.487
+G107Y:WT,0.601
+G44T:WT,0.904
+K43Q:WT,1.123
+L11E:WT,0.886
+L11Q:WT,1.108
+M108Y:WT,0.788
+N53P:WT,1.009
+N56E:WT,0.945
+Q113A:WT,0.814
+Q13P:WT,1.219
+Q3K:WT,1.023
+Q84V:WT,0.973
+R104Y:WT,0.919
+R19W:WT,0.75
+R52A:WT,0.951
+S101E:WT,1.014
+S25G:WT,0.804
+S63A:WT,0.993
+S65H:WT,0.891
+T115K:WT,0.354
+V2K:WT,1.02
+Y102P:WT,0.945
+Y33D:WT,1.274
+Y82R:WT,0.66
+WT:D49Y,0.148
+WT:E78A,0.604
+WT:E80S,0.397
+WT:F61Y,0.354
+WT:F82G,1.289
+WT:I105L,1.533
+WT:I105P,1.176
+WT:I26N,1.777
+WT:I26S,1.889
+WT:L21I,1.728
+WT:L45K,0.204
+WT:M32L,0.475
+WT:N52T,0.97
+WT:N52V,1.242
+WT:N91E,1.629
+WT:N91Y,0.534
+WT:Q36K,0.455
+WT:Q88M,1.562
+WT:R18T,1.523
+WT:S24R,1.358
+WT:S24T,1.231
+WT:S27Q,1.09
+WT:S28P,1.769
+WT:T50A,0.146
+WT:T71R,0.731
+WT:V29I,0.617
+WT:V54Q,0.327
+WT:V84T,1.762
+WT:W90Y,0.528
+WT:Y31W,0.608
+WT:Y33A,0.006
+WT:Y33C,0.462
+WT:Y70F,1.467
+WT:Y93P,1.463
+WT:Y93W,1.303
+WT:A42H,0.445
+WT:A59M,0.56
+WT:A59P,0.994
+WT:C23W,0.001
+WT:E104V,0.46
+WT:E78L,0.518
+WT:E80G,0.332
+WT:F82Y,1.167
+WT:F97Y,0.285
+WT:G100V,0.06
+WT:G99N,0.857
+WT:G99P,0.735
+WT:I105N,0.59
+WT:K38L,0.45
+WT:L46T,0.099
+WT:L4M,0.311
+WT:L53R,1.764
+WT:P43T,0.11
+WT:R44K,0.885
+WT:S12H,0.74
+WT:S27E,1.6
+WT:S30P,1.732
+WT:S30Y,0.356
+WT:S7Q,0.836
+WT:T20V,1.589
+WT:T92Y,0.965
+WT:T96D,0.538
+WT:W90L,1.342
+WT:Y33W,0.352
+WT:Y48R,0.62
+WT:Y93L,0.505
+A40P:F82G,1.338
+A40P:I105L,1.428
+A40P:I105P,1.422
+A40P:L21I,1.561
+A40P:N52T,1.092
+A40P:N52V,1.168
+A40P:R18T,1.551
+A40P:S24R,1.232
+A40P:S28P,1.785
+A40P:V84T,1.878
+A40P:Y70F,1.359
+A40P:Y93P,1.7
+A40P:Y93W,1.186
+A40P:F82Y,1.5
+A40P:L53R,2.078
+A40P:S30P,1.839
+A40P:T20V,1.388
+S51I:F82G,1.19
+S51I:I105L,1.048
+S51I:I105P,0.875
+S51I:L21I,1.43
+S51I:N52T,0.836
+S51I:N52V,0.8
+S51I:R18T,0.998
+S51I:S24R,1.076
+S51I:S28P,1.219
+S51I:V84T,1.711
+S51I:Y70F,1.28
+S51I:Y93P,1.482
+S51I:Y93W,1.313
+S51I:F82Y,0.847
+S51I:L53R,1.661
+S51I:S30P,1.569
+S51I:T20V,0.86
+T116L:F82G,1.54
+T116L:I105L,1.445
+T116L:I105P,1.209
+T116L:L21I,1.92
+T116L:N52T,0.984
+T116L:N52V,1.172
+T116L:R18T,1.704
+T116L:S24R,1.577
+T116L:S28P,2.098
+T116L:V84T,2.126
+T116L:Y70F,1.7
+T116L:Y93P,1.62
+T116L:Y93W,1.477
+T116L:F82Y,1.255
+T116L:L53R,2.41
+T116L:S30P,1.929
+T116L:T20V,1.606
+L11Q:F82G,1.913
+L11Q:I105L,1.707
+L11Q:I105P,1.222
+L11Q:L21I,1.98
+L11Q:N52T,1.143
+L11Q:N52V,1.491
+L11Q:R18T,1.548
+L11Q:S24R,1.525
+L11Q:S28P,2.08
+L11Q:V84T,1.98
+L11Q:Y70F,1.441
+L11Q:Y93P,1.14
+L11Q:Y93W,0.733
+L11Q:F82Y,1.139
+L11Q:L53R,1.958
+L11Q:S30P,1.666
+L11Q:T20V,1.258
+Q13P:F82G,1.502
+Q13P:I105L,1.548
+Q13P:I105P,1.272
+Q13P:L21I,1.73
+Q13P:N52T,0.917
+Q13P:N52V,1.245
+Q13P:R18T,1.216
+Q13P:S24R,1.462
+Q13P:S28P,1.842
+Q13P:V84T,2.395
+Q13P:Y70F,1.99
+Q13P:Y93P,1.895
+Q13P:Y93W,1.31
+Q13P:F82Y,1.325
+Q13P:L53R,2.547
+Q13P:S30P,1.959
+Q13P:T20V,1.659
+S63A:F82G,1.357
+S63A:I105L,1.728
+S63A:I105P,1.141
+S63A:L21I,1.63
+S63A:N52T,0.935
+S63A:N52V,1.105
+S63A:R18T,1.625
+S63A:S24R,1.435
+S63A:S28P,2.055
+S63A:V84T,2.517
+S63A:Y70F,1.623
+S63A:Y93P,1.938
+S63A:Y93W,1.44
+S63A:F82Y,1.129
+S63A:L53R,2.276
+S63A:S30P,1.869
+S63A:T20V,1.719
+S65H:F82G,1.34
+S65H:I105L,1.308
+S65H:I105P,0.763
+S65H:L21I,1.619
+S65H:N52T,0.722
+S65H:N52V,0.84
+S65H:R18T,1.346
+S65H:S24R,1.285
+S65H:S28P,2.09
+S65H:V84T,1.78
+S65H:Y70F,1.235
+S65H:Y93P,0.732
+S65H:Y93W,0.685
+S65H:F82Y,0.853
+S65H:L53R,1.669
+S65H:S30P,1.45
+S65H:T20V,1.203
+Q13P:L21I/S28P/L53R/V84T,2.018

data/example_multichain_protein/vh_chain1.fasta

@@ -0,0 +1,2 @@
+>vh_anticd122
+EVQLVESGGGLVQPGGSLRLSCAASGFTFSDFYMEWVRQAPGKGLEWIAASRNKANDYTTEYSASVKGRFIVSRDDSKNSLYLQMNSLKTEDTAVYYCARSYYRYDGMDYWGQGTTVTVSS

data/example_multichain_protein/vl_chain2.fasta

@@ -0,0 +1,2 @@
+>vl_anticd122
+EIVLTQSPATLSLSPGERATLSCSAISSVSYMYWYQQKPGQAPRLLIYDTSNLVSGVPARFSGSGSGTDYTLTISSLEPEDFAVYYCQQWNTYPYTFGGGTKVEIK

data/example_protein/APEX_33overhang.fasta

@@ -0,0 +1,2 @@
+>APEX_33overhang
+TCTCGCACCATGGACTACAAGGATGACGACGATAAGGGAAAGTCTTACCCAACTGTGAGTGCTGATTACCAGGACGCCGTTGAGAAGGCGAAGAAGAAGCTCAGAGGCTTCATCGCTGAGAAGAGATGCGCTCCTCTAATGCTCCGTTTGGCATTCCACTCTGCTGGAACCTTTGACAAGGGCACGAAGACCGGTGGACCCTTCGGAACCATCAAGCACCCTGCCGAACTGGCTCACAGCGCTAACAACGGTCTTGACATCGCTGTTAGGCTTTTGGAGCCACTCAAGGCGGAGTTCCCTATTTTGAGCTACGCCGATTTCTACCAGTTGGCTGGCGTTGTTGCCGTTGAGGTCACGGGTGGACCTAAGGTTCCATTCCACCCTGGAAGAGAGGACAAGCCTGAGCCACCACCAGAGGGTCGCTTGCCCGATGCCACTAAGGGTTCTGACCATTTGAGAGATGTGTTTGGCAAAGCTATGGGGCTTACTGACCAAGATATCGTTGCTCTATCTGGGGGTCACACTATTGGAGCTGCACACAAGGAGCGTTCTGGATTTGAGGGTCCCTGGACCTCTAATCCTCTTATTTTCGACAACTCATACTTCACGGAGTTGTTGAGTGGTGAGAAGGAAGGTCTCCTTCAGCTACCTTCTGACAAGGCTCTTTTGTCTGACCCTGTATTCCGCCCTCTCGTTGACAAATATGCAGCGGACGAAGATGCCTTCTTTGCTGATTACGCTGAGGCTCACCAAAAGCTTTCCGAGCTTGGGTTTGCTGATGCCCTGCAGCTGCCTCCCCTGGAGCGCCTGACCCTGGACTAGTAAGCGAGACGAAGACGAATTCTCTAGATATAG

data/example_protein/MULTI-assembly_input.csv

@@ -0,0 +1,15 @@
+N72A/A167R/T192K
+A167R/T192K/D222E
+S138A/A167R/T192K
+N72A/A167R/T192K/D222E
+N72A/S138A/A167R/T192K
+N72A/A167R/T192K/D229A
+N72A/S138A/A167R/T192K/D222E
+N72A/A167R/T192K/D222E/D229A
+K61R/N72A/A167R/T192K/D222E
+N72A/S138A/A167R/T192K/D222E/D229A
+K61R/N72A/S138A/A167R/T192K/D222E
+K61R/N72A/A167R/T192K/D222E/D229A
+K61R/N72A/S138A/A167R/T192K/D222E/D229A
+N72A/S138A/A148Q/A167R/T192K/D222E/D229A
+N72A/S138A/A167R/I185V/T192K/D222E/D229A

data/example_protein/apex.fasta

@@ -0,0 +1,2 @@
+>Lam2015_APEX_translation Expresses APEX2 in the cytosol of mammalian cells
+MGKSYPTVSADYQDAVEKAKKKLRGFIAEKRCAPLMLRLAFHSAGTFDKGTKTGGPFGTIKHPAELAHSANNGLDIAVRLLEPLKAEFPILSYADFYQLAGVVAVEVTGGPKVPFHPGREDKPEPPPEGRLPDATKGSDHLRDVFGKAMGLTDQDIVALSGGHTIGAAHKERSGFEGPWTSNPLIFDNSYFTELLSGEKEGLLQLPSDKALLSDPVFRPLVDKYAADEDAFFADYAEAHQKLSELGFADA

data/example_protein/combo_muts.csv

@@ -0,0 +1,15 @@
+T192V
+T192K
+A167R
+N72A
+D222E
+A148Q
+D229A
+S138A
+K61R
+S196A
+I185V
+L84V
+E87Q
+G50R
+L80M

data/example_protein/example_dataset.csv

@@ -0,0 +1,185 @@
+mutation,property_value,round
+WT,1,0
+A167R/T192V,5.003,1
+T192V,2.327,0
+T192K,2.637,0
+A167R,2.277,0
+N72A,2.157,0
+D222E,1.843,0
+A148Q,1.733,0
+D229A,1.73,0
+S138A,1.697,0
+K61R,1.66,0
+S196A,1.527,0
+I185V,1.443,0
+L84V,1.527,0
+E87Q,1.49,0
+G50R,1.877,0
+L80M,1.447,0
+K21R,1.31,0
+A148R,1.307,0
+S69G,1.3,0
+G50N,1.29,0
+Q240L,1.277,0
+P123A,1.267,0
+E237K,1.213,0
+K20R,1.19,0
+L84I,1.19,0
+S213T,1.19,0
+D143E,1.153,0
+A250D,1.14,0
+Q154K,1.1,0
+K21A,1.07,0
+G2A,0.997,0
+F47Y,0.993,0
+K22D,0.937,0
+E171D,0.897,0
+R31N,0.86,0
+S43E,0.83,0
+F26L,0.823,0
+I165L,0.803,0
+E128A,0.717,0
+G162A,0.603,0
+Y93W,0.557,0
+K49A,0.55,0
+N182T,0.49,0
+H116R,0.483,0
+Q240E,0.36,0
+H68Y,0.253,0
+K3H,0.24,0
+Q204L,0.223,0
+K170G,0.197,0
+F247C,0.15,0
+H239W,0.137,0
+H239L,0.033,0
+L220W,0.753,0
+A168C,0.297,0
+G50E,1.423,0
+V216G,0.903,0
+A148H,1.297,0
+G146I,1.387,0
+K49V,0.663,0
+R31P,0.237,0
+G2P,0.863,0
+P123L,0.987,0
+E120I,1.033,0
+L81C,0.417,0
+L66Q,0.34,0
+F57N,0.277,0
+L35I,1.247,0
+V102C,0.513,0
+A10P,1.233,0
+P123K,1.123,0
+T7P,1.133,0
+V216D,0.833,0
+D143Y,0.74,0
+E128P,1.067,0
+E65M,0.5,0
+R172Y,0.787,0
+G25E,1.24,0
+L37V,0.23,0
+K85M,0.49,0
+L141A,0.287,0
+A148M,0.923,0
+L205F,0.353,0
+A167R/T192K,7.267,1
+N72A/T192V,4.3,1
+N72A/T192K,6.273,1
+N72A/A167R,4.723,1
+T192V/D222E,3.81,1
+T192K/D222E,5.47,1
+A167R/D222E,5.91,1
+N72A/D222E,5.7,1
+A148Q/T192V,2.147,1
+A148Q/T192K,3.333,1
+A148Q/A167R,3.46,1
+N72A/A148Q,3.08,1
+A148Q/D222E,3.193,1
+T192V/D229A,3.17,1
+T192K/D229A,4.523,1
+A167R/D229A,4.557,1
+N72A/D229A,3.973,1
+D222E/D229A,3.023,1
+A148Q/D229A,2.32,1
+S138A/T192V,3.007,1
+S138A/T192K,5.277,1
+S138A/A167R,4.847,1
+N72A/S138A,4.453,1
+S138A/D222E,4.233,1
+S138A/D229A,3.477,1
+K61R/T192V,2.41,1
+K61R/T192K,4.253,1
+K61R/A167R,4.153,1
+K61R/N72A,3.037,1
+K61R/D222E,2.923,1
+K61R/A148Q,2.417,1
+K61R/D229A,2.873,1
+K61R/S138A,3.103,1
+T192V/S196A,3.627,1
+A167R/S196A,3.24,1
+N72A/S196A,2.697,1
+S196A/D222E,2.26,1
+A148Q/S196A,1.953,1
+S196A/D229A,1.85,1
+S138A/S196A,2.167,1
+K61R/S196A,2.533,1
+I185V/T192V,3.023,1
+I185V/T192K,4.017,1
+A167R/I185V,3.277,1
+N72A/I185V,2.313,1
+I185V/D222E,2.373,1
+A148Q/I185V,2.013,1
+I185V/D229A,1.977,1
+S138A/I185V,2.453,1
+K61R/I185V,2.293,1
+I185V/S196A,2.06,1
+L84V/T192V,1.963,1
+L84V/T192K,3.53,1
+L84V/A167R,3.033,1
+N72A/L84V,2.573,1
+L84V/D222E,2.43,1
+L84V/A148Q,1.967,1
+L84V/D229A,2.103,1
+L84V/S138A,2.417,1
+K61R/L84V,2.4,1
+L84V/S196A,1.327,1
+L84V/I185V,1.487,1
+E87Q/T192V,1.8,1
+E87Q/T192K,2.397,1
+E87Q/A167R,3.04,1
+N72A/E87Q,2.43,1
+E87Q/D222E,2.597,1
+E87Q/A148Q,1.99,1
+E87Q/D229A,1.943,1
+E87Q/S138A,1.297,1
+K61R/E87Q,1.307,1
+E87Q/S196A,1.32,1
+E87Q/I185V,1.05,1
+L84V/E87Q,1.113,1
+G50R/T192V,1.913,1
+G50R/T192K,2.997,1
+G50R/A167R,3.003,1
+G50R/N72A,2.48,1
+G50R/D222E,1.83,1
+G50R/A148Q,1.31,1
+G50R/D229A,1.283,1
+G50R/S138A,1.483,1
+G50R/K61R,1.36,1
+G50R/S196A,1.507,1
+G50R/I185V,1.483,1
+G50R/L84V,1.42,1
+G50R/E87Q,1.577,1
+L80M/T192V,1.84,1
+L80M/T192K,2.83,1
+L80M/A167R,2.517,1
+N72A/L80M,2.14,1
+L80M/D222E,2.127,1
+L80M/A148Q,1.797,1
+L80M/D229A,1.757,1
+L80M/S138A,1.897,1
+K61R/L80M,1.237,1
+L80M/S196A,1.25,1
+L80M/I185V,1.187,1
+L80M/L84V,1.237,1
+L80M/E87Q,0.827,1
+G50R/L80M,1.353,1

data/example_protein/run.sh

@@ -0,0 +1,43 @@
+#!/bin/bash
+module purge
+module load sghpc-mpi-gcc/26.3
+#module load compiler/dtk/25.04 
+source /public/software/apps/anaconda3/2023.09/etc/profile.d/conda.sh
+conda activate model_bio
+export WANDB_MODE=disabled
+export WANDB_SILENT=true
+export WANDB_DISABLED=true
+
+
+export HIP_VISIBLE_DEVICES=1
+
+#python -u ~/.conda/envs/model_bio/bin/p1_train.py \
+#    --experiment-name multievolve_example \
+#    --protein-name example_protein \
+#    --wt-files apex.fasta \
+#    --training-dataset-fname example_dataset.csv \
+#    --wandb-key dummy_key \
+#    --mode test
+
+#检查模型文件
+#find /public/home/scnb9biwet/jiangqq/MULTI-evolve-main/proteins/example_protein/model_cache -name "*.pth" -o -name "*.pt"
+
+## Step 2：Propose MULTI-evolve Variants，利用训练好的模型来推荐突变
+#python -u ~/.conda/envs/model_bio/bin/p2_propose.py \
+#--experiment-name multievolve_example \
+#--protein-name example_protein \
+#--wt-files apex.fasta \
+#--training-dataset example_dataset.csv \
+#--mutation-pool combo_muts.csv \
+#--top-muts-per-load 3 \
+#--export-name multievolve_proposals
+
+## step 3: 生成 MULTI-assembly 寡核苷酸
+python -u ~/.conda/envs/model_bio/bin/p3_assembly_design.py \
+--mutations-file ../../proteins/example_protein/multievolve_proposals.csv \
+--wt-fasta APEX_33overhang.fasta \
+--overhang 33 \
+--species human \
+--oligo-direction top \
+--tm 80 \
+--output design

env.yml

@@ -0,0 +1,26 @@
+# updated 5/21/2025
+name: multievolve
+channels:
+  - conda-forge
+dependencies:
+  - python=3.11
+  - pip
+  - tqdm
+  - ipykernel
+  - numpy=1.26
+  - pandas=2.2
+  - matplotlib
+  - seaborn=0.13
+  - scipy=1.15
+  - biopython=1.85
+  - scikit-learn=1.6
+  - scikit-optimize=0.10
+  - pip:
+      - torch==2.6.0
+      - wandb==0.19
+      - Levenshtein
+      - streamlit==1.45
+      - fair-esm
+      - biotite==0.41.2
+      - ankh
+      - protobuf

env_mac.yml

@@ -0,0 +1,27 @@
+# updated 5/21/2025
+name: multievolve
+channels:
+  - conda-forge
+dependencies:
+  - python=3.11
+  - pip
+  - tqdm
+  - ipykernel
+  - numpy=1.26
+  - pandas=2.2
+  - matplotlib
+  - seaborn=0.13
+  - scipy=1.15
+  - biopython=1.85
+  - scikit-learn=1.6
+  - scikit-optimize=0.10
+  - pip:
+      - torch==2.2.2
+      - transformers==4.42.4
+      - wandb==0.19
+      - Levenshtein
+      - streamlit==1.45
+      - fair-esm
+      - biotite==0.41.2
+      - ankh
+      - protobuf

hub/checkpoints/esm2_t36_3B_UR50D-contact-regression.pt

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

model/__init__.py

@@ -0,0 +1,6 @@
+# Import main components
+from model.splitters import *
+from model.predictors import *
+from model.proposers import *
+from model.utils import *
+from model.featurizers import *

model/featurizers/__init__.py

@@ -0,0 +1,7 @@
+from model.featurizers.base_featurizers import *
+from model.featurizers.esm_featurizers import *
+from model.featurizers.msa_featurizers import *
+from model.featurizers.zeroshot_featurizers import *
+from model.featurizers.combinatorial_featurizers import *
+from model.featurizers.ankh_featurizers import *
+from model.featurizers.prott5_featurizers import *

model/featurizers/ankh_featurizers.py

@@ -0,0 +1,82 @@
+import ankh
+import numpy as np
+import torch
+
+from model.featurizers.base_featurizers import BaseFeaturizer
+
+# alternate name: AnkhBaseFeaturizer
+class AnkhFeaturizer(BaseFeaturizer):
+    def __init__(self, 
+                protein=None, 
+                use_cache=False,
+                model_version=None, 
+                batch_size=968,
+                model_type="ankh",
+                **kwargs):
+            
+        super().__init__(model_type,protein, use_cache, **kwargs)
+
+        self.batch_size = batch_size
+        self.model_version = model_version
+
+    def featurize_ankh(self, seqs):
+
+        if self.model_version == 'large':
+            self.model, self.tokenizer = ankh.load_large_model()
+        elif self.model_version == 'base':
+            self.model, self.tokenizer = ankh.load_base_model()
+        else:
+            raise ValueError(f"Invalid model version: {self.model_version}")
+        self.model.eval()
+        self.model.to(self.device)
+
+        input_seqs = [list(seq) for seq in seqs]
+
+        seq_batch = []
+
+        for i in range(0, len(input_seqs), self.batch_size):
+            batch = input_seqs[i:i + self.batch_size]
+            outputs = self.tokenizer(
+                batch, 
+                add_special_tokens=True, 
+                padding=True, 
+                is_split_into_words=True, 
+                return_tensors="pt",
+            )
+            outputs = {key: val.to(self.device) for key, val in outputs.items()}
+            with torch.no_grad():
+                embeddings = self.model(input_ids=outputs['input_ids'], attention_mask=outputs['attention_mask'])
+            seq_batch.append(embeddings['last_hidden_state'].mean(axis=1).cpu().numpy())
+
+        return np.concatenate(seq_batch)
+
+# alternate name: AnkhBaseEmbedFeaturizer
+class AnkhBaseFeaturizer(AnkhFeaturizer):
+    def __init__(self, 
+                protein=None, 
+                use_cache=False,
+                model_version="base", 
+                batch_size=968,
+                model_type="ankh_base",
+                **kwargs):
+        super().__init__(protein, use_cache, model_version, batch_size, model_type, **kwargs)
+
+    def custom_featurizer(self, seqs):
+
+        X = self.featurize_ankh(seqs)
+        return X
+
+# alternate name: AnkhLargeEmbedFeaturizer
+class AnkhLargeFeaturizer(AnkhFeaturizer):
+    def __init__(self, 
+                protein=None, 
+                use_cache=False,
+                model_version="large", 
+                batch_size=968,
+                model_type="ankh_large",
+                **kwargs):
+        super().__init__(protein, use_cache, model_version, batch_size, model_type, **kwargs)
+
+    def custom_featurizer(self, seqs):
+        X = self.featurize_ankh(seqs)
+        return X

model/featurizers/base_featurizers.py

@@ -0,0 +1,266 @@
+from abc import ABC, abstractmethod
+import numpy as np
+import torch
+
+# root_folder = os.path.dirname(os.path.dirname(__file__))
+# sys.path.append(root_folder)
+
+from model.utils.other_utils import AAs
+from model.utils.featurizer_utils import seqs_to_georgiev, featurize_aa_idx
+from model.utils.cache_utils import load_cache, update_cache
+
+
+class BaseFeaturizer(ABC):
+    """Abstract base class for featurizers.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+
+    Example Usage:
+
+    featurizer = BaseFeaturizer(
+        model_type='onehot',      # Type of featurization used
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        flatten_features=False    # Whether to flatten output features
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(self, model_type=None, protein=None, use_cache=False, flatten_features=False, **kwargs):
+        """
+        Args:
+            model_type (str): Type of featurization model.
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            flatten_features (bool): Whether to flatten output features.
+            **kwargs: Additional keyword arguments.
+        """
+        self.model_type = model_type
+        self.name = str(model_type)
+        self.protein = protein
+        self.use_cache = use_cache
+        self.flatten_features = flatten_features
+        self.set_parameters(**kwargs)
+
+        if torch.backends.mps.is_available():
+            self.device = torch.device("mps")
+        elif torch.cuda.is_available():
+            self.device = torch.device("cuda:0")
+        else:
+            self.device = torch.device("cpu")
+
+    def set_parameters(self, **kwargs):
+        """Sets additional parameters from kwargs."""
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def load_features(self, seqs):
+        """
+        Loads cached features if they exist.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+
+        Returns:
+            tuple: (seq_to_feature dict, original sequences, unique sorted sequences)
+        """
+        print("Loading features...")
+
+        assert (
+            self.protein is not None
+        ), "No protein specified to cache. Either specify a protein or set use_cache to False."
+        original_seqs = seqs
+        cache = load_cache(self.model_type, self.protein)
+        seq_to_feature = {seq: cache[seq] for seq in seqs if seq in cache}
+        seqs = [seq for seq in seqs if seq not in cache]
+        print(f"Seqs in cache: {len(seq_to_feature)} | Seqs not in cache: {len(seqs)}")
+        del cache  # free up memory
+
+        unique_seqs = {}
+        for seq in seqs:
+            if seq not in unique_seqs:
+                unique_seqs[seq] = len(unique_seqs)
+
+        unique_seqs_sorted = sorted(unique_seqs.keys(), key=lambda k: unique_seqs[k])
+
+        return seq_to_feature, original_seqs, unique_seqs_sorted
+
+    @abstractmethod
+    def custom_featurizer(self, seqs, **kwargs):
+        """
+        Custom featurizer to be implemented in subclasses.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional keyword arguments.
+        """
+        pass
+
+    def featurize(self, seqs, **kwargs):
+        """
+        Featurizes input sequences.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: Array of featurized sequences.
+        """
+        if self.use_cache:
+            seqs_to_feature, original_seqs, unique_seqs_sorted = self.load_features(
+                seqs
+            )
+        else:
+            seqs_to_feature = {}
+            original_seqs, unique_seqs_sorted = seqs, seqs
+
+        if len(unique_seqs_sorted) > 0:
+            X_unique = self.custom_featurizer(unique_seqs_sorted, **kwargs)
+
+            for idx, seq in enumerate(unique_seqs_sorted):
+                seqs_to_feature[seq] = X_unique[idx]
+
+        if self.use_cache:
+            update_cache(self.model_type, self.protein, seqs_to_feature)
+
+        X = np.array([seqs_to_feature[seq] for seq in original_seqs])
+
+        if self.flatten_features == True:
+            X = X.reshape(len(X), -1)
+
+        return X
+
+
+class OneHotFeaturizer(BaseFeaturizer):
+    """Class for one-hot encoding of sequences.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+
+    Example Usage:
+
+    featurizer = OneHotFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        flatten_features=False    # Whether to flatten output features
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(model_type="onehot", **kwargs)
+
+    def custom_featurizer(self, seqs, **kwargs):
+        """
+        One-hot encodes sequences.
+
+        Args:
+            seqs (list): List of sequences to encode.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: One-hot encoded sequences.
+        """
+        data = [[char for char in seq] for seq in seqs]
+
+        from sklearn.preprocessing import OneHotEncoder
+
+        enc = OneHotEncoder(
+            categories=([AAs + ["X"]] * len(data[0])),
+            sparse_output=False,
+        )
+        X = enc.fit_transform(data).reshape(len(data), len(data[0]), len(AAs + ["X"]))
+
+        return X
+
+
+class GeorgievFeaturizer(BaseFeaturizer):
+    """Class for Georgiev encoding of sequences.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+
+    Example Usage:
+
+    featurizer = GeorgievFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        flatten_features=False    # Whether to flatten output features
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(model_type="georgiev", **kwargs)
+
+    def custom_featurizer(self, seqs, **kwargs):
+        """
+        Applies Georgiev encoding to sequences.
+
+        Args:
+            seqs (list): List of sequences to encode.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: Georgiev encoded sequences.
+        """
+        X = seqs_to_georgiev(seqs)
+
+        return X
+
+
+class AAIdxFeaturizer(BaseFeaturizer):
+    """Class for amino acid index encoding of sequences.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+
+    Example Usage:
+
+    featurizer = AAIdxFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        flatten_features=False    # Whether to flatten output features
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(model_type="aa_idx", **kwargs)
+
+    def custom_featurizer(self, seqs, **kwargs):
+        """
+        Applies amino acid index encoding to sequences.
+
+        Args:
+            seqs (list): List of sequences to encode.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: Amino acid index encoded sequences.
+        """
+        X = featurize_aa_idx(seqs)
+
+        return X

model/featurizers/combinatorial_featurizers.py

@@ -0,0 +1,425 @@
+from model.featurizers.base_featurizers import *
+from model.featurizers.esm_featurizers import *
+from model.featurizers.msa_featurizers import *
+from model.featurizers.zeroshot_featurizers import *
+
+from model.featurizers.model_choices import FEATURIZE_CHOICES
+
+FEATURIZE_CLASSES = {
+    # Dictionary of model names to model classes.
+    # Base Featurizers
+    "onehot": OneHotFeaturizer,
+    "georgiev": GeorgievFeaturizer,
+    "aa_idx": AAIdxFeaturizer,
+    # MSA Featurizers
+    "msa_embed": MSAEmbedFeaturizer,
+    "msa_sequence_embed": MSASequenceEmbedFeaturizer,
+    "msa_logits": MSALogitsFeaturizer,
+    # ESM Featurizers
+    "esm_logits": ESMLogitsFeaturizer,
+    "esm_embed_1v": ESM1vEmbedFeaturizer,
+    "esm_embed_2_3b": ESM2EmbedFeaturizer,
+    "esm_embed_2_15b": ESM2_15b_EmbedFeaturizer,
+    # Zeroshot Featurizers
+    "zeroshot_msa": ZeroshotMSAFeaturizer,
+    "zeroshot_esm": ZeroshotESMFeaturizer,
+    "zeroshot_prose": ZeroshotProseFeaturizer,
+    "zeroshot_cscs": ZeroshotCSCSFeaturizer,
+    "zeroshot_cscs_gram": ZeroshotCSCSGramFeaturizer,
+    "zeroshot_cscs_sem": ZeroshotCSCSSemFeaturizer,
+}
+
+
+class CombinatorialFeaturizer():
+    """Base class for combining multiple featurizers.
+
+    Attributes:
+        name (str): Name of the combined featurizer.
+        featurizers (dict): Dictionary mapping featurizer names to instances.
+
+    Example Usage:
+    
+    featurizer = CombinatorialFeaturizer(
+        featurize_methods=['onehot', 'georgiev'],  # List of featurizers to combine
+        protein='protein1',                        # Name of protein for caching
+        use_cache=True                            # Whether to cache results
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(self, featurize_methods, **kwargs):
+        """
+        Args:
+            featurize_methods (list): List of featurizer names to combine.
+            **kwargs: Additional arguments passed to each featurizer.
+        """
+        for featurize_method in featurize_methods:
+            assert (
+                featurize_method in FEATURIZE_CHOICES
+            ), f"{featurize_method} not in {FEATURIZE_CHOICES}"
+
+        model_type = "-".join(featurize_methods)
+        self.name = str(model_type)
+
+        self.featurizers = {
+            featurize_method: FEATURIZE_CLASSES[featurize_method](**kwargs)
+            for featurize_method in featurize_methods
+        }
+
+    def featurize(self, seqs, **kwargs):
+        """
+        Featurizes sequences using all component featurizers.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional arguments passed to each featurizer.
+
+        Returns:
+            np.ndarray: Combined features from all featurizers.
+        """
+        X = []
+        for featurizer in self.featurizers.values():
+            X.append(featurizer.featurize(seqs, **kwargs))
+
+        X = np.concatenate(X, axis=-1)
+
+        return X
+    
+class ESMAugmentedFeaturizer(CombinatorialFeaturizer):
+    """Class for combining ESM features with one-hot encoding.
+
+    Attributes:
+        name (str): Name of the combined featurizer.
+        featurizers (dict): Dictionary mapping featurizer names to instances.
+
+    Example Usage:
+    
+    featurizer = ESMAugmentedFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True           # Whether to cache results
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(self, featurize_methods=["zeroshot_esm", "onehot"], **kwargs):
+        """
+        Args:
+            featurize_methods (list): List of featurizer names to combine.
+            **kwargs: Additional arguments passed to each featurizer.
+        """
+        super().__init__(featurize_methods, **kwargs)
+    
+    def featurize(self, seqs, **kwargs):
+        """
+        Featurizes sequences using ESM and one-hot encoding.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional arguments passed to each featurizer.
+
+        Returns:
+            np.ndarray: Combined ESM and one-hot features.
+        """
+        X = []
+        
+        featurizer_0 = list(self.featurizers.values())[0]
+        X.append(featurizer_0.featurize(seqs, **kwargs))
+        
+        featurizer_1 = list(self.featurizers.values())[1]
+        onehot = featurizer_1.featurize(seqs, **kwargs)
+        X.append(onehot.reshape(onehot.shape[0], -1))
+        
+        X = np.concatenate(X, axis=1)
+
+        return X
+
+class MSAAugmentedFeaturizer(CombinatorialFeaturizer):
+    """Class for combining MSA features with one-hot encoding.
+
+    Attributes:
+        name (str): Name of the combined featurizer.
+        featurizers (dict): Dictionary mapping featurizer names to instances.
+
+    Example Usage:
+    
+    featurizer = MSAAugmentedFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True           # Whether to cache results
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(self, featurize_methods=["zeroshot_msa", "onehot"], **kwargs):
+        """
+        Args:
+            featurize_methods (list): List of featurizer names to combine.
+            **kwargs: Additional arguments passed to each featurizer.
+        """
+        super().__init__(featurize_methods, **kwargs)
+
+    def featurize(self, seqs, **kwargs):
+        """
+        Featurizes sequences using MSA and one-hot encoding.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional arguments passed to each featurizer.
+
+        Returns:
+            np.ndarray: Combined MSA and one-hot features.
+        """
+        X = []
+        
+        featurizer_0 = list(self.featurizers.values())[0]
+        X.append(featurizer_0.featurize(seqs, **kwargs))
+        
+        featurizer_1 = list(self.featurizers.values())[1]
+        onehot = featurizer_1.featurize(seqs, **kwargs)
+        X.append(onehot.reshape(onehot.shape[0], -1))
+        
+        X = np.concatenate(X, axis=1)
+
+        return X
+
+class OnehotAndGeorgievFeaturizer(CombinatorialFeaturizer):
+    """Class for combining one-hot and Georgiev encodings.
+
+    Attributes:
+        name (str): Name of the combined featurizer.
+        featurizers (dict): Dictionary mapping featurizer names to instances.
+
+    Example Usage:
+    
+    featurizer = OnehotAndGeorgievFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True           # Whether to cache results
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(self, featurize_methods=["onehot", "georgiev"], **kwargs):
+        """
+        Args:
+            featurize_methods (list): List of featurizer names to combine.
+            **kwargs: Additional arguments passed to each featurizer.
+        """
+        super().__init__(featurize_methods, **kwargs)
+
+class OnehotAndAAIdxFeaturizer(CombinatorialFeaturizer):
+    """Class for combining one-hot and amino acid index encodings.
+
+    Attributes:
+        name (str): Name of the combined featurizer.
+        featurizers (dict): Dictionary mapping featurizer names to instances.
+
+    Example Usage:
+    
+    featurizer = OnehotAndAAIdxFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True           # Whether to cache results
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(self, featurize_methods=["onehot", "aa_idx"], **kwargs):
+        """
+        Args:
+            featurize_methods (list): List of featurizer names to combine.
+            **kwargs: Additional arguments passed to each featurizer.
+        """
+        super().__init__(featurize_methods, **kwargs)      
+
+class OnehotAndESMLogitsFeaturizer(CombinatorialFeaturizer):
+    """Class for combining one-hot encoding with ESM logits.
+
+    Attributes:
+        name (str): Name of the combined featurizer.
+        featurizers (dict): Dictionary mapping featurizer names to instances.
+
+    Example Usage:
+    
+    featurizer = OnehotAndESMLogitsFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True           # Whether to cache results
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(self, featurize_methods=["onehot", "esm_logits"], **kwargs):
+        """
+        Args:
+            featurize_methods (list): List of featurizer names to combine.
+            **kwargs: Additional arguments passed to each featurizer.
+        """
+        super().__init__(featurize_methods, **kwargs)
+    
+    def featurize(self, seqs, **kwargs):
+        """
+        Featurizes sequences using one-hot encoding and ESM logits.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional arguments passed to each featurizer.
+
+        Returns:
+            np.ndarray: Combined one-hot and ESM logits features.
+        """
+        X = []
+        
+        featurizer_0 = list(self.featurizers.values())[0]
+        x = featurizer_0.featurize(seqs, **kwargs)
+        zero_vectors = np.zeros((x.shape[0], 1, x.shape[2]))
+        X.append(np.concatenate((zero_vectors, x, zero_vectors), axis=1))
+        
+        featurizer_1 = list(self.featurizers.values())[1]
+        X.append(featurizer_1.featurize(seqs, **kwargs))
+        
+        X = np.concatenate(X, axis=-1)
+
+        return X
+
+class OnehotAndESMMSALogitsFeaturizer(CombinatorialFeaturizer):
+    """Class for combining one-hot encoding with ESM-MSA logits.
+
+    Attributes:
+        name (str): Name of the combined featurizer.
+        featurizers (dict): Dictionary mapping featurizer names to instances.
+
+    Example Usage:
+    
+    featurizer = OnehotAndESMMSALogitsFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True           # Whether to cache results
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(self, featurize_methods=["onehot", "msa_logits"], **kwargs):
+        """
+        Args:
+            featurize_methods (list): List of featurizer names to combine.
+            **kwargs: Additional arguments passed to each featurizer.
+        """
+        super().__init__(featurize_methods, **kwargs)
+    
+    def featurize(self, seqs, **kwargs):
+        """
+        Featurizes sequences using one-hot encoding and ESM-MSA logits.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional arguments passed to each featurizer.
+
+        Returns:
+            np.ndarray: Combined one-hot and ESM-MSA logits features.
+        """
+        X = []
+        
+        featurizer_0 = list(self.featurizers.values())[0]
+        x = featurizer_0.featurize(seqs, **kwargs)
+        zero_vectors = np.zeros((x.shape[0], 1, x.shape[2]))
+        X.append(np.concatenate((zero_vectors, x), axis=1))
+        
+        featurizer_1 = list(self.featurizers.values())[1]
+        X.append(featurizer_1.featurize(seqs, **kwargs))
+        
+        X = np.concatenate(X, axis=-1)
+
+        return X
+    
+
+class OnehotAndESM2EmbedFeaturizer(CombinatorialFeaturizer):
+    """Class for combining one-hot encoding with ESM2 embeddings.
+
+    Attributes:
+        name (str): Name of the combined featurizer.
+        featurizers (dict): Dictionary mapping featurizer names to instances.
+
+    Example Usage:
+    
+    featurizer = OnehotAndESM2EmbedFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True           # Whether to cache results
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(self, featurize_methods=["onehot", "esm_embed_2_3b"], **kwargs):
+        """
+        Args:
+            featurize_methods (list): List of featurizer names to combine.
+            **kwargs: Additional arguments passed to each featurizer.
+        """
+        super().__init__(featurize_methods, **kwargs)
+
+    def featurize(self, seqs, **kwargs):
+        """
+        Featurizes sequences using one-hot encoding and ESM2 embeddings.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional arguments passed to each featurizer.
+
+        Returns:
+            np.ndarray: Combined one-hot and ESM2 embedding features.
+        """
+        X = []
+        
+        featurizer_0 = list(self.featurizers.values())[0]
+        onehot = featurizer_0.featurize(seqs, **kwargs)
+        X.append(onehot.reshape(onehot.shape[0], -1))
+        
+        featurizer_1 = list(self.featurizers.values())[1]
+        X.append(featurizer_1.featurize(seqs, **kwargs))
+        
+        X = np.concatenate(X, axis=1)
+
+        return X
+
+class OnehotAndESM2_15bEmbedFeaturizer(CombinatorialFeaturizer):
+    """Class for combining one-hot encoding with ESM2 embeddings.
+
+    Attributes:
+        name (str): Name of the combined featurizer.
+        featurizers (dict): Dictionary mapping featurizer names to instances.
+
+    Example Usage:
+    
+    featurizer = OnehotAndESM2EmbedFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True           # Whether to cache results
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(self, featurize_methods=["onehot", "esm_embed_2_15b"], **kwargs):
+        """
+        Args:
+            featurize_methods (list): List of featurizer names to combine.
+            **kwargs: Additional arguments passed to each featurizer.
+        """
+        super().__init__(featurize_methods, **kwargs)
+
+    def featurize(self, seqs, **kwargs):
+        """
+        Featurizes sequences using one-hot encoding and ESM2 embeddings.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional arguments passed to each featurizer.
+
+        Returns:
+            np.ndarray: Combined one-hot and ESM2 embedding features.
+        """
+        X = []
+        
+        featurizer_0 = list(self.featurizers.values())[0]
+        onehot = featurizer_0.featurize(seqs, **kwargs)
+        X.append(onehot.reshape(onehot.shape[0], -1))
+        
+        featurizer_1 = list(self.featurizers.values())[1]
+        X.append(featurizer_1.featurize(seqs, **kwargs))
+        
+        X = np.concatenate(X, axis=1)
+
+        return X

model/featurizers/esm_featurizers.py

@@ -0,0 +1,620 @@
+import numpy as np
+import torch
+
+from model.featurizers.model_choices import FEATURE_MODELS
+from model.featurizers.base_featurizers import BaseFeaturizer
+from concurrent.futures import ProcessPoolExecutor
+import concurrent.futures
+
+class ForgeESMFeaturizer(BaseFeaturizer):
+    """Class for generating ESM Forge-based protein embeddings or log probabilities.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        model (str): ESM model name.
+        url (str): Forge API URL.
+        token (str): API token.
+        output_type (str): Output format type.
+
+    Example Usage:
+
+    featurizer = ForgeESMFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        flatten_features=False,   # Whether to flatten output features
+        model='esm2_t33_650M',   # ESM model to use
+        token='api_token',       # Forge API token
+        output_type='sequence_representations'  # Output type
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        flatten_features=False,
+        model=None,
+        url="https://forge.evolutionaryscale.ai",
+        token=None,
+        output_type=None,
+        model_type="esm",
+        **kwargs,
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            flatten_features (bool): Whether to flatten output features.
+            model (str): ESM model name.
+            url (str): Forge API URL.
+            token (str): API token.
+            output_type (str): Output format type.
+            model_type (str): Type of featurization model.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(model_type=model_type, protein=protein, use_cache=use_cache, flatten_features=flatten_features, **kwargs)
+        self.model = model
+        self.url = url
+        if token is None:
+            raise ValueError("ESM Forge API token must be provided")
+        self.token = token
+        self.output_type = output_type
+
+    def process_single_protein(self, sequence, model, url, token, output_type):
+        """Process a single protein sequence using ESM Forge.
+        
+        Args:
+            sequence (str): Protein sequence to process.
+            model (str): ESM model name.
+            url (str): Forge API URL.
+            token (str): Forge API access token.
+            output_type (str): Either "log_probabilities" or "sequence_representations".
+            
+        Returns:
+            numpy.ndarray: Protein embeddings or log probabilities.
+            
+        Raises:
+            ValueError: If output_type is invalid.
+            RuntimeError: If API call fails.
+        """
+        if output_type not in ["log_probabilities", "sequence_representations"]:
+            raise ValueError("output_type must be 'log_probabilities' or 'sequence_representations'")
+         
+        try:
+            
+            from esm.sdk.forge import ESM3ForgeInferenceClient
+            from esm.sdk.api import ESMProtein, LogitsConfig
+
+            forge_client = ESM3ForgeInferenceClient(model=model, url=url, token=token)
+            protein = ESMProtein(sequence=sequence)
+            protein_tensor = forge_client.encode(protein)
+            logits_output = forge_client.logits(
+                protein_tensor, LogitsConfig(sequence=True, return_embeddings=True)
+            )
+        except Exception as e:
+            raise RuntimeError(f"ESM Forge API call failed: {str(e)}")
+
+        if output_type == "log_probabilities":
+            return np.array(logits_output.logits.sequence.numpy())
+        elif output_type == "sequence_representations":
+            embeddings_float32 = logits_output.embeddings.squeeze().float()
+            return (np.array(embeddings_float32.numpy()).mean(axis=0))
+        
+
+    def process_proteins_parallel(self, seqs, model, url, token, output_type):
+        """Process a list of protein sequences in parallel using ESM Forge.
+        
+        Args:
+            seqs (list): List of protein sequences to process.
+            model (str): ESM model name to use.
+            url (str): Forge API URL.
+            token (str): Forge API access token.
+            output_type (str): Type of output features.
+            
+        Returns:
+            list: List of protein embeddings as numpy arrays.
+        """
+        with ProcessPoolExecutor(max_workers=16) as executor:
+            future_to_index = {
+                executor.submit(self.process_single_protein, seq, model, url, token, output_type): i 
+                for i, seq in enumerate(seqs)
+            }
+            
+            results = [None] * len(seqs)
+            
+            for future in concurrent.futures.as_completed(future_to_index):
+                index = future_to_index[future]
+                results[index] = future.result()
+                
+        return results
+        
+    def custom_featurizer(self, seqs, **kwargs):
+        """
+        Featurizes sequences using ESM Forge.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: Array of featurized sequences.
+        """
+        X = self.process_proteins_parallel(seqs, model=self.model, url=self.url, token=self.token, output_type=self.output_type)
+        return X
+
+
+class Forge_ESMC_6B_EmbedFeaturizer(ForgeESMFeaturizer):
+    """Class for generating ESM-C 6B embeddings using Forge.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        model (str): ESM model name.
+        url (str): Forge API URL.
+        token (str): API token.
+        output_type (str): Output format type.
+
+    Example Usage:
+
+    featurizer = Forge_ESMC_6B_EmbedFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        flatten_features=False,   # Whether to flatten output features
+        token='api_token'        # Forge API token
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        flatten_features=False,
+        model="esmc-6b-2024-12",
+        url="https://forge.evolutionaryscale.ai",
+        token=None,
+        output_type="sequence_representations",
+        **kwargs,
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            flatten_features (bool): Whether to flatten output features.
+            model (str): ESM model name.
+            url (str): Forge API URL.
+            token (str): API token.
+            output_type (str): Output format type.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(model_type="esmc_6b", protein=protein, use_cache=use_cache, flatten_features=flatten_features, output_type=output_type, model=model, url=url, token=token, **kwargs)
+
+
+class ESMBaseFeaturizer(BaseFeaturizer):
+    """Base class for ESM model featurizers.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        batch_size (int): Batch size for processing.
+        model_locations (list): Paths to ESM model files.
+
+    Example Usage:
+
+    featurizer = ESMBaseFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        batch_size=968,          # Processing batch size
+        model_locations=['path/to/model']  # Model file paths
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        model_locations=None,
+        batch_size=968,
+        model_type="esm",
+        **kwargs,
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            model_locations (list): Paths to ESM model files.
+            batch_size (int): Batch size for processing.
+            model_type (str): Type of featurization model.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(model_type, protein, use_cache, **kwargs)
+        self.batch_size = batch_size
+        self.model_locations = model_locations
+
+    def eval_esm(self, model, batch_tokens, sequence_data, output_type):
+        """
+        Evaluates sequences using ESM model.
+
+        Args:
+            model (torch.nn.Module): ESM model.
+            batch_tokens (torch.Tensor): Tokenized sequences.
+            sequence_data (list): Original sequence data.
+            output_type (str): Type of output features.
+
+        Returns:
+            np.ndarray: Model outputs.
+        """
+        if output_type == "log_probabilities":
+            # Featurize as sequence log probabilities.
+            with torch.no_grad():
+                token_probs = torch.log_softmax(
+                    model(batch_tokens.to(self.device))["logits"], dim=-1
+                )
+            output = token_probs.cpu().numpy()
+
+        elif output_type == "sequence_representations":
+            # Featurize as sequence embeddings (last hidden layer).
+            last_layer = len(model.layers)
+            with torch.no_grad():
+                results = model(
+                    batch_tokens.to(self.device), repr_layers=[last_layer]
+                )
+            token_representations = results["representations"][last_layer]
+
+            sequence_representations = []
+            for i, (_, seq) in enumerate(sequence_data):
+                seq_embed = token_representations[i, 1 : len(seq) + 1].mean(0)
+                sequence_representations.append(seq_embed.cpu().numpy().ravel())
+
+            output = np.array(sequence_representations)
+
+        return output
+
+    def featurize_esm(self, seqs, output_type):
+        """
+        Featurizes sequences using ESM model.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            output_type (str): Type of output features.
+
+        Returns:
+            np.ndarray: Array of featurized sequences.
+        """
+        from esm import pretrained
+        model_loc_to_model = {}
+        for model_location in self.model_locations:
+            model, alphabet = pretrained.load_model_and_alphabet(model_location)
+            model.eval()
+            if torch.backends.mps.is_available() or torch.cuda.is_available():
+                model = model.to(self.device)
+            else:
+                print("GPU device not available")
+                return
+            model_loc_to_model[model_location] = model
+
+        output = []
+
+        batch_size = self.batch_size
+        n_batches = ((len(seqs[0]) - 1) // batch_size) + 1
+        for batchi in range(n_batches):
+            start = batchi * batch_size
+            end = (batchi + 1) * batch_size
+
+            model_features = []
+            for model_location in self.model_locations:
+                model = model_loc_to_model[model_location]
+
+                seq_batch = []
+                sbatch_size = 3
+                n_sbatches = ((len(seqs) - 1) // sbatch_size) + 1
+                for batchj in range(n_sbatches):
+                    sb_start = batchj * sbatch_size
+                    sb_end = (batchj + 1) * sbatch_size
+
+                    sequence_data = [
+                        (f"protein{sbidx}", seq[start:end])
+                        if output_type == "log_probabilities"
+                        else (f"protein{sbidx}", seq[start:end].replace("X", ""))
+                        for sbidx, seq in enumerate(seqs[sb_start:sb_end])
+                    ]
+
+                    batch_converter = alphabet.get_batch_converter()
+                    batch_labels, batch_strs, batch_tokens = batch_converter(
+                        sequence_data
+                    )
+
+                    if n_batches > 1:
+                        if batchi == 0:
+                            batch_tokens = batch_tokens[:, :-1]
+                        elif batchi == n_batches - 1:
+                            batch_tokens = batch_tokens[:, 1:]
+                        else:
+                            batch_tokens = batch_tokens[:, 1:-1]
+
+                    seq_batch.append(
+                        self.eval_esm(model, batch_tokens, sequence_data, output_type)
+                    )
+
+                model_features.append(np.concatenate(seq_batch))
+
+            output.append(np.mean(model_features, axis=0))
+
+        X = np.hstack(output)
+
+        return X
+
+
+class ESMLogitsFeaturizer(ESMBaseFeaturizer):
+    """Class for generating ESM model log probabilities.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        batch_size (int): Batch size for processing.
+        model_locations (list): Paths to ESM model files.
+
+    Example Usage:
+
+    featurizer = ESMLogitsFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        batch_size=968           # Processing batch size
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        model_locations=FEATURE_MODELS["esm_logits"],
+        batch_size=968,
+        model_type="esm_logits",
+        **kwargs,
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            model_locations (list): Paths to ESM model files.
+            batch_size (int): Batch size for processing.
+            model_type (str): Type of featurization model.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(
+            protein,
+            use_cache,
+            model_locations,
+            batch_size,
+            model_type,
+            **kwargs,
+        )
+
+    def custom_featurizer(self, seqs, **kwargs):
+        """
+        Featurizes sequences using ESM model log probabilities.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: Array of log probabilities.
+        """
+        X = self.featurize_esm(seqs, output_type="log_probabilities")
+        return X
+
+
+class ESM1vEmbedFeaturizer(ESMBaseFeaturizer):
+    """Class for generating ESM-1v model embeddings.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        batch_size (int): Batch size for processing.
+        model_locations (list): Paths to ESM model files.
+
+    Example Usage:
+
+    featurizer = ESM1vEmbedFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        batch_size=968           # Processing batch size
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        model_locations=FEATURE_MODELS["esm_embed_1v"],
+        batch_size=968,
+        model_type="esm_embed_1v",
+        **kwargs,
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            model_locations (list): Paths to ESM model files.
+            batch_size (int): Batch size for processing.
+            model_type (str): Type of featurization model.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(
+            protein,
+            use_cache,
+            model_locations,
+            batch_size,
+            model_type,
+            **kwargs,
+        )
+
+    def custom_featurizer(self, seqs, **kwargs):
+        """
+        Featurizes sequences using ESM-1v model embeddings.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: Array of sequence embeddings.
+        """
+        X = self.featurize_esm(seqs, output_type="sequence_representations")
+        return X
+
+
+class ESM2EmbedFeaturizer(ESMBaseFeaturizer):
+    """Class for generating ESM-2 model embeddings.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        batch_size (int): Batch size for processing.
+        model_locations (list): Paths to ESM model files.
+
+    Example Usage:
+
+    featurizer = ESM2EmbedFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        batch_size=968           # Processing batch size
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        model_locations=FEATURE_MODELS["esm_embed_2_3b"],
+        batch_size=968,
+        model_type="esm_embed_2_3b",
+        **kwargs,
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            model_locations (list): Paths to ESM model files.
+            batch_size (int): Batch size for processing.
+            model_type (str): Type of featurization model.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(
+            protein,
+            use_cache,
+            model_locations,
+            batch_size,
+            model_type,
+            **kwargs,
+        )
+
+    def custom_featurizer(self, seqs, **kwargs):
+        """
+        Featurizes sequences using ESM-2 model embeddings.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: Array of sequence embeddings.
+        """
+        X = self.featurize_esm(seqs, output_type="sequence_representations")
+        return X
+
+
+class ESM2_15b_EmbedFeaturizer(ESMBaseFeaturizer):
+    """Class for generating ESM-2 15B model embeddings.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        batch_size (int): Batch size for processing.
+        model_locations (list): Paths to ESM model files.
+
+    Example Usage:
+
+    featurizer = ESM2_15b_EmbedFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        batch_size=968           # Processing batch size
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        model_locations=FEATURE_MODELS["esm_embed_2_15b"],
+        batch_size=968,
+        model_type="esm_embed_2_15b",
+        **kwargs,
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            model_locations (list): Paths to ESM model files.
+            batch_size (int): Batch size for processing.
+            model_type (str): Type of featurization model.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(
+            protein,
+            use_cache,
+            model_locations,
+            batch_size,
+            model_type,
+            **kwargs,
+        )
+
+    def custom_featurizer(self, seqs, **kwargs):
+        """
+        Featurizes sequences using ESM-2 15B model embeddings.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: Array of sequence embeddings.
+        """
+        X = self.featurize_esm(seqs, output_type="sequence_representations")
+        return X

model/featurizers/model_choices.py

@@ -0,0 +1,70 @@
+from model.featurizers.model_locations import (
+    msa_models,
+    esm_models,
+    prose_models,
+    prose_models_cas13,
+)
+
+FEATURIZE_CHOICES = [
+    # All the implemented featurizers
+    # Base Featurizers
+    'onehot',
+    'georgiev',
+    'aa_idx',
+
+    # MSA Featurizers
+    'msa_embed',
+    'msa_sequence_embed',
+    'msa_logits',
+    'msa_augmented',
+
+    # ESM Featurizers
+    'esm_logits',
+    'esm_embed_1v',
+    'esm_embed_2_3b',
+    'esm_embed_2_15b',
+    'esm_augmented',
+
+    # Zeroshot Featurizers
+    'zeroshot_msa',
+    'zeroshot_esm',
+    'zeroshot_cscs',
+    'zeroshot_cscs_gram',
+    'zeroshot_cscs_sem',
+    'zeroshot_prose',
+    'zeroshot_esmif',
+]
+
+FEATURE_MODELS = {
+    # Dictionary of model names to model locations.
+    # Base Featurizers
+    'onehot': [ None ],
+    'georgiev': [ None ],
+    'aa_idx': [ None ],
+
+    # MSA Featurizers
+    'msa_embed': msa_models,
+    'msa_sequence_embed': msa_models,
+    'msa_logits': msa_models,
+    'msa_augmented': msa_models,
+
+    # ESM Featurizers
+    'esm_logits': esm_models[:1],
+    'esm_embed_1v': esm_models[:1],
+    'esm_embed_2_3b': esm_models[5:6],
+    'esm_embed_2_15b': esm_models[7:8],
+    'esm_augmented': esm_models[:5],
+
+    # Zeroshot Featurizers
+    'zeroshot_msa': msa_models,
+    'zeroshot_esm': esm_models[:6],
+    'zeroshot_cscs': [ None ],
+    'zeroshot_cscs_gram': [ None ],
+    'zeroshot_cscs_sem': [ None ],
+    'zeroshot_prose': prose_models_cas13,
+    'zeroshot_esmif': esm_models[6:7],
+
+    # Prose Featurizers
+    'prose_embedmean': prose_models,
+    'prose_augmented': prose_models_cas13,
+}

model/featurizers/model_locations.py

@@ -0,0 +1,38 @@
+esm_models = [
+    'esm1v_t33_650M_UR90S_1',
+    'esm1v_t33_650M_UR90S_2',
+    'esm1v_t33_650M_UR90S_3',
+    'esm1v_t33_650M_UR90S_4',
+    'esm1v_t33_650M_UR90S_5',
+    'esm2_t36_3B_UR50D',
+    'esm_if1_gvp4_t16_142M_UR50',
+    'esm2_t48_15B_UR50D',
+]
+
+msa_models = [
+    'esm_msa1b_t12_100M_UR50S',
+]
+
+prose_models = [
+    'data/prose_pretrained_models/prose_dlm_3x1024.sav',
+]
+
+prose_models_cas13 = [
+    'target/cas13/prose_rd512_nl3_dr0_ns2000000_si100000_le1280_mr0.1_bs100_wd0_lr0.0001_cs0.98_0/_iter0500000.sav',
+    'target/cas13/prose_rd512_nl3_dr0_ns2000000_si100000_le1280_mr0.1_bs100_wd0_lr0.0001_cs0.98_1/_iter0400000.sav',
+    'target/cas13/prose_rd512_nl3_dr0_ns2000000_si100000_le1280_mr0.1_bs100_wd0_lr0.0001_cs0.98_3/_iter0600000.sav',
+    'target/cas13/prose_rd512_nl3_dr0_ns2000000_si100000_le1280_mr0.1_bs100_wd0_lr0.0001_cs0.98_4/_iter0600000.sav',
+    'target/cas13/prose_rd512_nl3_dr0_ns2000000_si100000_le1280_mr0.1_bs100_wd0_lr0.0001_cs0.98_5/_iter0600000.sav',
+]
+
+prose_models_cas13_old1 = [
+    'target/cas13/prose_rd512_nl3_dr0_ns2000000_si100000_le1280_mr0.1_bs100_wd0_lr0.0001_cs0.98_0/_iter0500000.sav',
+    'target/cas13/prose_rd512_nl3_dr0_ns2000000_si100000_le1280_mr0.1_bs100_wd0_lr0.0001_cs0.98_1/_iter0400000.sav',
+    'target/cas13/prose_rd512_nl3_dr0_ns2000000_si100000_le1280_mr0.1_bs100_wd0_lr0.0001_cs0.98_2/_iter0300000.sav',
+]
+
+prose_models_cas13_old = [
+    'target/cas13_old/prose_rd512_nl3_dr0_ns2000000_si100000_le1280_mr0.1_bs100_wd0_lr0.0001_cs0.98_1/_iter0300000.sav',
+    'target/cas13_old/prose_rd512_nl3_dr0_ns2000000_si100000_le1280_mr0.1_bs100_wd0_lr0.0001_cs0.98_2/_iter0200000.sav',
+    'target/cas13_old/prose_rd512_nl3_dr0_ns2000000_si100000_le1280_mr0.1_bs100_wd0_lr0.0001_cs0.98_3/_iter0300000.sav',
+]

model/featurizers/msa_featurizers.py

@@ -0,0 +1,381 @@
+from Bio import SeqIO
+import numpy as np
+import torch
+
+from model.featurizers.base_featurizers import BaseFeaturizer
+from model.featurizers.model_choices import FEATURE_MODELS
+from model.utils.other_utils import read_msa, greedy_select, msa_splicer
+
+
+class MSABaseFeaturizer(BaseFeaturizer):
+    """Base class for MSA-based featurizers.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        batch_size (int): Batch size for processing.
+        model_locations (list): Paths to model files.
+        msa_file (str): Path to MSA file.
+
+    Example Usage:
+
+    featurizer = MSABaseFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        msa_file='msa.fasta',    # Path to MSA file
+        batch_size=968           # Processing batch size
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        model_locations=None,
+        msa_file=None,
+        batch_size=968,
+        model_type="msa",
+        **kwargs
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            model_locations (list): Paths to model files.
+            msa_file (str): Path to MSA file.
+            batch_size (int): Batch size for processing.
+            model_type (str): Type of featurization model.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(model_type, protein, use_cache, **kwargs)
+        self.batch_size = batch_size
+        self.model_locations = model_locations
+        self.msa_file = msa_file
+        self.device = torch.device("cpu")  # MSAs might be too big for GPU
+
+        
+
+    def featurize_msa(self, seqs, msa_file, output_type, **kwargs):
+        """
+        Featurizes sequences using MSA Transformer model.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            msa_file (str): Path to MSA file.
+            output_type (str): Type of output features to extract.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: Array of sequence features.
+        """
+        from esm import pretrained
+        
+        torch.set_grad_enabled(False)
+        # Check to see if there is an MSA file in **kwargs.
+        if msa_file is None:
+            assert self.msa_file is not None, "No MSA file provided."
+            msa_file = self.msa_file
+        assert output_type in [
+            "sequence_representations",
+            "msa_representations",
+            "log_probabilities",
+        ]
+        msa = read_msa(msa_file)
+
+        # Instantiate the model
+        (
+            msa_transformer,
+            msa_transformer_alphabet,
+        ) = pretrained.esm_msa1b_t12_100M_UR50S()
+        msa_transformer = msa_transformer.eval().to(self.device)
+        msa_transformer_batch_converter = msa_transformer_alphabet.get_batch_converter()
+
+        # Prep the MSA, making the appropriate mutations
+        inputs = greedy_select(
+            msa, num_seqs=128
+        )  # can change this to pass more/fewer sequences
+        # This splices the MSA to exclude gaps in the first sequence, due to MSATransformer context window
+        # size limit of 1024. If your MSA width is less than 1024, then you don't need to do this
+        inputs = msa_splicer(inputs)
+        name, wt_seq = inputs[0][0], inputs[0][1]
+        reps = []
+
+        # Batch processing
+        batch_size = self.batch_size
+        num_batches = len(seqs) // batch_size + (len(seqs) % batch_size != 0)
+        for batch_i in range(num_batches):  # TODO: change this to num_batches
+            start_idx = batch_i * batch_size
+            end_idx = start_idx + batch_size
+            batch_seqs = seqs[start_idx:end_idx]
+
+            input_msas = []
+            for seq in batch_seqs:
+                # Replace the first sequence in the MSA with the mutant sequence
+                assert len(wt_seq) == len(seq)
+                inputs[0] = (name, seq)
+                input_msas.append(inputs)
+            # print(len(input_msas), len(input_msas[0]), len(input_msas[0][0][1]), len(input_msas[0][1][1]))
+            # Run the MSA Transformer
+            (
+                msa_transformer_batch_labels,
+                msa_transformer_batch_strs,
+                msa_transformer_batch_tokens,
+            ) = msa_transformer_batch_converter(input_msas)
+            num_msas = len(msa_transformer_batch_tokens)
+            msa_transformer_batch_tokens = msa_transformer_batch_tokens.to(
+                next(msa_transformer.parameters()).device
+            )
+            msa_transformer_predictions = msa_transformer.forward(
+                msa_transformer_batch_tokens, repr_layers=[12]
+            )
+
+            # Extract features
+            msa_reps = (
+                msa_transformer_predictions["representations"][12]
+                .detach()
+                .cpu()
+                .numpy()
+            )
+            for i in range(num_msas):
+                logits = (
+                    msa_transformer_predictions["logits"][i][0].detach().cpu().numpy()
+                )
+                avg_msa_rep = msa_reps[i].mean((0, 1))
+                sequence_rep = msa_reps[i][0].mean(0)
+
+                if output_type == "msa_representations":
+                    reps.append(avg_msa_rep)
+
+                elif output_type == "sequence_representations":
+                    reps.append(sequence_rep)
+
+                elif output_type == "log_probabilities":
+                    reps.append(logits)
+
+        X = np.array(reps)
+    
+        return X
+
+
+class MSAEmbedFeaturizer(MSABaseFeaturizer):
+    """Class for generating MSA embedding features.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        batch_size (int): Batch size for processing.
+        model_locations (list): Paths to model files.
+        msa_file (str): Path to MSA file.
+
+    Example Usage:
+
+    featurizer = MSAEmbedFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        msa_file='msa.fasta',    # Path to MSA file
+        batch_size=968           # Processing batch size
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        msa_file=None,
+        model_locations=FEATURE_MODELS["msa_embed"],
+        batch_size=968,
+        model_type="msa_embed",
+        **kwargs
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            msa_file (str): Path to MSA file.
+            model_locations (list): Paths to model files.
+            batch_size (int): Batch size for processing.
+            model_type (str): Type of featurization model.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(
+            protein,
+            use_cache,
+            model_locations,
+            msa_file,
+            batch_size,
+            model_type,
+            **kwargs
+        )
+
+    def custom_featurizer(self, seqs, msa_file=None, **kwargs):
+        """
+        Featurizes sequences using MSA embeddings.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            msa_file (str): Path to MSA file.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: Array of MSA embedding features.
+        """
+        X = self.featurize_msa(
+            seqs, msa_file, output_type="msa_representations", **kwargs
+        )
+        return X
+
+
+class MSASequenceEmbedFeaturizer(MSABaseFeaturizer):
+    """Class for generating MSA sequence embedding features.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        batch_size (int): Batch size for processing.
+        model_locations (list): Paths to model files.
+        msa_file (str): Path to MSA file.
+
+    Example Usage:
+
+    featurizer = MSASequenceEmbedFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        msa_file='msa.fasta',    # Path to MSA file
+        batch_size=968           # Processing batch size
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        msa_file=None,
+        model_locations=FEATURE_MODELS["msa_sequence_embed"],
+        batch_size=968,
+        model_type="msa_sequence_embed",
+        **kwargs
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            msa_file (str): Path to MSA file.
+            model_locations (list): Paths to model files.
+            batch_size (int): Batch size for processing.
+            model_type (str): Type of featurization model.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(
+            protein,
+            use_cache,
+            model_locations,
+            msa_file,
+            batch_size,
+            model_type,
+            **kwargs
+        )
+
+    def custom_featurizer(self, seqs, msa_file=None, **kwargs):
+        """
+        Featurizes sequences using MSA sequence embeddings.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            msa_file (str): Path to MSA file.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: Array of MSA sequence embedding features.
+        """
+        X = self.featurize_msa(
+            seqs, msa_file, output_type="sequence_representations", **kwargs
+        )
+        return X
+
+
+class MSALogitsFeaturizer(MSABaseFeaturizer):
+    """Class for generating MSA logits features.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        batch_size (int): Batch size for processing.
+        model_locations (list): Paths to model files.
+        msa_file (str): Path to MSA file.
+
+    Example Usage:
+
+    featurizer = MSALogitsFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        msa_file='msa.fasta',    # Path to MSA file
+        batch_size=968           # Processing batch size
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        msa_file=None,
+        model_locations=FEATURE_MODELS["msa_logits"],
+        batch_size=968,
+        model_type="msa_logits",
+        **kwargs
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            msa_file (str): Path to MSA file.
+            model_locations (list): Paths to model files.
+            batch_size (int): Batch size for processing.
+            model_type (str): Type of featurization model.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(
+            protein,
+            use_cache,
+            model_locations,
+            msa_file,
+            batch_size,
+            model_type,
+            **kwargs
+        )
+
+    def custom_featurizer(self, seqs, msa_file=None, **kwargs):
+        """
+        Featurizes sequences using MSA logits.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            msa_file (str): Path to MSA file.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: Array of MSA logits features.
+        """
+        X = self.featurize_msa(
+            seqs, msa_file, output_type="log_probabilities", **kwargs
+        )
+        return X

model/featurizers/prott5_featurizers.py

@@ -0,0 +1,65 @@
+import ankh
+import numpy as np
+import torch
+
+from model.featurizers.base_featurizers import BaseFeaturizer
+from transformers import T5Tokenizer, T5EncoderModel
+import re
+
+class ProtT5BaseFeaturizer(BaseFeaturizer):
+    def __init__(self, 
+                protein=None, 
+                use_cache=False,
+                model_version=None, 
+                batch_size=968,
+                model_type="ProtT5",
+                **kwargs):
+            
+        super().__init__(model_type,protein, use_cache, **kwargs)
+
+        self.batch_size = batch_size
+        self.model_version = model_version
+
+    def featurize_prott5(self, seqs):
+
+        if self.model_version == 'prot_t5_xl_u50':
+            self.tokenizer = T5Tokenizer.from_pretrained('Rostlab/prot_t5_xl_half_uniref50-enc', do_lower_case=False)
+            self.model = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_half_uniref50-enc").to(self.device)
+        else:
+            raise ValueError(f"Invalid model version: {self.model_version}")
+
+        input_seqs = [" ".join(list(re.sub(r"[UZOB]", "X", sequence))) for sequence in seqs]
+
+        seq_batch = []
+
+        for i in range(0, len(input_seqs), self.batch_size):
+            batch = input_seqs[i:i + self.batch_size]
+            # tokenize sequences and pad up to the longest sequence in the batch
+            ids = self.tokenizer(batch, add_special_tokens=True, padding="longest")
+
+            input_ids = torch.tensor(ids['input_ids']).to(self.device)
+            attention_mask = torch.tensor(ids['attention_mask']).to(self.device)
+
+            # generate embeddings
+            with torch.no_grad():
+                embeddings = self.model(input_ids=input_ids, attention_mask=attention_mask)
+
+            seq_batch.append(embeddings['last_hidden_state'].mean(axis=1).cpu().numpy())
+
+        return np.concatenate(seq_batch)
+
+
+class ProtT5_XL_U50_EmbedFeaturizer(ProtT5BaseFeaturizer):
+    def __init__(self, 
+                protein=None, 
+                use_cache=False,
+                model_version="prot_t5_xl_u50", 
+                batch_size=968,
+                model_type="ProtT5_XL_U50_Embed",
+                **kwargs):
+        super().__init__(protein, use_cache, model_version, batch_size, model_type, **kwargs)
+
+    def custom_featurizer(self, seqs):
+
+        X = self.featurize_prott5(seqs)
+        return X

model/featurizers/unirep_featurizers.py

@@ -0,0 +1,42 @@
+from jax_unirep.featurize import get_reps
+from jax_unirep.utils import load_params
+
+from model.featurizers.base_featurizers import BaseFeaturizer
+
+UNIREP_MODEL_SIZES = [1900, 256, 64]
+
+
+class UnirepBaseFeaturizer(BaseFeaturizer):
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        model_locations=None,
+        model_type="unirep",
+        model_size=1900,
+        **kwargs,
+    ):
+        super().__init__(model_type, protein, use_cache, **kwargs)
+        self.model_locations = model_locations
+        self.update_model_name(model_size)
+        self.load_params()
+
+    def update_model_name(self, model_size):
+        # Validate model size
+        assert model_size in UNIREP_MODEL_SIZES, "Model size must be 1900, 256, or 64."
+        self.model_size = model_size
+        self.model_type = self.model_type + str(model_size)
+
+    def load_params(self):
+        self.params = load_params(self.model_locations, self.model_size)[1]
+
+    def custom_featurizer(self, seqs, **kwargs):
+        h_avg, h_final, c_final = get_reps(
+            seqs=seqs, params=self.params, mlstm_size=self.model_size
+        )
+
+        return h_avg
+
+class EvotunedUnirepFeaturizer(UnirepBaseFeaturizer):
+    def __init__(self, model_type="evotuned_unirep", **kwargs):
+        super().__init__(model_type=model_type, **kwargs)

model/featurizers/zeroshot_featurizers.py

@@ -0,0 +1,632 @@
+from Bio import SeqIO
+
+from model.featurizers.base_featurizers import BaseFeaturizer
+from model.featurizers.model_choices import FEATURE_MODELS
+from model.utils.data_utils import find_mutations_multithreaded
+
+class ZeroshotBaseFeaturizer(BaseFeaturizer):
+    """Base class for zero-shot featurizers.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        model_locations (list): Paths to model files.
+        wt_file (str): Path to wild-type sequence file.
+        wt_seq (str): Wild-type protein sequence.
+
+    Example Usage:
+
+    featurizer = ZeroshotBaseFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        wt_file='wt.fasta',      # Path to wild-type sequence
+        model_locations=[]        # Paths to model files
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        model_locations=None,
+        wt_file=None,
+        model_type="zeroshot",
+        **kwargs
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            model_locations (list): Paths to model files.
+            wt_file (str): Path to wild-type sequence file.
+            model_type (str): Type of featurization model.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(model_type, protein, use_cache, **kwargs)
+        self.model_locations = model_locations
+        self.wt_file = wt_file
+        self.wt_seq = str(SeqIO.read(self.wt_file, "fasta").seq)
+
+    def featurize_zeroshot(
+        self, seqs, model_locations, wt_file, zeroshot_model, **kwargs
+    ):
+        """
+        Featurizes sequences using zero-shot prediction.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            model_locations (list): Paths to model files.
+            wt_file (str): Path to wild-type sequence file.
+            zeroshot_model (callable): Zero-shot prediction function.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: Zero-shot prediction scores.
+        """
+        assert (self.wt_file is not None) or (
+            wt_file is not None
+        ), "No wt sequence provided."
+        assert (self.model_locations is not None) or (
+            model_locations is not None
+        ), "No model locations provided."
+
+        wt_file = wt_file or self.wt_file
+        model_locations = model_locations or self.model_locations
+
+        wt_seq = str(SeqIO.read(self.wt_file, "fasta").seq)
+        model_locations = self.model_locations
+
+        mutations = find_mutations_multithreaded(wt_seq, seqs)
+
+        # make sure to remove model_locations and sequence from kwargs
+        kwargs.pop("model_locations", None)
+        kwargs.pop("sequence", None)
+        kwargs['device'] = self.device
+        X = zeroshot_model(
+            mutations, model_locations=model_locations, sequence=wt_seq, **kwargs
+        )
+
+        X = X.reshape(-1, 1)
+
+        return X
+
+
+class ZeroshotESMFeaturizer(ZeroshotBaseFeaturizer):
+    """Class for ESM zero-shot featurization.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        model_locations (list): Paths to model files.
+        wt_file (str): Path to wild-type sequence file.
+        wt_seq (str): Wild-type protein sequence.
+
+    Example Usage:
+
+    featurizer = ZeroshotESMFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        wt_file='wt.fasta'       # Path to wild-type sequence
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        model_locations=FEATURE_MODELS["zeroshot_esm"],
+        wt_file=None,
+        model_type="zeroshot_esm",
+        **kwargs
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            model_locations (list): Paths to model files.
+            wt_file (str): Path to wild-type sequence file.
+            model_type (str): Type of featurization model.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(
+            protein=protein,
+            use_cache=use_cache,
+            model_locations=model_locations,
+            wt_file=wt_file,
+            model_type=model_type,
+            **kwargs
+        )
+
+    def custom_featurizer(self, seqs, **kwargs):
+        """
+        Featurizes sequences using ESM zero-shot prediction.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: ESM zero-shot prediction scores.
+        """
+        from model.utils.zeroshot_utils import zero_shot_esm as zero_shot
+
+        X = self.featurize_zeroshot(
+            seqs,
+            model_locations=self.model_locations,
+            wt_file=self.wt_file,
+            zeroshot_model=zero_shot,
+            **kwargs
+        )
+        return X
+
+
+class ZeroshotMSAFeaturizer(ZeroshotBaseFeaturizer):
+    """Class for MSA zero-shot featurization.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        model_locations (list): Paths to model files.
+        wt_file (str): Path to wild-type sequence file.
+        wt_seq (str): Wild-type protein sequence.
+        msa_file (str): Path to MSA file.
+
+    Example Usage:
+
+    featurizer = ZeroshotMSAFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        wt_file='wt.fasta',      # Path to wild-type sequence
+        msa_file='msa.fasta'     # Path to MSA file
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        model_locations=FEATURE_MODELS["zeroshot_msa"],
+        wt_file=None,
+        msa_file=None,
+        model_type="zeroshot_msa",
+        **kwargs
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            model_locations (list): Paths to model files.
+            wt_file (str): Path to wild-type sequence file.
+            msa_file (str): Path to MSA file.
+            model_type (str): Type of featurization model.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(
+            protein=protein,
+            use_cache=use_cache,
+            model_locations=model_locations,
+            wt_file=wt_file,
+            model_type=model_type,
+            **kwargs
+        )
+        self.msa_file = msa_file
+
+    def custom_featurizer(self, seqs, **kwargs):
+        """
+        Featurizes sequences using MSA zero-shot prediction.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: MSA zero-shot prediction scores.
+        """
+        from model.utils.zeroshot_utils import zero_shot_msa as zero_shot
+
+        X = self.featurize_zeroshot(
+            seqs,
+            model_locations=self.model_locations,
+            wt_file=self.wt_file,
+            zeroshot_model=zero_shot,
+            msa_file=self.msa_file,
+            **kwargs
+        )
+        return X
+
+
+class ZeroshotCSCSFeaturizer(ZeroshotBaseFeaturizer):
+    """Class for CSCS zero-shot featurization.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        model_locations (list): Paths to model files.
+        wt_file (str): Path to wild-type sequence file.
+        wt_seq (str): Wild-type protein sequence.
+
+    Example Usage:
+
+    featurizer = ZeroshotCSCSFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        wt_file='wt.fasta'       # Path to wild-type sequence
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        model_locations=None,
+        wt_file=None,
+        model_type="zeroshot_cscs",
+        **kwargs
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            model_locations (list): Paths to model files.
+            wt_file (str): Path to wild-type sequence file.
+            model_type (str): Type of featurization model.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(
+            protein=protein,
+            use_cache=use_cache,
+            model_locations=model_locations,
+            wt_file=wt_file,
+            model_type=model_type,
+            **kwargs
+        )
+
+    def custom_featurizer(self, seqs, **kwargs):
+        """
+        Featurizes sequences using CSCS zero-shot prediction.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: CSCS zero-shot prediction scores.
+        """
+        from model.utils.zeroshot_utils import zero_shot_cscs as zero_shot
+
+        X = self.featurize_zeroshot(
+            seqs,
+            model_locations=self.model_locations,
+            wt_file=self.wt_file,
+            zeroshot_model=zero_shot,
+            **kwargs
+        )
+        return X
+
+
+class ZeroshotCSCSGramFeaturizer(ZeroshotBaseFeaturizer):
+    """Class for CSCS-Gram zero-shot featurization.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        model_locations (list): Paths to model files.
+        wt_file (str): Path to wild-type sequence file.
+        wt_seq (str): Wild-type protein sequence.
+
+    Example Usage:
+
+    featurizer = ZeroshotCSCSGramFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        wt_file='wt.fasta'       # Path to wild-type sequence
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        model_locations=None,
+        wt_file=None,
+        model_type="zeroshot_cscs_gram",
+        **kwargs
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            model_locations (list): Paths to model files.
+            wt_file (str): Path to wild-type sequence file.
+            model_type (str): Type of featurization model.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(
+            protein=protein,
+            use_cache=use_cache,
+            model_locations=model_locations,
+            wt_file=wt_file,
+            model_type=model_type,
+            **kwargs
+        )
+
+    def custom_featurizer(self, seqs, **kwargs):
+        """
+        Featurizes sequences using CSCS-Gram zero-shot prediction.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: CSCS-Gram zero-shot prediction scores.
+        """
+        from model.utils.zeroshot_utils import zero_shot_cscs_gram as zero_shot
+
+        X = self.featurize_zeroshot(
+            seqs,
+            model_locations=self.model_locations,
+            wt_file=self.wt_file,
+            zeroshot_model=zero_shot,
+            **kwargs
+        )
+        return X
+
+
+class ZeroshotCSCSSemFeaturizer(ZeroshotBaseFeaturizer):
+    """Class for CSCS-Sem zero-shot featurization.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        model_locations (list): Paths to model files.
+        wt_file (str): Path to wild-type sequence file.
+        wt_seq (str): Wild-type protein sequence.
+
+    Example Usage:
+
+    featurizer = ZeroshotCSCSSemFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        wt_file='wt.fasta'       # Path to wild-type sequence
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        model_locations=None,
+        wt_file=None,
+        model_type="zeroshot_cscs_sem",
+        **kwargs
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            model_locations (list): Paths to model files.
+            wt_file (str): Path to wild-type sequence file.
+            model_type (str): Type of featurization model.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(
+            protein=protein,
+            use_cache=use_cache,
+            model_locations=model_locations,
+            wt_file=wt_file,
+            model_type=model_type,
+            **kwargs
+        )
+
+    def custom_featurizer(self, seqs, **kwargs):
+        """
+        Featurizes sequences using CSCS-Sem zero-shot prediction.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: CSCS-Sem zero-shot prediction scores.
+        """
+        from model.utils.zeroshot_utils import zero_shot_cscs_sem as zero_shot
+
+        X = self.featurize_zeroshot(
+            seqs,
+            model_locations=self.model_locations,
+            wt_file=self.wt_file,
+            zeroshot_model=zero_shot,
+            **kwargs
+        )
+        return X
+
+
+class ZeroshotProseFeaturizer(ZeroshotBaseFeaturizer):
+    """Class for ProSE zero-shot featurization.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        model_locations (list): Paths to model files.
+        wt_file (str): Path to wild-type sequence file.
+        wt_seq (str): Wild-type protein sequence.
+
+    Example Usage:
+
+    featurizer = ZeroshotProseFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        wt_file='wt.fasta'       # Path to wild-type sequence
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        model_locations=FEATURE_MODELS["zeroshot_prose"],
+        wt_file=None,
+        model_type="zeroshot_prose",
+        **kwargs
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            model_locations (list): Paths to model files.
+            wt_file (str): Path to wild-type sequence file.
+            model_type (str): Type of featurization model.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(
+            protein=protein,
+            use_cache=use_cache,
+            model_locations=model_locations,
+            wt_file=wt_file,
+            model_type=model_type,
+            **kwargs
+        )
+
+    def custom_featurizer(self, seqs, **kwargs):
+        """
+        Featurizes sequences using ProSE zero-shot prediction.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: ProSE zero-shot prediction scores.
+        """
+        from model.utils.zeroshot_utils import zero_shot_prose as zero_shot
+
+        X = self.featurize_zeroshot(
+            seqs,
+            model_locations=self.model_locations,
+            wt_file=self.wt_file,
+            zeroshot_model=zero_shot,
+            **kwargs
+        )
+        return X
+
+
+class ZeroshotESMIFFeaturizer(ZeroshotBaseFeaturizer):
+    """Class for ESM-IF zero-shot featurization.
+
+    Attributes:
+        model_type (str): Type of featurization model to use.
+        name (str): Name of the featurizer.
+        protein (str): Name of protein being featurized.
+        use_cache (bool): Whether to cache featurization results.
+        flatten_features (bool): Whether to flatten output features.
+        device (torch.device): Device to use for computation.
+        model_locations (list): Paths to model files.
+        wt_file (str): Path to wild-type sequence file.
+        wt_seq (str): Wild-type protein sequence.
+        pdb_file (str): Path to PDB structure file.
+        chain_id (str): Chain identifier in PDB file.
+
+    Example Usage:
+
+    featurizer = ZeroshotESMIFFeaturizer(
+        protein='protein1',       # Name of protein for caching
+        use_cache=True,          # Whether to cache results
+        wt_file='wt.fasta',      # Path to wild-type sequence
+        pdb_file='struct.pdb',   # Path to structure file
+        chain_id='A'             # Chain identifier
+    )
+    features = featurizer.featurize(sequences)
+    """
+
+    def __init__(
+        self,
+        protein=None,
+        use_cache=False,
+        model_locations=FEATURE_MODELS["zeroshot_esmif"],
+        wt_file=None,
+        model_type="zeroshot_esmif",
+        pdb_file=None,
+        chain_id='A',
+        **kwargs
+    ):
+        """
+        Args:
+            protein (str): Name of protein being featurized.
+            use_cache (bool): Whether to cache results.
+            model_locations (list): Paths to model files.
+            wt_file (str): Path to wild-type sequence file.
+            model_type (str): Type of featurization model.
+            pdb_file (str): Path to PDB structure file.
+            chain_id (str): Chain identifier in PDB file.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(
+            protein=protein,
+            use_cache=use_cache,
+            model_locations=model_locations,
+            wt_file=wt_file,
+            model_type=model_type,
+            pdb_file=pdb_file,
+            chain_id=chain_id,
+            **kwargs
+        )
+
+    def custom_featurizer(self, seqs, **kwargs):
+        """
+        Featurizes sequences using ESM-IF zero-shot prediction.
+
+        Args:
+            seqs (list): List of sequences to featurize.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: ESM-IF zero-shot prediction scores.
+        """
+        from model.utils.zeroshot_utils import zero_shot_esm_if as zero_shot
+
+        X = self.featurize_zeroshot(
+            seqs,
+            model_locations=self.model_locations,
+            wt_file=self.wt_file,
+            zeroshot_model=zero_shot,
+            pdb_file=self.pdb_file,
+            chain_id=self.chain_id,
+            **kwargs
+        )
+        return X

model/predictors/__init__.py

@@ -0,0 +1,3 @@
+from model.predictors.base_regressors import *
+from model.predictors.gaussian_process_regressors import *
+from model.predictors.neural_net_regressors import *

model/predictors/base_regressors.py

@@ -0,0 +1,550 @@
+import os
+import pickle
+from abc import ABC, abstractmethod
+
+from matplotlib import pyplot as plt
+import numpy as np
+import pandas as pd
+import scipy.stats as ss
+from sklearn.metrics import make_scorer
+from sklearn.linear_model import Ridge, LinearRegression
+from sklearn.model_selection import cross_val_score
+from sklearn.preprocessing import StandardScaler
+from sklearn.ensemble import RandomForestRegressor as RFRegressor
+
+from model.utils.other_utils import performance_report
+
+def run_model_experiments(splits, 
+                          features, 
+                          models, 
+                          experiment_name,
+                          use_cache=False,
+                          show_plots=True):
+    """
+    Trains multiple models with various data splits and features, evaluates their performance,
+    and compiles the results into a CSV file.
+
+    Args:
+        splits (list): A list of data splits to use for training
+        features (list): A list of features to use for training
+        models (list): A list of model instances to train
+        experiment_name (str): Name of the experiment for saving results
+        use_cache (bool, optional): Whether to use cached models. Defaults to False.
+        show_plots (bool, optional): Whether to show matplotlib plots. Defaults to True.
+
+    Returns:
+        pd.DataFrame: DataFrame containing evaluation results for all models
+
+    Example Usage:
+
+    run_model_experiments(splits, 
+                          features, 
+                          models, 
+                          experiment_name,
+                          use_cache=False)
+    """
+
+    names = []
+    stats = []
+
+    for split in splits:
+        for feature in features:
+            for model in models:
+
+                instance = model(split, feature, use_cache=use_cache, show_plots=show_plots)
+                
+                # Train and evaluate model
+                stat = instance.run_model()
+                names.append(instance.file_attrs['model_name'].split('__'))
+                stats.append(list(stat.values()))
+
+    # Return results for all training permutations
+    stats_array = np.array(stats)
+    names_array = np.array(names)
+    combined_array = np.concatenate([names_array, stats_array], axis=1)
+    columns = ['Data Split', 'Feature', 'Model'] + list(stat.keys())
+    table = pd.DataFrame(combined_array, columns=columns)
+
+    # Check if the directory exists, create it if it doesn't
+    dir_path = f'{instance.file_attrs["model_dir"]}/' + "results"
+    if not os.path.exists(dir_path):
+        os.makedirs(dir_path)
+
+    # Save the results
+    if use_cache:
+        table.to_csv(f"{dir_path}/{experiment_name}.csv", index=False)
+
+    return table
+
+
+class BaseRegressor(ABC):
+    """
+    Abstract base class for regression models.
+
+    Args:
+        data_splitter: Object containing train/test splits
+        featurizer: Object that converts sequences to features
+        model (str, optional): Name of model. Defaults to 'Base'
+        use_cache (bool, optional): Whether to use cached models. Defaults to False
+        show_plots (bool, optional): Whether to show matplotlib plots. Defaults to True
+        **kwargs: Additional keyword arguments
+
+    Attributes:
+        model_name (str): Name of the model
+        featurizer: Featurizer object
+        use_cache (bool): Whether to use cached models
+        kwargs (dict): Additional keyword arguments
+        X_train: Training features
+        X_test: Test features  
+        y_train: Training labels
+        y_test: Test labels
+        split_method (str): Name of data split method
+        file_attrs (dict): Dictionary of file attributes and paths
+        show_plots (bool): Whether to show matplotlib plots. Defaults to True
+    
+    Example Usage:
+
+    regressor = BaseRegressor(data_splitter, featurizer, model='Linear', use_cache=False, show_plots=True)
+    regressor.run_model()
+    """
+
+    def __init__(self, 
+                 data_splitter, 
+                 featurizer, 
+                 model='Base', 
+                 use_cache=False, 
+                 show_plots=True,
+                 **kwargs):
+        
+        # Set variables
+        self.model_name = model
+        self.featurizer = featurizer
+        self.use_cache = use_cache
+        self.show_plots = show_plots
+        self.kwargs = kwargs
+
+        # Setup data
+        self.X_train = data_splitter.splits['X_train']
+        self.X_test = data_splitter.splits['X_test']
+        self.y_train = data_splitter.splits['y_train']
+        self.y_test = data_splitter.splits['y_test']
+        self.split_method = data_splitter.splits['split_name']
+
+        # Check if 'X_val' is not a key in data_splitter
+        if 'X_val' in data_splitter.splits:
+            print("Validation sets do not need to be present in data splits for non-neural network models.")
+
+        # Set model directory
+        self.file_attrs = data_splitter.file_attrs
+        self.file_attrs['model_name'] = self.split_method + ' __ ' + self.featurizer.name + ' __ ' + self.model_name
+        self.file_attrs['model_dir'] = os.path.join(data_splitter.file_attrs["dataset_dir"], 'model_cache', data_splitter.file_attrs["dataset_name"])
+        self.file_attrs['model_path'] = os.path.join(self.file_attrs['model_dir'], 'objects', f'{self.file_attrs["model_name"]}.pkl')
+        
+        
+        # Load model if available
+        if self.file_attrs['model_path'] is not None and os.path.exists(self.file_attrs['model_path']) and self.use_cache:
+            self.load_model(self.file_attrs['model_path'])
+
+    def run_model(self, eval=True):
+        """
+        Runs model training and evaluation.
+
+        Args:
+            eval (bool, optional): Whether to evaluate the model. Defaults to True.
+
+        Returns:
+            If eval=True:
+                dict: Dictionary of evaluation statistics
+            If eval=False:
+                None
+        """
+
+        if self.file_attrs['model_path'] is not None and os.path.exists(self.file_attrs['model_path']) and self.use_cache:
+            pass
+        else:
+            print(f'Training model for {self.file_attrs["model_name"]}')
+            X = self.preprocess_data(self.X_train)
+            self.train(X, self.y_train)
+            
+            if self.use_cache:
+                self.save_model()
+        
+        if eval:
+            return self.evaluate()
+        else:
+            return None
+
+            
+    def load_model(self, model_path=None):
+        """
+        Loads a pre-trained model from a pkl file.
+
+        Args:
+            model_path (str, optional): Path to model file. Defaults to None.
+        """
+
+        # set location to load model
+        model_path = self.file_attrs['model_path'] if model_path is None else model_path
+        print(f'Loading model from {model_path}')
+
+        try:
+            with open(model_path, 'rb') as file:
+                self.model = pickle.load(file)
+            print("Model loaded successfully.")
+        except FileNotFoundError:
+            print(f"Error: The file at {model_path} was not found.")
+        except PermissionError:
+            print(f"Error: Permission denied when trying to read {model_path}.")
+        except Exception as e:
+            print(f"An error occurred while loading the model: {str(e)}")
+
+    def save_model(self, model_path=None):
+        """
+        Saves the model to a pkl file.
+
+        Args:
+            model_path (str, optional): Path to save model to. Defaults to None.
+        """
+        
+        # set location to save model
+        model_path = self.file_attrs['model_path'] if model_path is None else model_path
+
+        dir_path = os.path.join(self.file_attrs['model_dir'], 'objects')
+        # Check if the directory exists, create it if it doesn't
+        if not os.path.exists(dir_path):
+            os.makedirs(dir_path)
+
+        # Save the model
+        print(f'Saving model to {model_path}')
+        pickle.dump(self.model, open(model_path, 'wb'))
+
+    def featurize(self, X):
+        """
+        Featurizes a list of sequences.
+
+        Args:
+            X (list): List of sequences to featurize
+
+        Returns:
+            array: Featurized sequences
+        """
+
+        X_featurized = self.featurizer.featurize(X)
+        return X_featurized
+
+    def preprocess_data(self, X):
+        """
+        Featurizes and scales input data.
+
+        Args:
+            X (list): List of sequences to preprocess
+
+        Returns:
+            array: Preprocessed data
+        """
+        X = self.featurizer.featurize(X)
+
+        X = X.reshape(X.shape[0],-1)
+
+        return X
+        
+    @abstractmethod
+    def train(self, X, y):
+        """
+        Trains the model.
+
+        Args:
+            X (array): Input features
+            y (array): Target values
+
+        Returns:
+            Trained model, also stored in self.model
+        """
+        pass
+
+    def evaluate(self):
+        """
+        Evaluates model on test set.
+
+        Returns:
+            dict: Dictionary of evaluation statistics
+        """
+        
+        # Evaluate model
+        y_pred = self.predict(self.X_test)
+
+        # Reshape data and get correlation stats
+        y, y_pred = np.array(self.y_test), np.array(y_pred)
+        y, y_pred = y.reshape(-1), y_pred.reshape(-1)
+
+        # Get stats
+        stats = performance_report(y, y_pred)
+
+        # Set the default parameters
+        plt.rcParams['font.size'] = 7
+        plt.rcParams['lines.linewidth'] = 0.5
+
+        # Plotting Results
+        fig, ax = plt.subplots(figsize=(4, 3))  # Adjust size as needed
+        
+        ## Mark data points that have activity less than 0 or greater than 1.2x the max experimental y value
+        y_max = max(y.max() * 1.2, y_pred.max() * 1.2)
+        y_min = min(y.min() * 0.8, y_pred.min() * 0.8)
+        # colors = np.where(y_pred > y_max, 'crimson', np.where(y_pred < 0, 'crimson', 'dodgerblue'))
+        # y_pred_adjusted = np.clip(y_pred, 0, y_max)
+
+        ## Scatter plot for main graph
+        ax.scatter(y_pred, y, c='dodgerblue', alpha=0.4, edgecolors='w', linewidth=0.5)
+
+        ## Draw x=y line
+        ax.plot([y_min, y_max], [y_min, y_max], 'k--', linewidth=0.5)
+
+        ## Set labels and title for main graph
+        ax.text(0.9, 0.1, f'Pearson r={stats["Pearson r"]:.2f}', fontsize=7, ha='right', va='bottom', transform=ax.transAxes)
+        ax.text(0.9, 0.2, f'Spearman r={stats["Spearman r"]:.2f}', fontsize=7, ha='right', va='bottom', transform=ax.transAxes)
+        ax.set_xlabel('Predicted Score', fontsize=7)
+        ax.set_ylabel('True Score', fontsize=7)
+        ax.set_title('Model Performance', fontsize=7)
+        ax.set_xlim(y_min, y_max)
+
+        ## Display model parameters using legend
+        model_params = self.file_attrs["model_name"].split('__')  # Assuming '|' separates different parameters
+        param_text = '\n'.join(model_params)
+        props = dict(boxstyle='square', facecolor='wheat', alpha=0.2)
+        ax.text(0.02, 0.98, param_text, transform=ax.transAxes, fontsize=7, verticalalignment='top', bbox=props)
+
+        # Adjust tick parameters
+        ax.tick_params(axis='both', which='major', labelsize=7)
+
+        # Show figure
+        if self.show_plots:
+            plt.show()
+        plt.close(fig)
+        
+        # Return the stats
+        return stats
+
+    @abstractmethod
+    def custom_predictor(self, X):
+        """
+        Custom prediction method to be implemented in subclasses.
+        Inputs have been filtered by self.predict()
+
+        Args:
+            X (array): Featurized sequences
+
+        Returns:
+            array: Model predictions
+        """
+        pass
+        
+    def predict(self, X):
+        """
+        Gets model predictions. Runs checks and calls custom_predictor.
+
+        Args:
+            X (list): List of sequences
+
+        Returns:
+            array: Model predictions
+        """
+
+        X_featurized = self.featurizer.featurize(X)
+
+        X_featurized = X_featurized.reshape(X_featurized.shape[0],-1)
+
+        predictions = self.custom_predictor(X_featurized)
+
+        return predictions
+
+class IdentityRegressor(BaseRegressor):
+    """
+    Identity regressor that returns all 1's.
+
+    Args:
+        data_splitter: Object containing train/test splits
+        featurizer: Object that converts sequences to features
+        model (str, optional): Name of model. Defaults to 'Linear'
+        use_cache (bool, optional): Whether to use cached models. Defaults to False
+        **kwargs: Additional keyword arguments
+    """
+ 
+    def train(self, X, y):
+        pass
+
+    def custom_predictor(self, X):
+        return [1 for _ in range(len(X))]
+
+class LinearRegressor(BaseRegressor):
+    """
+    Linear regression model.
+
+    Args:
+        data_splitter: Object containing train/test splits
+        featurizer: Object that converts sequences to features
+        model (str, optional): Name of model. Defaults to 'Linear'
+        use_cache (bool, optional): Whether to use cached models. Defaults to False
+        **kwargs: Additional keyword arguments
+    """
+    def __init__(self, data_splitter, featurizer, model='Linear', use_cache=False, **kwargs):
+        super().__init__(data_splitter, featurizer, model, use_cache, **kwargs)
+        
+    def train(self, X, y):
+        model = LinearRegression(
+            fit_intercept=True,
+            copy_X=True,
+            n_jobs=10,
+        )
+        model.fit(X, y)
+        self.model = model
+
+    def custom_predictor(self, X):
+        return self.model.predict(X)
+    
+class RandomForestRegressor(BaseRegressor):
+    """
+    Random Forest regression model.
+
+    Args:
+        data_splitter: Object containing train/test splits
+        featurizer: Object that converts sequences to features
+        model (str, optional): Name of model. Defaults to 'RandomForest'
+        use_cache (bool, optional): Whether to use cached models. Defaults to False
+        n_estimators (int, optional): Number of trees. Defaults to 100
+        criterion (str, optional): Split criterion. Defaults to 'friedman_mse'
+        max_depth (int, optional): Max tree depth. Defaults to None
+        min_samples_split (int, optional): Min samples for split. Defaults to 2
+        min_samples_leaf (int, optional): Min samples in leaf. Defaults to 1
+        min_weight_fraction_leaf (float, optional): Min weight fraction in leaf. Defaults to 0.0
+        max_features (float, optional): Max features to consider. Defaults to 1.0
+        max_leaf_nodes (int, optional): Max leaf nodes. Defaults to None
+        min_impurity_decrease (float, optional): Min impurity decrease. Defaults to 0.0
+        bootstrap (bool, optional): Whether to bootstrap. Defaults to True
+        oob_score (bool, optional): Whether to use out-of-bag score. Defaults to False
+        n_jobs (int, optional): Number of parallel jobs. Defaults to 6
+        random_state (int, optional): Random seed. Defaults to 1
+        verbose (int, optional): Verbosity level. Defaults to 0
+        warm_start (bool, optional): Whether to reuse solution. Defaults to False
+        ccp_alpha (float, optional): Complexity parameter. Defaults to 0.0
+        max_samples (int, optional): Max samples for bootstrap. Defaults to None
+        **kwargs: Additional keyword arguments
+    """
+    def __init__(self, data_splitter, featurizer, model='RandomForest', use_cache=False,
+                 n_estimators=100, 
+                 criterion='friedman_mse',
+                 max_depth=None,  
+                 min_samples_split=2, 
+                 min_samples_leaf=1,
+                 min_weight_fraction_leaf=0.0, 
+                 max_features=1.0,
+                 max_leaf_nodes=None, 
+                 min_impurity_decrease=0.0, 
+                 bootstrap=True, 
+                 oob_score=False,
+                 n_jobs=6, # change this based on number of cores 
+                 random_state=1, 
+                 verbose=0, 
+                 warm_start=False, 
+                 ccp_alpha=0.0,
+                 max_samples=None,
+                 **kwargs
+                 ):
+        
+        self.n_estimators = n_estimators
+        self.criterion = criterion
+        self.max_depth = max_depth
+        self.min_samples_split = min_samples_split
+        self.min_samples_leaf = min_samples_leaf
+        self.min_weight_fraction_leaf = min_weight_fraction_leaf
+        self.max_features = max_features
+        self.max_leaf_nodes = max_leaf_nodes
+        self.min_impurity_decrease = min_impurity_decrease
+        self.bootstrap = bootstrap
+        self.oob_score = oob_score
+        self.n_jobs = n_jobs
+        self.random_state = random_state
+        self.verbose = verbose
+        self.warm_start = warm_start
+        self.ccp_alpha = ccp_alpha
+        self.max_samples = max_samples
+
+        super().__init__(data_splitter, featurizer, model=model, use_cache=use_cache, **kwargs)
+    
+    def train(self, X, y):
+        self.model = RFRegressor(n_estimators=self.n_estimators,
+                                 criterion=self.criterion,
+                                 max_depth=self.max_depth,  
+                                 min_samples_split=self.min_samples_split,
+                                 min_samples_leaf=self.min_samples_leaf, 
+                                 min_weight_fraction_leaf=self.min_weight_fraction_leaf,
+                                 max_features=self.max_features, 
+                                 max_leaf_nodes=self.max_leaf_nodes,
+                                 min_impurity_decrease=self.min_impurity_decrease, 
+                                 bootstrap=self.bootstrap, 
+                                 oob_score=self.oob_score,
+                                 n_jobs=self.n_jobs, 
+                                 random_state=self.random_state, 
+                                 verbose=self.verbose,
+                                 warm_start=self.warm_start,
+                                 ccp_alpha=self.ccp_alpha,
+                                 max_samples=self.max_samples
+                                )
+        
+        self.model.fit(X, y)
+
+    def custom_predictor(self, X):
+        return self.model.predict(X)
+
+class RidgeRegressor(BaseRegressor):
+    """
+    Ridge regression model.
+
+    Args:
+        data_splitter: Object containing train/test splits
+        featurizer: Object that converts sequences to features
+        model (str, optional): Name of model. Defaults to 'Ridge'
+        use_cache (bool, optional): Whether to use cached models. Defaults to False
+        reg_coef (float, optional): Ridge regularization coefficient. If None, use CV. Defaults to None
+        linear_model_cls (class, optional): Sklearn linear model class. Defaults to Ridge
+        reg_coef_list (list, optional): List of regularization strengths for CV. Defaults to [0.1, 1.0, 2.0]
+        **kwargs: Additional keyword arguments
+    """
+    # [TODO] edit cv to include modifiable splits
+    def __init__(self, data_splitter, featurizer, model='Ridge', use_cache=False, reg_coef=None, linear_model_cls=Ridge, reg_coef_list=None, **kwargs):
+        """
+        Args:
+            - reg_coef: Ridge regression coefficient. If none, then train with CV
+            - linear_model_cls: sklearn linear model class
+            - reg_coef_list: list of ridge regression regularization strength (default: [0.1, 1.0])
+        """
+        self.reg_coef = reg_coef
+        self.linear_model_cls = linear_model_cls
+        self.reg_coef_list = reg_coef_list if reg_coef_list is not None else [0.1, 1.0, 2.0]
+        super().__init__(data_splitter, featurizer, model=model, use_cache=use_cache, **kwargs)
+    
+    def train(self, X, y):
+        def spearman(y_pred, y_true):
+            y_pred = np.array(y_pred)
+            y_true = np.array(y_true)
+
+            y_pred = y_pred.reshape(-1)
+            y_true = y_true.reshape(-1)
+            
+            if np.var(y_pred) < 1e-6 or np.var(y_true) < 1e-6:
+                return 0.0
+            return ss.spearmanr(y_pred, y_true).correlation
+        if self.reg_coef is None or self.reg_coef == 'CV':
+            best_reg_coef, best_score = None, -np.inf
+            for sample_reg_coef in self.reg_coef_list:
+                model = self.linear_model_cls(alpha=sample_reg_coef)
+                score = cross_val_score(model, X, y,
+                                        cv=5,  
+                                        scoring=make_scorer(spearman)).mean()
+                if score > best_score:
+                    best_reg_coef = sample_reg_coef
+                    best_score = score
+            self.model = self.linear_model_cls(alpha=best_reg_coef)
+            self.model.fit(X, y)
+
+    def custom_predictor(self, X):
+        return self.model.predict(X)

model/predictors/gaussian_process_regressors.py

@@ -0,0 +1,503 @@
+from joblib import Parallel, delayed
+from math import ceil
+
+from matplotlib import pyplot as plt
+import matplotlib.colors as mcolors
+import numpy as np
+from scipy.stats import iqr
+from sklearn.gaussian_process.kernels import (
+    ConstantKernel as C,
+    DotProduct as DP,
+    RBF
+)
+
+from model.predictors.base_regressors import BaseRegressor
+from model.utils.other_utils import performance_report
+
+def parallel_predict(model, X, batch_num, n_batches, verbose):
+    """
+    Makes predictions in parallel using batches.
+
+    Args:
+        model: The trained model to make predictions with
+        X (array-like): Input features to predict on
+        batch_num (int): Current batch number
+        n_batches (int): Total number of batches
+        verbose (bool): Whether to print progress messages
+
+    Returns:
+        tuple: (mean predictions, prediction variances)
+    """
+    mean, var = model.predict(X, return_std=True)
+    if verbose:
+        print('Finished predicting batch number {}/{}'
+               .format(batch_num + 1, n_batches))
+    return mean, var
+
+class GPRegressor(BaseRegressor):
+    """Base Gaussian Process regressor class.
+
+    Attributes:
+        n_restarts_ (int): Number of restarts for optimizer
+        kernel_ (sklearn.gaussian_process.kernels): Kernel function
+        normalize_y_ (bool): Whether to normalize target values
+        backend_ (str): Backend framework to use ('sklearn', 'gpy', or 'gpytorch')
+        batch_size_ (int): Batch size for predictions
+        n_jobs_ (int): Number of parallel jobs
+        verbose_ (bool): Whether to print progress messages
+        model: The trained GP model
+        uncertainties_ (array): Prediction uncertainties
+    """
+
+    def __init__(self, 
+                data_splitter, 
+                featurizer,
+                model='GPR',
+                n_restarts=0,
+                kernel=None,
+                normalize_y=True,
+                backend='sklearn',
+                batch_size=1000,
+                n_jobs=1,
+                verbose=False,
+                **kwargs
+                ):
+        """
+        Args:
+            data_splitter: Object to split data into train/test sets
+            featurizer: Object to convert sequences to numerical features
+            model (str): Model identifier string
+            n_restarts (int): Number of restarts for optimizer
+            kernel: Kernel function for GP
+            normalize_y (bool): Whether to normalize target values
+            backend (str): Framework to use ('sklearn', 'gpy', or 'gpytorch')
+            batch_size (int): Batch size for predictions
+            n_jobs (int): Number of parallel jobs
+            verbose (bool): Whether to print progress messages
+            **kwargs: Additional keyword arguments
+        """
+        self.n_restarts_ = n_restarts
+        self.kernel_ = kernel
+        self.normalize_y_ = normalize_y
+        self.backend_ = backend
+        self.batch_size_ = batch_size
+        self.n_jobs_ = n_jobs
+        self.verbose_ = verbose
+        super().__init__(data_splitter, featurizer, model, **kwargs)
+    
+    def train(self, X, y):
+        """
+        Train the GP model.
+
+        Args:
+            X (array-like): Training features
+            y (array-like): Training target values
+
+        Returns:
+            self: The trained model instance
+        """
+        n_samples, n_features = X.shape
+
+        if self.verbose_:
+            print('Fitting GP model on {} data points with dimension {}...'
+                   .format(*X.shape))
+
+        # scikit-learn backend.
+        if self.backend_ == 'sklearn':
+            from sklearn.gaussian_process import GaussianProcessRegressor
+            self.model = GaussianProcessRegressor(
+                kernel=self.kernel_,
+                normalize_y=self.normalize_y_,
+                alpha=1e-0,
+                n_restarts_optimizer=self.n_restarts_,
+                copy_X_train=False,
+            ).fit(X, y)
+
+        # GPy backend.
+        elif self.backend_ == 'gpy':
+            import GPy
+            if self.kernel_ == 'rbf':
+                kernel = GPy.kern.RBF(
+                    input_dim=n_features, variance=1., lengthscale=1.
+                )
+            else:
+                raise ValueError('Kernel value {} not supported'
+                                 .format(self.kernel_))
+
+            self.model = GPy.models.SparseGPRegression(
+                X, y.reshape(-1, 1), kernel=kernel,
+                num_inducing=min(self.n_inducing_, n_samples)
+            )
+            self.model.Z.unconstrain()
+            self.model.optimize(messages=self.verbose_)
+
+        # GPyTorch with CUDA backend.
+        elif self.backend_ == 'gpytorch':
+            import gpytorch
+            import torch
+
+            class GPyTorchRegressor(gpytorch.models.ExactGP):
+                def __init__(self, X, y, likelihood):
+                    super(GPyTorchRegressor, self).__init__(X, y, likelihood)
+                    self.mean_module = gpytorch.means.ConstantMean()
+                    self.covar_module = gpytorch.kernels.ScaleKernel(
+                        gpytorch.kernels.RBFKernel()
+                    )
+            
+                def forward(self, X):
+                    mean_X = self.mean_module(X)
+                    covar_X = self.covar_module(X)
+                    return gpytorch.distributions.MultivariateNormal(mean_X, covar_X)
+            
+            X = torch.Tensor(X).contiguous().cuda()
+            y = torch.Tensor(y).contiguous().cuda()
+
+            likelihood = gpytorch.likelihoods.GaussianLikelihood().cuda()
+            model = GPyTorchRegressor(X, y, likelihood).cuda()
+
+            model.train()
+            likelihood.train()
+
+            # Use the Adam optimizer.
+            #optimizer = torch.optim.LBFGS([ {'params': model.parameters()} ])
+            optimizer = torch.optim.Adam([
+                {'params': model.parameters()}, # Includes GaussianLikelihood parameters.
+            ], lr=1.)
+
+            # Loss for GPs is the marginal log likelihood.
+            mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)
+
+            training_iterations = 100
+            for i in range(training_iterations):
+                optimizer.zero_grad()
+                output = model(X)
+                loss = -mll(output, y)
+                loss.backward()
+                if self.verbose_:
+                    print('Iter {}/{} - Loss: {:.3f}'
+                           .format(i + 1, training_iterations, loss.item()))
+                optimizer.step()
+
+            self.model = model
+            self.likelihood_ = likelihood
+
+        if self.verbose_:
+            print('Done fitting GP model.')
+
+        return self
+
+    def evaluate(self):
+        """
+        Evaluates the model on a test set.
+
+        Returns:
+            tuple: (dict of evaluation metrics, matplotlib figure)
+        """
+        
+        # Reshape data and get correlation stats
+        y_pred = self.predict(self.X_test)
+        y, y_pred = np.array(self.y_test), np.array(y_pred)
+        y, y_pred = y.reshape(-1), y_pred.reshape(-1)
+        stats = performance_report(y, y_pred)
+
+        # Plotting
+        fig, ax = plt.subplots(figsize=(8, 5))  # Adjust size as needed
+
+        # Clip data points that have activity less than 0 or greater than 1.2x the max experimental y value
+        y_max = y.max()*1.1
+        y_pred_adjusted = np.clip(y_pred, 0, y_max)
+
+        # Scale uncertainties using IQR and color data points based on uncertainty
+        scaled_uncertainties = (self.uncertainties_ - np.percentile(self.uncertainties_, 25)) / iqr(self.uncertainties_)
+        cmap = plt.cm.viridis
+        colors = cmap(scaled_uncertainties)
+        colors[y_pred > y_max] = mcolors.to_rgba('crimson')
+        colors[y_pred < 0] = mcolors.to_rgba('crimson')
+
+        # Scatter plot for main graph
+        scatter = ax.scatter(y_pred_adjusted, y, c=colors, alpha = 0.4)
+
+        # Uncertainty colorbar
+        cbar = plt.colorbar(scatter, ax=ax)
+        cbar.set_label('Uncertainty')
+
+        # Draw x=y line
+        ax.plot([0, y_max], [0, y_max], 'k--', linewidth=2)
+
+        # Set labels and title for main graph
+        ax.text(0.9, 0.1, f'Pearson r={stats["Pearson r"]:.2f}', fontsize=12, ha='right', va='bottom', transform=ax.transAxes)
+        ax.set_xlabel('Predicted Score')
+        ax.set_ylabel('True Score')
+        ax.set_title(f'Model Performance')
+        ax.set_xlim(0, y_max)
+
+        # Display model parameters using legend
+        model_params = self.name.split('|')  # Assuming '|' separates different parameters
+        param_text = '\n'.join(model_params)
+        props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
+        ax.text(0.05, 0.95, param_text, transform=ax.transAxes, fontsize=9,
+                verticalalignment='top', bbox=props)
+
+        # Return the figure and axes object
+        return stats, fig
+
+    def custom_predictor(self, X):
+        """
+        Makes predictions using the trained GP model.
+
+        Args:
+            X (array-like): Features to predict on
+
+        Returns:
+            array: Mean predictions
+        """
+        if self.verbose_:
+            print('Finding GP model predictions on {} data points...'
+                   .format(X.shape[0]))
+
+        if self.backend_ == 'sklearn':
+            n_batches = int(ceil(float(X.shape[0]) / self.batch_size_))
+            results = Parallel(n_jobs=self.n_jobs_)(#, max_nbytes=None)(
+                delayed(parallel_predict)(
+                    self.model,
+                    X[batch_num*self.batch_size_:(batch_num+1)*self.batch_size_],
+                    batch_num, n_batches, self.verbose_
+                )
+                for batch_num in range(n_batches)
+            )
+            mean = np.concatenate([ result[0] for result in results ])
+            var = np.concatenate([ result[1] for result in results ])
+
+        elif self.backend_ == 'gpy':
+            mean, var = self.model.predict(X, full_cov=False)
+
+        elif self.backend_ == 'gpytorch':
+            import gpytorch
+            import torch
+
+            X = torch.Tensor(X).contiguous().cuda()
+
+            # Set into eval mode.
+            self.model.eval()
+            self.likelihood_.eval()
+
+            with torch.no_grad(), \
+                 gpytorch.settings.fast_pred_var(), \
+                 gpytorch.settings.max_root_decomposition_size(35):
+                preds = self.model(X)
+
+            mean = preds.mean.detach().cpu().numpy()
+            var = preds.variance.detach().cpu().numpy()
+
+        if self.verbose_:
+            print('Done predicting with GP model.')
+
+        self.uncertainties_ = var.flatten()
+        return mean.flatten()
+
+class SparseGPRegressor(BaseRegressor):
+    """Sparse Gaussian Process regressor using inducing points.
+
+    Attributes:
+        n_inducing_ (int): Number of inducing points
+        method_ (str): Method for selecting inducing points ('uniform' or 'geosketch')
+        n_restarts_ (int): Number of restarts for optimizer
+        kernel_ (sklearn.gaussian_process.kernels): Kernel function
+        backend_ (str): Backend framework to use
+        batch_size_ (int): Batch size for predictions
+        n_jobs_ (int): Number of parallel jobs
+        verbose_ (bool): Whether to print progress messages
+        gpr_: The trained GP model
+    """
+
+    def __init__(
+            self, 
+            data_splitter, 
+            featurizer,
+            model='SparseGPRegressor', 
+            n_inducing=1000,
+            method='geoskech',
+            n_restarts=0,
+            kernel=None,
+            backend='sklearn',
+            batch_size=1000,
+            n_jobs=1,
+            verbose=False,
+            **kwargs
+    ):
+        """
+        Args:
+            data_splitter: Object to split data into train/test sets
+            featurizer: Object to convert sequences to numerical features
+            model (str): Model identifier string
+            n_inducing (int): Number of inducing points
+            method (str): Method for selecting inducing points
+            n_restarts (int): Number of restarts for optimizer
+            kernel: Kernel function for GP
+            backend (str): Framework to use
+            batch_size (int): Batch size for predictions
+            n_jobs (int): Number of parallel jobs
+            verbose (bool): Whether to print progress messages
+            **kwargs: Additional keyword arguments
+        """
+        self.n_inducing_ = n_inducing
+        self.method_ = method
+        self.n_restarts_ = n_restarts
+        self.kernel_ = kernel
+        self.backend_ = backend
+        self.batch_size_ = batch_size
+        self.n_jobs_ = n_jobs
+        self.verbose_ = verbose
+        super().__init__(data_splitter, featurizer,model, **kwargs)
+
+    def train(self, X, y):
+        """
+        Train the sparse GP model.
+
+        Args:
+            X (array-like): Training features
+            y (array-like): Training target values
+        """
+        X, y = self.X, self.y
+        if X.shape[0] > self.n_inducing_:
+            if self.method_ == 'uniform':
+                uni_idx = np.random.choice(X.shape[0], self.n_inducing_,
+                                           replace=False)
+                X_sketch = X[uni_idx]
+                y_sketch = y[uni_idx]
+
+            elif self.method_ == 'geosketch':
+                from fbpca import pca
+                from geosketch import gs
+
+                U, s, _ = pca(X, k=100)
+                X_dimred = U[:, :100] * s[:100]
+                gs_idx = gs(X_dimred, self.n_inducing_, replace=False)
+                X_sketch = X[gs_idx]
+                y_sketch = y[gs_idx]
+
+        else:
+            X_sketch, y_sketch = X, y
+
+        self.gpr_ = GPRegressor(
+            n_restarts=self.n_restarts_,
+            kernel=self.kernel_,
+            backend=self.backend_,
+            batch_size=self.batch_size_,
+            n_jobs=self.n_jobs_,
+            verbose=self.verbose_,
+        ).fit(X_sketch, y_sketch)
+
+
+    def custom_predictor(self, X):
+        """
+        Makes predictions using the trained sparse GP model.
+
+        Args:
+            X (array-like): Features to predict on
+
+        Returns:
+            array: Mean predictions
+        """
+        y_pred = self.gpr_.predict(X)
+        self.uncertainties_ = self.gpr_.uncertainties_
+        return y_pred
+
+class GPLinearRegressor(GPRegressor):
+    """Gaussian Process regressor with linear kernel."""
+
+    def __init__(self, 
+                 data_splitter, 
+                 featurizer,
+                 model='GPLinearRegressor',
+                 n_restarts=0,
+                 kernel = C(1., 'fixed') * DP(1., 'fixed'),
+                 normalize_y=True,
+                 backend='sklearn',
+                 batch_size=1000,
+                 n_jobs=1,
+                 verbose=False,
+                 **kwargs
+                 ):
+        """
+        Args:
+            data_splitter: Object to split data into train/test sets
+            featurizer: Object to convert sequences to numerical features
+            model (str): Model identifier string
+            n_restarts (int): Number of restarts for optimizer
+            kernel: Linear kernel function
+            normalize_y (bool): Whether to normalize target values
+            backend (str): Framework to use
+            batch_size (int): Batch size for predictions
+            n_jobs (int): Number of parallel jobs
+            verbose (bool): Whether to print progress messages
+            **kwargs: Additional keyword arguments
+        """
+        
+        super().__init__(data_splitter, featurizer, model, n_restarts, kernel, normalize_y, backend, batch_size, n_jobs, verbose, **kwargs)
+
+class GPQuadRegressor(GPRegressor):
+    """Gaussian Process regressor with quadratic kernel."""
+
+    def __init__(self,
+                data_splitter, 
+                featurizer,
+                model='GPQuadRegressor',
+                n_restarts=0,
+                kernel = C(1., 'fixed') * (DP(1, 'fixed') ** 2),
+                normalize_y=True,
+                backend='sklearn',
+                batch_size=1000,
+                n_jobs=1,
+                verbose=False,
+                **kwargs
+                ):
+        """
+        Args:
+            data_splitter: Object to split data into train/test sets
+            featurizer: Object to convert sequences to numerical features
+            model (str): Model identifier string
+            n_restarts (int): Number of restarts for optimizer
+            kernel: Quadratic kernel function
+            normalize_y (bool): Whether to normalize target values
+            backend (str): Framework to use
+            batch_size (int): Batch size for predictions
+            n_jobs (int): Number of parallel jobs
+            verbose (bool): Whether to print progress messages
+            **kwargs: Additional keyword arguments
+        """
+    
+        super().__init__(data_splitter, featurizer, model, n_restarts, kernel, normalize_y, backend, batch_size, n_jobs, verbose, **kwargs)
+
+class GPRBFRegressor(GPRegressor):
+    """Gaussian Process regressor with RBF kernel."""
+
+    def __init__(self, 
+                 data_splitter, 
+                 featurizer,
+                model='GPRBFRegressor',
+                n_restarts=0,
+                kernel = C(1., 'fixed') * RBF(1., 'fixed'),
+                normalize_y=True,
+                backend='sklearn',
+                batch_size=1000,
+                n_jobs=1,
+                verbose=False,
+                **kwargs
+                ):
+        """
+        Args:
+            data_splitter: Object to split data into train/test sets
+            featurizer: Object to convert sequences to numerical features
+            model (str): Model identifier string
+            n_restarts (int): Number of restarts for optimizer
+            kernel: RBF kernel function
+            normalize_y (bool): Whether to normalize target values
+            backend (str): Framework to use
+            batch_size (int): Batch size for predictions
+            n_jobs (int): Number of parallel jobs
+            verbose (bool): Whether to print progress messages
+            **kwargs: Additional keyword arguments
+        """
+    
+        super().__init__(data_splitter, featurizer, model, n_restarts, kernel, normalize_y, backend, batch_size, n_jobs, verbose, **kwargs)

model/predictors/neural_net_regressors.py

@@ -0,0 +1,746 @@
+import torch, wandb
+from torch import nn, optim
+import numpy as np
+import matplotlib.pyplot as plt
+import yaml
+import os
+import pandas as pd
+
+from model.utils.other_utils import performance_report, log_results
+from model.utils.data_utils import TorchDataProcessor
+
+# Get the directory where the script is located
+script_dir = os.path.dirname(__file__)
+
+# Master Functions to Train and Evaluate Models
+def run_nn_model_experiments(splits,
+                             features,
+                             models,  # Fcn, Cnn
+                             experiment_name,
+                             use_cache=False,
+                             sweep_depth="standard",  # standard, custom, test
+                             search_method="grid",  # grid, bayes, test
+                             count=10,
+                             show_plots=True
+                             ):
+    """Run neural network model experiments with hyperparameter sweeps.
+
+    Args:
+        splits (list): List of DataSplitter objects containing train/val/test splits
+        features (list): List of feature types to use (e.g. ['onehot', 'esm'])
+        models (list): List of model classes to run (e.g. [Fcn, Cnn]) 
+        experiment_name (str): Name for the W&B experiment
+        use_cache (bool, optional): Whether to cache results. Defaults to False.
+        sweep_depth (str, optional): Sweep type - 'standard', 'custom', 'test'. Defaults to 'standard'.
+        search_method (str, optional): Search method - 'grid', 'bayes', 'test'. Defaults to 'grid'.
+        count (int, optional): Number of runs per sweep. Defaults to 10.
+        show_plots (bool, optional): Whether to show matplotlib plots. Defaults to True.
+
+    Returns:
+        None: Results are logged to W&B
+
+    Example:
+        >>> splits = [DataSplitter(data, 'random')]
+        >>> features = ['onehot']
+        >>> models = [Fcn, Cnn]
+        >>> run_nn_model_experiments(splits, 
+        ...                         features, 
+        ...                         models, 
+        ...                         experiment_name='my_experiment',
+        ...                         sweep_depth='selective',
+        ...                         search_method='bayes',
+        ...                         show_plots=True)
+    """
+        
+    for split in splits:
+        for feature in features:
+            for model in models:
+        
+                """Define sweep configuration."""
+                config_map = {
+                    ("Fcn", "standard", "grid"): "fcn_standard_grid_sweep.yaml",
+                    ("Fcn", "standard", "bayes"): "fcn_standard_bayes_sweep.yaml",
+                    ("Fcn", "custom", "grid"): "fcn_custom_grid_sweep.yaml",
+                    ("Fcn", "test", "test"): "fcn_test_sweep.yaml",
+                    ("Cnn", "standard", "grid"): "cnn_standard_grid_sweep.yaml",
+                    ("Cnn", "standard", "bayes"): "cnn_standard_bayes_sweep.yaml",
+                    ("Cnn", "custom", "grid"): "cnn_custom_grid_sweep.yaml",
+                    ("Cnn", "test", "test"): "cnn_test_sweep.yaml",
+                }
+
+                yaml_file = config_map.get((model.__name__, sweep_depth, search_method))
+                if yaml_file is None:
+                    print(
+                        f"Invalid sweep configuration: model={model}, sweep_depth={sweep_depth}, search_method={search_method}."
+                    )
+                    return
+
+                working_script_dir = script_dir
+
+                # Assuming 'script_dir' is defined earlier in your code
+                yaml_file_path = os.path.join(working_script_dir, "sweep_configs", yaml_file)
+
+                with open(yaml_file_path, "r") as file:
+                    sweep_config = yaml.safe_load(file)
+
+
+                """initialize the sweep."""
+                #sweep_id = wandb.sweep(sweep=sweep_config, project=experiment_name)
+
+                # Define a train function for hyperparameter sweeps with WANDB
+
+                #def train_function():
+                #    with wandb.init() as run:
+
+                        # Grab config
+                #        config = run.config
+
+                        # Specify model
+                #        instance = model(split, feature, use_cache=use_cache, config=config, show_plots=show_plots)
+
+                        # Train and evaluate model
+                #        stat = instance.run_model()
+
+                #if search_method == "grid" or search_method == "test":
+                #    wandb.agent(sweep_id, train_function)
+                #elif search_method == "bayes":
+                #    wandb.agent(sweep_id, train_function, count=count)
+                          #LL 2026      # 手动从 sweep_config 中提取一组默认参数（取每个参数的第一个值）
+                default_config = {}
+                if 'parameters' in sweep_config:
+                    for param_name, param_values in sweep_config['parameters'].items():
+                        if 'value' in param_values:
+                            default_config[param_name] = param_values['value']
+                        elif 'values' in param_values:
+                            default_config[param_name] = param_values['values'][0]
+                        else:
+                            default_config[param_name] = None
+
+                # 直接训练模型，不通过 wandb
+                instance = model(split, feature, use_cache=use_cache, config=default_config, show_plots=show_plots)
+                stat = instance.run_model()
+
+
+# Neural network classes
+class BaseNN(nn.Module):
+    """Base neural network class implementing common functionality.
+    
+    This class provides the base implementation for neural network models including
+    data loading, training loops, evaluation, and model saving/loading.
+
+    Args:
+        data_splitter: DataSplitter object containing train/val/test splits
+        featurizer: Featurizer object for processing sequences
+        nn_arch: Neural network architecture specification
+        model (str): Model name identifier. Defaults to "Base"
+        use_cache (bool): Whether to use model caching. Defaults to False
+        show_plots (bool): Whether to show matplotlib plots. Defaults to True
+        **kwargs: Additional keyword arguments
+
+    Attributes:
+        model_name (str): Name of the model
+        featurizer: Featurizer object
+        use_cache (bool): Whether caching is enabled
+        kwargs (dict): Additional arguments
+        nn_arch (str): Architecture specification string
+        device (torch.device): Device to run model on (CPU/GPU)
+        show_plots (bool): Whether to show matplotlib plots. Defaults to True
+    Example:
+        >>> splitter = DataSplitter(data, 'random')
+        >>> featurizer = OneHotFeaturizer()
+        >>> model = BaseNN(splitter, featurizer, [64,32], model='test', show_plots=True)
+    """
+
+    def __init__(self, data_splitter, featurizer, nn_arch, model="Base", use_cache=False, show_plots=True, **kwargs):
+        super(BaseNN, self).__init__()
+
+        # Set variables
+        self.model_name = model
+        self.featurizer = featurizer
+        self.use_cache = use_cache
+        self.kwargs = kwargs
+        self.nn_arch = "-".join([str(x) for x in nn_arch])
+        self.show_plots = show_plots
+
+        # Setup data
+        self.nn_data_processor = TorchDataProcessor(data_splitter, self.featurizer, self.kwargs["config"]["batch_size"])
+        #[TODO] remove this and only process data once required
+
+        self.split_method = self.nn_data_processor.split_name 
+
+        # set model directory
+        self.file_attrs = data_splitter.file_attrs
+        self.file_attrs['model_dir'] = os.path.join(data_splitter.file_attrs["dataset_dir"], 'model_cache', data_splitter.file_attrs["dataset_name"])
+
+        """Set variables."""
+        if torch.backends.mps.is_available():
+            self.device = torch.device("mps")
+            print("MPS available. Using Apple Silicon GPU for Neural Network.")
+        elif torch.cuda.is_available():
+            self.device = torch.device("cuda:0")
+            print("CUDA available. Using Nvidia GPU for Neural Network.")
+        else:
+            self.device = torch.device("cpu")
+            print("Neither MPS nor CUDA is available. Using CPU for Neural Network.")
+
+    def setup_model(self):
+        """Set up the model by initializing hyperparameters and loading cached model if available."""
+
+        # Retrieve hyperparameters from current run config
+        self.set_hyperparams()
+
+        # Define model
+        self.file_attrs['model_name'] = (
+                     self.split_method + " __ " + 
+                     self.featurizer.name + " __ " + 
+                     self.model_name + " __ " +
+                     self.nn_arch + " __ " +
+                     str(self.kwargs["config"]["learning_rate"]) + " __ " +
+                     str(self.kwargs["config"]["batch_size"]) + " __ " +
+                     self.kwargs["config"]["optimizer"]
+                     )
+        
+        self.model_path = os.path.join(self.file_attrs['model_dir'], 'objects', f'{self.file_attrs["model_name"]}.pth')
+
+        # Load model if available 
+        if self.model_path is not None and os.path.exists(self.model_path) and self.use_cache:
+            self.load_model(model_path=None)
+            self.to(self.device)
+        else: 
+            self.to(self.device)
+
+    def run_model(self, eval=True):
+        """Run the full model training and evaluation pipeline.
+        
+        This method handles:
+        1. Loading cached model if available
+        2. Training the model if needed
+        3. Evaluating on test set
+        4. Saving model if caching enabled
+        
+        Returns:
+            dict: Dictionary of model performance statistics
+        """
+
+        if self.model_path is not None and os.path.exists(self.model_path):
+            model = self
+            train_loss = self.train_loop_eval_mode(model)
+            val_loss = self.val_loop(model)
+
+        else:
+            model = self
+            
+            # Train model
+
+            for epoch in range(self.epochs):
+                train_loss = self.train_loop(model)
+                val_loss = self.val_loop(model)
+
+                # Log data
+                if wandb.run is not None:
+                    wandb.log({"Train Loss": train_loss, "Val Loss": val_loss})
+
+                # Check for early stopping
+                if self.early_stopping_check(val_loss, epoch) == True:
+                    break
+                else:
+                    continue
+
+            # Save model
+            if self.use_cache:
+                self.save_model(model, model_path=None)
+
+        # Test model
+        if eval == True:
+            return self.evaluate(model)
+        else:
+            return None
+    
+    def load_model(self, model_path=None):
+        """Load a pre-trained model from disk.
+
+        Args:
+            model_path (str, optional): Path to model file. If None, uses default path.
+        """
+
+        # set location to load model
+        model_path = self.model_path if model_path is None else model_path
+        print(f"Loading model from {model_path}")
+        # Load the trained model parameters
+        self.load_state_dict(torch.load(model_path, map_location=self.device, weights_only=True))
+
+    def save_model(self, model, model_path=None):
+        """Save model to disk.
+
+        Args:
+            model: Model to save
+            model_path (str, optional): Path to save model to. If None, uses default path.
+        """
+        
+        # set location to save model
+        model_path = self.model_path if model_path is None else model_path
+
+        dir_path = os.path.join(self.file_attrs['model_dir'], 'objects')
+        # Check if the directory exists, create it if it doesn't
+        if not os.path.exists(dir_path):
+            os.makedirs(dir_path)
+
+        # Save the model
+        print(f"Saving model to {self.model_path}")
+        torch.save(model.state_dict(), self.model_path)
+
+    def forward(self, x):
+        """Forward pass through the network.
+        
+        Args:
+            x: Input tensor
+            
+        Returns:
+            Output tensor
+        """
+        return x
+
+    def train_loop(self, model):
+        """Training loop for one epoch.
+        
+        Args:
+            model: Model to train
+            
+        Returns:
+            float: Average training loss for the epoch
+        """
+        model.train()
+        total_train_loss = 0
+        total_samples = 0
+
+        # [TODO] new function to set up train loader if not already done
+        if not hasattr(self, 'train_loader'):
+            self.train_loader = self.nn_data_processor.setup_train_loader()
+
+        for batch in self.train_loader:
+            inputs, targets, __ = batch
+            inputs, targets = inputs.to(self.device), targets.to(self.device)
+            targets = targets.unsqueeze(1)
+            self.optimizer.zero_grad()
+            outputs = model(inputs)
+            loss = self.criterion(outputs, targets)
+            total_train_loss += loss.item()
+            loss.backward()
+            self.optimizer.step()
+            total_samples += inputs.size(0)
+
+        train_loss = total_train_loss / total_samples
+
+        return train_loss
+    
+    def train_loop_eval_mode(self, model):
+        """Training loop in evaluation mode (no gradients).
+        
+        Args:
+            model: Model to evaluate
+            
+        Returns:
+            float: Average training loss
+        """
+        model.eval()
+        with torch.no_grad():
+            total_train_loss = 0
+            total_samples = 0
+
+            # [TODO] new function to set up train loader if not already done
+            if not hasattr(self, 'train_loader'):
+                self.train_loader = self.nn_data_processor.setup_train_loader()
+
+            for batch in self.train_loader:
+                inputs, targets, __ = batch
+                inputs, targets = inputs.to(self.device), targets.to(self.device)
+                targets = targets.unsqueeze(1)
+                outputs = model(inputs)
+                total_train_loss += self.criterion(outputs, targets).item()
+                total_samples += inputs.size(0)
+
+            train_loss = total_train_loss / total_samples
+
+        return train_loss
+
+    def val_loop(self, model):
+        """Validation loop.
+        
+        Args:
+            model: Model to evaluate
+            
+        Returns:
+            float: Average validation loss
+        """
+        model.eval()
+        with torch.no_grad():
+            total_val_loss = 0
+            total_samples = 0
+            # [TODO] new function to set up val loader if not already done
+            if not hasattr(self, 'val_loader'):
+                self.val_loader = self.nn_data_processor.setup_val_loader()
+                
+            for batch in self.val_loader:
+                inputs, targets, __ = batch
+                inputs, targets = inputs.to(self.device), targets.to(self.device)
+                targets = targets.unsqueeze(1)
+                outputs = model(inputs)
+                total_val_loss += self.criterion(outputs, targets).item()
+                total_samples += inputs.size(0)
+
+            val_loss = total_val_loss / total_samples
+
+        return val_loss
+
+    def evaluate(self, model):
+        """Evaluate model on test set.
+        
+        Args:
+            model: Model to evaluate
+            
+        Returns:
+            dict: Dictionary of performance statistics
+        """
+        
+        # Evaluate model, get metrics for validation and test set
+        model.eval()
+
+        stats_dict = {
+            "val": {},
+            "test": {}
+        }
+        
+        loader_names = ["val", "test"]
+
+        with torch.no_grad():
+
+            # [TODO] new function to set up val and test loaders if not already done
+            if not hasattr(self, 'val_loader'):
+                self.val_loader = self.nn_data_processor.setup_val_loader()
+            if not hasattr(self, 'test_loader'):
+                self.test_loader = self.nn_data_processor.setup_test_loader()
+                
+            for index, loader in enumerate([self.val_loader, self.test_loader]):
+                loader_name = loader_names[index]
+                total_loss = 0
+                total_samples = 0
+                y = []
+                y_pred = []
+                original_sequences_list = []
+
+                for batch in loader:
+                    inputs, targets, original_sequences = batch
+                    inputs, targets = inputs.to(self.device), targets.to(self.device)
+                    targets = targets.unsqueeze(1)
+                    outputs = model(inputs)
+                    total_loss += self.criterion(outputs, targets).item()
+
+                    # Move to CPU and convert to numpy
+                    y.extend(targets.cpu().detach().numpy())
+                    y_pred.extend(outputs.cpu().detach().numpy())
+                    original_sequences_list.extend(original_sequences)
+
+                    total_samples += inputs.size(0)
+
+
+                # Reshape data and get correlation stats
+                y = np.concatenate(y).ravel()
+                y_pred = np.concatenate(y_pred).ravel()
+
+                # Get stats
+                stats_dict[loader_name] = performance_report(y, y_pred)
+
+        # graph results for test set
+        # Set the default parameters
+        plt.rcParams['font.size'] = 7
+        plt.rcParams['lines.linewidth'] = 0.5
+
+        fig, ax = plt.subplots(figsize=(4, 3))
+        
+        # Mark data points that have activity less than 0 or greater than 1.2x the max experimental y value
+        y_max = max(y.max(), y_pred.max()) * 1.2
+        colors = np.where(y_pred > y_max, 'crimson', np.where(y_pred < 0, 'crimson', 'dodgerblue'))
+        y_pred_adjusted = np.clip(y_pred, 0, y_max)
+
+        # Scatter plot for main graph
+        ax.scatter(y_pred_adjusted, y, c=colors, alpha=0.4, edgecolors='w', linewidth=0.5)
+
+        # Draw x=y line
+        ax.plot([0, y_max], [0, y_max], 'k--', linewidth=0.5)
+
+        # Set labels and title for main graph
+        ax.text(0.9, 0.1, f'Pearson r={stats_dict["test"]["Pearson r"]:.2f}', fontsize=7, ha='right', va='bottom', transform=ax.transAxes)
+        ax.text(0.9, 0.2, f'Spearman r={stats_dict["test"]["Spearman r"]:.2f}', fontsize=7, ha='right', va='bottom', transform=ax.transAxes)
+        ax.set_xlabel('Predicted Score', fontsize=7)
+        ax.set_ylabel('True Score', fontsize=7)
+        ax.set_title('Model Performance', fontsize=7)
+        ax.set_xlim(0, y_max)
+
+        # Display model parameters using legend
+        model_params = self.file_attrs['model_name'].split('__')  # Assuming '|' separates different parameters
+        param_text = '\n'.join(model_params)
+        props = dict(boxstyle='square', facecolor='wheat', alpha=0.2)
+        ax.text(0.02, 0.98, param_text, transform=ax.transAxes, fontsize=7, verticalalignment='top', bbox=props)
+
+        # Adjust tick parameters
+        ax.tick_params(axis='both', which='major', labelsize=7)
+
+        self.fig = fig
+
+        if self.show_plots:
+            plt.show()
+        plt.close(fig)
+            
+        # Log data
+        log_results(stats_dict, self)
+
+        # Save predictions for test set as a table
+        if self.use_cache:
+            dir_path = os.path.join(self.file_attrs['model_dir'], 'results')
+            # Check if the directory exists, create it if it doesn't
+            if not os.path.exists(dir_path):
+                os.makedirs(dir_path)
+            pred_results = pd.DataFrame({"original_sequences": original_sequences_list, "y": list(y), "y_pred": list(y_pred)})
+            pred_results.to_csv(f"{dir_path}/{self.file_attrs['model_name']}.csv", index=False)
+
+        return stats_dict['test']
+
+    def early_stopping_check(self, val_loss, epoch):
+        """Check if early stopping criteria are met.
+        
+        Args:
+            val_loss (float): Current validation loss
+            epoch (int): Current epoch number
+            
+        Returns:
+            bool: True if training should stop, False otherwise
+        """
+        # modify epoch count
+        val_loss_delta = self.val_loss_min - val_loss
+        if val_loss_delta > self.val_loss_delta_min:
+            self.val_loss_min = val_loss
+            self.epochs_no_improve = 0
+        else:
+            self.epochs_no_improve += 1
+
+        # check epoch count
+        if self.epochs_no_improve == self.patience:
+            print(f"Early stopping after {epoch} epochs with {self.val_loss_min}.")
+            return True
+        else:
+            return False
+
+    def set_hyperparams(self):
+        """Set model hyperparameters from config."""
+
+        self.criterion = nn.MSELoss()
+        self.lr = self.kwargs["config"]["learning_rate"]
+        if self.kwargs["config"]["optimizer"] == "adam":
+            self.optimizer = optim.Adam(self.parameters(), lr=self.lr)
+        elif self.kwargs["config"]["optimizer"] == "sgd":
+            self.optimizer = optim.SGD(self.parameters(), lr=self.lr)
+        self.epochs = self.kwargs["config"]["epochs"]
+
+        # early stopping
+        self.patience = 15
+        self.val_loss_min = float("inf")
+        self.val_loss_delta_min = 0.00001
+        self.epochs_no_improve = 0  # initialize epochs_no_improve for early stopping
+
+    def custom_predictor(self, X):
+        """Make predictions on input data.
+        
+        Args:
+            X: Input features
+            
+        Returns:
+            numpy.ndarray: Model predictions
+        """
+        
+        model = self
+        inputs = torch.from_numpy(X.astype(np.float32)).to(self.device)
+        
+        model.eval()
+        with torch.no_grad():
+            outputs = model(inputs)
+            
+        outputs_np = outputs.cpu().numpy()
+        return outputs_np
+
+    def predict(self, X, batch_size=10000):
+        """Make predictions on sequences in batches.
+        
+        Args:
+            X (list): List of sequences to predict
+            
+        Returns:
+            numpy.ndarray: Array of predictions
+        """
+        batch_size = batch_size
+        predictions = []
+        
+        # Process in batches
+        for i in range(0, len(X), batch_size):
+            batch = X[i:i + batch_size]
+            X_featurized = self.featurizer.featurize(batch)
+            X_featurized = X_featurized.reshape(X_featurized.shape[0], -1)
+            batch_predictions = self.custom_predictor(X_featurized)
+            predictions.append(batch_predictions)
+
+        return np.concatenate(predictions).ravel()
+
+class Fcn(BaseNN):
+    """Fully connected neural network model.
+    
+    Args:
+        data_splitter: DataSplitter object containing train/val/test splits
+        feature: Featurizer object for processing sequences
+        model (str): Model name identifier. Defaults to "fcn"
+        use_cache (bool): Whether to use model caching. Defaults to False
+        show_plots (bool): Whether to show matplotlib plots. Defaults to True
+        **kwargs: Additional keyword arguments including network architecture
+        
+    Example:
+        >>> splitter = DataSplitter(data, 'random')
+        >>> featurizer = OneHotFeaturizer()
+        >>> model = Fcn(splitter, featurizer, config=config, use_cache=True, show_plots=True)
+    """
+
+    def __init__(self, data_splitter, feature, model="fcn", use_cache=False, show_plots=True, **kwargs):
+        
+        # Specify network architecture
+        nn_arch = [kwargs["config"]["layer_size"]] * kwargs["config"]["num_layers"]
+
+        super().__init__(data_splitter, feature, nn_arch, model, use_cache=use_cache, show_plots=show_plots, **kwargs)
+
+        # [TODO] new function to extract input features if not already done
+
+        X_train_feat_example = self.nn_data_processor.featurize([self.nn_data_processor.X_train[0]])[0]
+        input_features = X_train_feat_example.flatten().shape[0]
+        self.flatten = nn.Flatten()
+        self.layers = nn.ModuleList()
+
+        # First layer
+        self.layers.append(nn.Linear(input_features, nn_arch[0]))
+        self.layers.append(nn.LeakyReLU(negative_slope=0.2))
+        self.layers.append(nn.Dropout(p=0.2))
+
+        for i in range(0, len(nn_arch)):
+            if i < len(nn_arch) - 1:
+                self.layers.append(nn.Linear(nn_arch[i], nn_arch[i + 1]))
+                self.layers.append(nn.LeakyReLU(negative_slope=0.2))
+                self.layers.append(nn.Dropout(p=0.2))
+            if i == len(nn_arch) - 1:
+                self.layers.append(nn.Linear(nn_arch[i], 1))
+
+        # set model hyperparameters
+        self.setup_model()
+
+    def forward(self, x):
+        """Forward pass through the network.
+        
+        Args:
+            x: Input tensor
+            
+        Returns:
+            Output tensor
+        """
+        x = self.flatten(x)
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+class Cnn(BaseNN):
+    """Convolutional neural network model.
+    
+    Args:
+        data_splitter: DataSplitter object containing train/val/test splits
+        feature: Featurizer object for processing sequences
+        model (str): Model name identifier. Defaults to "cnn"
+        use_cache (bool): Whether to use model caching. Defaults to False
+        show_plots (bool): Whether to show matplotlib plots. Defaults to True
+        **kwargs: Additional keyword arguments including network architecture
+        
+    Example:
+        >>> splitter = DataSplitter(data, 'random')
+        >>> featurizer = OneHotFeaturizer()
+        >>> model = Cnn(splitter, featurizer, use_cache=True, show_plots=True)
+    """
+
+    def __init__(self, data_splitter, feature, model="cnn", use_cache=False, show_plots=True, **kwargs):
+
+        # Specify network architecture
+        nn_arch = [kwargs["config"]["kernel_size"]] + [
+            int(x) for x in kwargs["config"]["layersize_filtersize"].split("-")
+        ]
+
+        super().__init__(data_splitter, feature, nn_arch, model, use_cache=use_cache, show_plots=show_plots, **kwargs)
+
+        # [TODO] new function to extract input features if not already done
+        X_train_feat_example = self.nn_data_processor.featurize([self.nn_data_processor.X_train[0]])[0]
+        protein_len = X_train_feat_example.shape[0]
+        encoding_len = X_train_feat_example.shape[1]
+        kernel_size_dim1, layers, out_channels = nn_arch
+
+        in_channels = 1
+
+        self.conv_layers = nn.ModuleList()
+
+        for i in range(layers):
+            conv2d_layer = nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=(kernel_size_dim1, encoding_len if i == 0 else 1),
+                stride=(1, 1),
+            )
+            self.conv_layers.append(conv2d_layer)
+            self.conv_layers.append(nn.LeakyReLU(negative_slope=0.2))
+            in_channels = out_channels
+
+        # Dynamically calculate the input size for the fully connected layer
+        self._init_fc_layers(protein_len, encoding_len, out_channels)
+
+        # set model hyperparameters
+        self.setup_model()
+
+    def _init_fc_layers(self, protein_len, encoding_len, out_channels):
+        """Initialize fully connected layers.
+        
+        Args:
+            protein_len (int): Length of protein sequence
+            encoding_len (int): Length of sequence encoding
+            out_channels (int): Number of output channels
+        """
+        # Dummy input for calculating size
+        dummy_input = torch.randn(1, 1, protein_len, encoding_len)
+        for layer in self.conv_layers:
+            dummy_input = layer(dummy_input)
+
+        output_size = dummy_input.view(dummy_input.size(0), -1).size(1)
+
+        self.flatten = nn.Flatten()
+        self.fc1 = nn.Linear(output_size, 100)
+        self.leaky_relu = nn.LeakyReLU(negative_slope=0.2)
+        self.dropout = nn.Dropout(p=0.2)
+        self.fc2 = nn.Linear(100, 1)
+
+    def forward(self, x):
+        """Forward pass through the network.
+        
+        Args:
+            x: Input tensor
+            
+        Returns:
+            Output tensor
+        """
+        x = x.unsqueeze(1)  # add a channel dimension of 1 to the data
+        for layer in self.conv_layers:
+            x = layer(x)
+
+        x = self.flatten(x)
+        x = self.fc1(x)
+        x = self.leaky_relu(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+
+        return x

model/predictors/sweep_configs/cnn_custom_grid_sweep.yaml

@@ -0,0 +1,18 @@
+program: train.py
+method: grid
+metric:
+ name: Test Loss
+ goal: minimize
+parameters:
+ layersize_filtersize:
+   values: ["1-32", "1-128", "1-256", "1-1024"]
+ kernel_size:
+   values: [3, 7]
+ learning_rate:
+   values: [0.0001, 0.001]
+ batch_size:
+   values: [4, 8]
+ optimizer:
+   value: "adam"
+ epochs:
+   value: 300

model/predictors/sweep_configs/cnn_standard_bayes_sweep.yaml

@@ -0,0 +1,18 @@
+program: train.py
+method: bayes
+metric:
+ name: Test Loss
+ goal: minimize
+parameters:
+ layersize_filtersize:
+   values: ["1-32", "1-128", "1-256", "1-1024", "2-128", "3-128", "5-128"]
+ kernel_size:
+   values: [3, 7, 13, 17]
+ learning_rate:
+   values: [0.0001, 0.001, 0.01]
+ batch_size:
+   values: [4, 8, 16, 32]
+ optimizer:
+   value: "adam"
+ epochs:
+   value: 300

model/predictors/sweep_configs/cnn_standard_grid_sweep.yaml

@@ -0,0 +1,18 @@
+program: train.py
+method: grid
+metric:
+ name: Test Loss
+ goal: minimize
+parameters:
+ layersize_filtersize:
+   values: ["1-32", "1-128", "1-256", "1-1024", "2-128", "3-128", "5-128"]
+ kernel_size:
+   values: [3, 7, 13, 17]
+ learning_rate:
+   values: [0.0001, 0.001, 0.01]
+ batch_size:
+   values: [4, 8, 16, 32]
+ optimizer:
+   value: "adam"
+ epochs:
+   value: 300

model/predictors/sweep_configs/cnn_test_sweep.yaml

@@ -0,0 +1,18 @@
+program: train.py
+method: grid
+metric:
+ name: Test Loss
+ goal: minimize
+parameters:
+ layersize_filtersize:
+   values: ["1-32"]
+ kernel_size:
+   values: [3]
+ learning_rate:
+   values: [0.0001]
+ batch_size:
+   values: [32]
+ optimizer:
+   value: "adam"
+ epochs:
+   value: 300

model/predictors/sweep_configs/fcn_custom_grid_sweep.yaml

@@ -0,0 +1,18 @@
+program: train.py
+method: grid
+metric:
+ name: Test Loss
+ goal: minimize
+parameters:
+ layer_size:
+   value: 100
+ num_layers:
+   values: [1, 2, 3, 4, 5]
+ learning_rate:
+   values: [0.0001, 0.001, 0.01]
+ batch_size:
+   values: [16, 32, 64, 128, 256, 512]
+ optimizer:
+   value: "adam"
+ epochs:
+   value: 500

model/predictors/sweep_configs/fcn_standard_bayes_sweep.yaml

@@ -0,0 +1,18 @@
+program: train.py
+method: bayes
+metric:
+ name: Test Loss
+ goal: minimize
+parameters:
+ layer_size:
+   value: 100
+ num_layers:
+   values: [1, 2, 3, 4, 5]
+ learning_rate:
+   values: [0.0001, 0.001, 0.01]
+ batch_size:
+   values: [4, 8, 16, 32]
+ optimizer:
+   value: "adam"
+ epochs:
+   value: 300

model/predictors/sweep_configs/fcn_standard_grid_sweep.yaml

@@ -0,0 +1,18 @@
+program: train.py
+method: grid
+metric:
+ name: Test Loss
+ goal: minimize
+parameters:
+ layer_size:
+   value: 100
+ num_layers:
+   values: [1, 2, 3, 4, 5]
+ learning_rate:
+   values: [0.0001, 0.001, 0.01]
+ batch_size:
+   values: [4, 8, 16, 32]
+ optimizer:
+   value: "adam"
+ epochs:
+   value: 300

model/predictors/sweep_configs/fcn_test_sweep.yaml

@@ -0,0 +1,18 @@
+program: train.py
+method: grid
+metric:
+ name: Test Loss
+ goal: minimize
+parameters:
+ layer_size:
+   value: 100
+ num_layers:
+   values: [1]
+ learning_rate:
+   values: [0.0001]
+ batch_size:
+   values: [32]
+ optimizer:
+   value: "adam"
+ epochs:
+   value: 300

model/proposers/__init__.py

@@ -0,0 +1 @@
+from model.proposers.base_proposers import *