Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,47 @@
 *.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/src/banner_final.png filter=lfs diff=lfs merge=lfs -text

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 2024 Adibvafa Fallahpour
+
+   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.

README.md

@@ -0,0 +1,444 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- SCNet
+- DCU
+- codon-optimization
+- protein-design
+- sequence-generation
+- transformer
+frameworks:
+- PyTorch
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">CodonTransformer</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+CodonTransformer is a deep learning model for multispecies codon optimization. Given an input protein sequence and a target host organism, it generates a host-specific DNA coding sequence.
+
+Because of the degeneracy of the genetic code, the same protein can be encoded by many different DNA sequences, while different hosts have distinct preferences for synonymous codons. CodonTransformer uses a Transformer to model the contextual relationships among proteins, codons, and host organisms, generating DNA sequences that better match the natural codon distribution of the target host while preserving the translated protein sequence as much as possible.
+
+Paper:
+> **CodonTransformer: a multispecies codon optimizer using context-aware neural networks**  
+> *Nature Communications*, 2025
+
+# Model Description
+
+CodonTransformer is a Transformer model designed for multispecies codon optimization. Its core network is based on a BigBird masked language model and incorporates the STREAM representation proposed by the authors, which encodes target host information, amino acid information, and codon information into a unified sequence representation.
+
+Unlike traditional optimization methods based on global codon frequencies, CodonTransformer considers not only host preferences for individual codons but also uses the contextual modeling capability of Transformers to learn local dependencies between neighboring codons. This enables it to generate DNA sequences that more closely resemble the natural coding patterns of the target host. The model can be used for multispecies codon optimization, heterologous protein expression sequence design, and further fine-tuning on custom DNA-protein-host datasets.
+
+The official model was trained on more than one million DNA-protein paired samples spanning 164 species, including bacteria, archaea, plants, animals, and fungi.
+
+# Use Cases
+
+| Scenario | Description |
+| --- | --- |
+| Codon optimization | Redesign protein-coding DNA for a target host |
+| Heterologous protein expression | Generate coding sequences that better match host codon preferences for different hosts |
+| Multispecies sequence design | Switch the target organism among multiple supported hosts |
+| Multiple candidate sequence generation | Generate multiple distinct DNA candidate sequences through temperature sampling |
+| Batch codon optimization | Perform batch inference for multiple protein-host combinations |
+| Model fine-tuning | Continue fine-tuning using custom DNA-protein-host data |
+| Model pretraining | Train the model from scratch using large-scale processed training data |
+
+# Usage
+
+## 1. Using OneCode
+
+You can use the OneCode online environment for intelligent one-click AI4S programming:
+
+[Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+### Hardware Requirements
+
+- CodonTransformer supports inference on both CPUs and accelerator devices.
+- A CPU can be used for single protein sequence inference.
+- GPU/DCU devices are recommended for batch inference, long-sequence inference, fine-tuning, and pretraining.
+
+### Environment Setup
+
+#### DCU/SCNet Environment
+
+```bash
+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
+```
+
+- If you encounter missing dependencies or version incompatibilities during execution, refer to the dependency versions specified in `requirements.txt` at the repository root and install or adjust the corresponding dependencies as needed.
+
+### Model and Data Preparation
+
+#### 1) CodonTransformer Model Weights
+
+The official model is available on Hugging Face:
+
+```text
+https://huggingface.co/adibvafa/CodonTransformer
+```
+
+In an online environment, `from_pretrained` automatically downloads the model weights when they are not available in the local cache. In an offline environment, the weights must be downloaded in advance to the cache directory corresponding to `HF_HOME`. The recommended cache location is:
+
+```text
+/path/to/.cache/huggingface/hub/
+```
+
+Set the following environment variables when running the scripts:
+
+```bash
+export HF_HOME=/path/to/.cache/huggingface
+export HF_HUB_OFFLINE=1
+export TRANSFORMERS_OFFLINE=1
+```
+
+The inference and fine-tuning scripts in this repository support using the local cache through `HF_HOME` and `OFFLINE=1`.
+
+#### 2) Training Dataset
+
+The official training data can be obtained from Zenodo or Hugging Face Datasets:
+
+```text
+https://zenodo.org/records/12509224
+https://huggingface.co/datasets/adibvafa/CodonTransformer
+```
+
+It is recommended to place the downloaded raw data in:
+
+```text
+scripts/data/raw/
+```
+
+Accordingly, the complete dataset file is recommended to be saved as:
+
+```text
+scripts/data/raw/dataset.csv
+```
+
+The processed training JSON files are recommended to be saved in:
+
+```text
+scripts/data/processed/
+```
+## 3. Quick Start
+
+### Download the Model Package
+
+```bash
+hf download OneScience-Group/CodonTransformer \
+  --local-dir ./CodonTransformer
+
+cd CodonTransformer
+```
+
+# Inference Examples
+
+## Single-Sequence Codon Optimization
+
+Run the script:
+
+```bash
+bash scripts/slurm/run_inference_single.sh
+```
+
+The default parameters are:
+
+```bash
+PROTEIN="MFWY"
+ORGANISM="Escherichia coli general"
+OFFLINE=1
+```
+
+To change the input protein and host:
+
+```bash
+PROTEIN="MALWMRLLPLLALLALWGPDPAAA" \
+ORGANISM="Homo sapiens" \
+bash scripts/slurm/run_inference_single.sh
+```
+
+## Generate Multiple Candidate DNA Sequences
+
+Run:
+
+```bash
+bash scripts/slurm/run_inference_multiple.sh
+```
+
+By default, multiple candidate DNA sequences are generated for the same protein sequence and saved to:
+
+```text
+outputs/multiple_predictions.csv
+```
+
+The main parameters include:
+
+```text
+deterministic=False
+temperature=0.5
+top_p=0.95
+num_sequences=5
+match_protein=True
+```
+
+where:
+
+- `deterministic=False` enables probabilistic sampling.
+- `temperature` controls sampling diversity, with a typical range of `0.2 ~ 0.8`.
+- `top_p` controls nucleus sampling.
+- `num_sequences` specifies the number of candidate sequences to generate.
+- `match_protein=True` constrains the generated DNA to translate to the same input protein sequence.
+
+## Change the Target Host
+
+The target host is specified directly through the `organism` parameter, for example:
+
+```python
+organism = "Escherichia coli general"
+```
+
+It can be changed to:
+
+```python
+organism = "Homo sapiens"
+```
+
+or:
+
+```python
+organism = "Saccharomyces cerevisiae"
+```
+
+Then rerun inference to obtain a codon-optimized DNA sequence for the corresponding host.
+
+## Batch Inference
+
+The example data is currently located at:
+
+```text
+scripts/demo/sample_dataset.csv
+```
+
+Run directly:
+
+```bash
+bash scripts/slurm/run_inference_batch.sh
+```
+
+The default output is saved to:
+
+```text
+outputs/sample_predictions.csv
+```
+
+The input CSV must contain at least:
+
+```text
+protein_sequence
+organism
+```
+
+To use a custom CSV:
+
+```bash
+INPUT_CSV=/path/to/input.csv \
+OUTPUT_CSV=/path/to/output.csv \
+bash scripts/slurm/run_inference_batch.sh
+```
+
+# Training
+
+## Fine-Tuning Data Preparation
+
+For custom fine-tuning, first prepare your own CSV file. It is recommended to place it at:
+
+```text
+scripts/data/raw/your_data.csv
+```
+
+It must contain at least:
+
+```text
+dna
+protein
+organism
+```
+
+Then run the data preprocessing script provided in the repository:
+
+```bash
+INPUT_CSV=$PWD/scripts/data/raw/your_data.csv \
+OUTPUT_JSON=$PWD/scripts/data/processed/finetune_data.json \
+bash scripts/slurm/prepare_finetune_data.sh
+```
+
+## CodonTransformer Fine-Tuning
+
+Run the script:
+
+```bash
+bash scripts/slurm/run_finetune.sh
+```
+
+By default, it reads:
+
+```text
+scripts/data/processed/finetune_data.json
+```
+
+and saves checkpoints to:
+
+```text
+weight/checkpoints/finetune
+```
+
+To fine-tune using custom data, run:
+
+```bash
+DATASET_JSON=$PWD/scripts/data/processed/finetune_data.json \
+CHECKPOINT_DIR=$PWD/weight/checkpoints/finetune \
+CHECKPOINT_FILENAME=finetune.ckpt \
+BATCH_SIZE=6 \
+MAX_EPOCHS=15 \
+NUM_WORKERS=5 \
+ACCUMULATE_GRAD_BATCHES=1 \
+NUM_GPUS=4 \
+LEARNING_RATE=0.00005 \
+WARMUP_FRACTION=0.1 \
+SAVE_EVERY_N_STEPS=512 \
+SEED=123 \
+DEBUG=0 \
+bash scripts/slurm/run_finetune.sh
+```
+
+Here, `NUM_GPUS=4`, `BATCH_SIZE=6`, and `MAX_EPOCHS=15` are the default training settings. In practice, these parameters should be adjusted according to the number of allocated GPU/DCU devices, available device memory, and dataset size.
+
+## Export the Fine-Tuned Model and Run Inference
+
+After fine-tuning is complete, the following script can be used to export the checkpoint into a model file suitable for inference:
+
+```text
+scripts/slurm/export_finetuned_model.sh
+```
+
+Run it as follows:
+
+```bash
+CHECKPOINT_PATH=/path/to/finetuned_checkpoint.ckpt \
+OUTPUT_MODEL_PATH=/path/to/output_finetuned_model.pt \
+NUM_ORGANISMS=164 \
+bash scripts/slurm/export_finetuned_model.sh
+```
+
+After export, the following script can be used to load the fine-tuned model for inference:
+
+```text
+scripts/slurm/run_inference_finetuned.sh
+```
+
+Run it as follows:
+
+```bash
+PROTEIN="MFWY" \
+ORGANISM="Escherichia coli general" \
+MODEL_PATH=/path/to/output_finetuned_model.pt \
+bash scripts/slurm/run_inference_finetuned.sh
+```
+
+## Pretraining
+
+The current pretraining entry point is:
+
+```text
+scripts/pretrain.py
+```
+
+Pretraining is a complete model training workflow and requires large-scale processed DNA-protein-host data. Before full pretraining, the complete `dataset.csv` must first be converted into a JSONL file that can be read by the training script.
+
+The repository provides the following preprocessing script:
+
+```text
+scripts/slurm/prepare_pretrain_data.sh
+```
+
+By default, this script reads:
+
+```text
+scripts/data/raw/dataset.csv
+```
+
+and outputs:
+
+```text
+scripts/data/processed/pretrain_data.json
+```
+
+Therefore, before full pretraining, first run:
+
+```bash
+bash scripts/slurm/prepare_pretrain_data.sh
+```
+
+If the raw data or output directory differs from the default path, modify it using environment variables:
+
+```bash
+INPUT_CSV=$PWD/scripts/data/raw/dataset.csv \
+OUTPUT_JSON=$PWD/scripts/data/processed/pretrain_data.json \
+bash scripts/slurm/prepare_pretrain_data.sh
+```
+
+The complete generated pretraining data is recommended to be saved as:
+
+```text
+scripts/data/processed/pretrain_data.json
+```
+
+Full pretraining can be started with:
+
+```bash
+TRAIN_DATA_PATH=$PWD/scripts/data/processed/pretrain_data.json \
+CHECKPOINT_DIR=$PWD/weight/checkpoints/pretrain \
+BATCH_SIZE=6 \
+MAX_EPOCHS=5 \
+NUM_WORKERS=5 \
+ACCUMULATE_GRAD_BATCHES=1 \
+NUM_GPUS=16 \
+LEARNING_RATE=0.00005 \
+WARMUP_FRACTION=0.1 \
+SAVE_INTERVAL=5 \
+SEED=123 \
+DEBUG=0 \
+bash scripts/slurm/run_pretrain.sh
+```
+
+Here, `NUM_GPUS=16` corresponds to the default multi-device setting in the pretraining script. In practice, adjust `NUM_GPUS` according to the number of allocated GPU/DCU devices. Pretraining checkpoints are saved by default to:
+
+```text
+weight/checkpoints/pretrain
+```
+# OneScience Official Information
+
+| Platform | OneScience Main 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
+
+- The official CodonTransformer source code repository is licensed under the **Apache License 2.0**. See the `LICENSE` file in the repository root for details.
+- CodonTransformer model weights are distributed separately through Hugging Face, while the training data is available through Zenodo and Hugging Face Datasets. The model weights, training data, and related third-party resources must each be used in accordance with the licenses and terms of use specified on their respective pages.
+- This repository is the **DCU-adapted version** of CodonTransformer. The use of the repository code, model weights, and related data remains subject to the licenses and terms of use of their respective original projects.

conf/CODE_OF_CONDUCT.md

@@ -0,0 +1,128 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, religion, or sexual identity
+and orientation.
+
+We pledge to act and interact in ways that contribute to an open, welcoming,
+diverse, inclusive, and healthy community.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment for our
+community include:
+
+* Demonstrating empathy and kindness toward other people
+* Being respectful of differing opinions, viewpoints, and experiences
+* Giving and gracefully accepting constructive feedback
+* Accepting responsibility and apologizing to those affected by our mistakes,
+  and learning from the experience
+* Focusing on what is best not just for us as individuals, but for the
+  overall community
+
+Examples of unacceptable behavior include:
+
+* The use of sexualized language or imagery, and sexual attention or
+  advances of any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or email
+  address, without their explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Enforcement Responsibilities
+
+Community leaders are responsible for clarifying and enforcing our standards of
+acceptable behavior and will take appropriate and fair corrective action in
+response to any behavior that they deem inappropriate, threatening, offensive,
+or harmful.
+
+Community leaders have the right and responsibility to remove, edit, or reject
+comments, commits, code, wiki edits, issues, and other contributions that are
+not aligned to this Code of Conduct, and will communicate reasons for moderation
+decisions when appropriate.
+
+## Scope
+
+This Code of Conduct applies within all community spaces, and also applies when
+an individual is officially representing the community in public spaces.
+Examples of representing our community include using an official e-mail address,
+posting via an official social media account, or acting as an appointed
+representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the community leaders responsible for enforcement at
+Adibvafa.fallahpour@mail.utoronto.ca.
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+**Community Impact**: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+**Consequence**: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+**Community Impact**: A violation through a single incident or series
+of actions.
+
+**Consequence**: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or
+permanent ban.
+
+### 3. Temporary Ban
+
+**Community Impact**: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+**Consequence**: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+**Community Impact**: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior,  harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+**Consequence**: A permanent ban from any sort of public interaction within
+the community.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 2.0, available at
+https://www.contributor-covenant.org/version/2/0/code_of_conduct.html.
+
+Community Impact Guidelines were inspired by [Mozilla's code of conduct
+enforcement ladder](https://github.com/mozilla/diversity).
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see the FAQ at
+https://www.contributor-covenant.org/faq. Translations are available at
+https://www.contributor-covenant.org/translations.

conf/Makefile

@@ -0,0 +1,9 @@
+# Makefile
+
+.PHONY: test
+test:
+	python -m unittest discover -s tests
+
+.PHONY: test_with_coverage
+test_with_coverage:
+	coverage run -m unittest discover -s tests

config.json

@@ -0,0 +1,28 @@
+{
+  "name": "CodonTransformer",
+  "version": "1.6.7",
+  "layout": {
+    "conf": "Repository configuration and auxiliary project files.",
+    "model": "CodonTransformer Python package and bundled model resources.",
+    "scripts": "Executable training, inference, SLURM, notebook, test, and demo files.",
+    "weight": "Local model weights and checkpoints."
+  },
+  "entrypoints": {
+    "prepare_pretrain_data": "scripts/slurm/prepare_pretrain_data.sh",
+    "prepare_finetune_data": "scripts/slurm/prepare_finetune_data.sh",
+    "export_finetuned_model": "scripts/slurm/export_finetuned_model.sh",
+    "inference_finetuned": "scripts/slurm/run_inference_finetuned.sh",
+    "finetune": "scripts/finetune.py",
+    "pretrain": "scripts/pretrain.py",
+    "inference_single": "scripts/slurm/run_inference_single.sh",
+    "inference_multiple": "scripts/slurm/run_inference_multiple.sh",
+    "inference_batch": "scripts/slurm/run_inference_batch.sh"
+  },
+  "resources": {
+    "tokenizer": "model/src/CodonTransformerTokenizer.json",
+    "demo_dataset": "scripts/demo/sample_dataset.csv",
+    "raw_training_data": "scripts/data/raw",
+    "processed_training_data": "scripts/data/processed",
+    "checkpoints": "weight/checkpoints"
+  }
+}

model/CodonTransformer/CodonData.py

@@ -0,0 +1,682 @@
+"""
+File: CodonData.py
+---------------------
+Includes helper functions for preprocessing NCBI or Kazusa databases and
+preparing the data for training and inference of the CodonTransformer model.
+"""
+
+import json
+import os
+import random
+from typing import Dict, List, Optional, Tuple, Union
+
+import pandas as pd
+import python_codon_tables as pct
+from Bio import SeqIO
+from Bio.Seq import Seq
+from sklearn.utils import shuffle as sk_shuffle
+from tqdm import tqdm
+
+from CodonTransformer.CodonUtils import (
+    AMBIGUOUS_AMINOACID_MAP,
+    AMINO2CODON_TYPE,
+    AMINO_ACIDS,
+    ORGANISM2ID,
+    START_CODONS,
+    STOP_CODONS,
+    STOP_SYMBOL,
+    STOP_SYMBOLS,
+    ProteinConfig,
+    find_pattern_in_fasta,
+    get_taxonomy_id,
+    sort_amino2codon_skeleton,
+)
+
+
+def prepare_training_data(
+    dataset: Union[str, pd.DataFrame], output_file: str, shuffle: bool = True
+) -> None:
+    """
+    Prepare a JSON dataset for training the CodonTransformer model.
+
+    Input dataset should have columns below:
+        - dna: str (DNA sequence)
+        - protein: str (Protein sequence)
+        - organism: Union[int, str] (ID or Name of the organism)
+
+    The output JSON dataset will have the following format:
+        {"idx": 0, "codons": "M_ATG R_AGG L_TTG L_CTA R_CGA __TAG", "organism": 51}
+        {"idx": 1, "codons": "M_ATG K_AAG C_TGC F_TTT F_TTC __TAA", "organism": 59}
+
+    Args:
+        dataset (Union[str, pd.DataFrame]): Input dataset in CSV or DataFrame format.
+        output_file (str): Path to save the output JSON dataset.
+        shuffle (bool, optional): Whether to shuffle the dataset before saving.
+            Defaults to True.
+
+    Returns:
+        None
+    """
+    if isinstance(dataset, str):
+        dataset = pd.read_csv(dataset)
+
+    required_columns = {"dna", "protein", "organism"}
+    if not required_columns.issubset(dataset.columns):
+        raise ValueError(f"Input dataset must have columns: {required_columns}")
+
+    # Prepare the dataset for finetuning
+    dataset["codons"] = dataset.apply(
+        lambda row: get_merged_seq(row["protein"], row["dna"], separator="_"), axis=1
+    )
+
+    # Replace organism str with organism id using ORGANISM2ID
+    dataset["organism"] = dataset["organism"].apply(
+        lambda org: process_organism(org, ORGANISM2ID)
+    )
+
+    # Save the dataset to a JSON file
+    dataframe_to_json(dataset[["codons", "organism"]], output_file, shuffle=shuffle)
+
+
+def dataframe_to_json(df: pd.DataFrame, output_file: str, shuffle: bool = True) -> None:
+    """
+    Convert pandas DataFrame to JSON file format suitable for training CodonTransformer.
+
+    This function takes a preprocessed DataFrame and writes it to a JSON file
+    where each line is a JSON object representing a single record.
+
+    Args:
+        df (pd.DataFrame): The input DataFrame with 'codons' and 'organism' columns.
+        output_file (str): Path to the output JSON file.
+        shuffle (bool, optional): Whether to shuffle the dataset before saving.
+            Defaults to True.
+
+    Returns:
+        None
+
+    Raises:
+        ValueError: If the required columns are not present in the DataFrame.
+    """
+    required_columns = {"codons", "organism"}
+    if not required_columns.issubset(df.columns):
+        raise ValueError(f"DataFrame must contain columns: {required_columns}")
+
+    print(f"\nStarted writing to {output_file}...")
+
+    # Shuffle the DataFrame if requested
+    if shuffle:
+        df = sk_shuffle(df)
+
+    # Write the DataFrame to a JSON file
+    with open(output_file, "w") as f:
+        for idx, row in tqdm(
+            df.iterrows(), total=len(df), desc="Writing JSON...", unit=" records"
+        ):
+            doc = {"idx": idx, "codons": row["codons"], "organism": row["organism"]}
+            f.write(json.dumps(doc) + "\n")
+
+    print(f"\nTotal Entries Saved: {len(df)}, JSON data saved to {output_file}")
+
+
+def process_organism(organism: Union[str, int], organism_to_id: Dict[str, int]) -> int:
+    """
+    Process and validate the organism input, converting it to a valid organism ID.
+
+    This function handles both string (organism name) and integer (organism ID) inputs.
+    It validates the input against a provided mapping of organism names to IDs.
+
+    Args:
+        organism (Union[str, int]): Input organism, either as a name (str) or ID (int).
+        organism_to_id (Dict[str, int]): Dictionary mapping organism names to their
+            corresponding IDs.
+
+    Returns:
+        int: The validated organism ID.
+
+    Raises:
+        ValueError: If the input is an invalid organism name or ID.
+        TypeError: If the input is neither a string nor an integer.
+    """
+    if isinstance(organism, str):
+        if organism not in organism_to_id:
+            raise ValueError(f"Invalid organism name: {organism}")
+        return organism_to_id[organism]
+
+    elif isinstance(organism, int):
+        if organism not in organism_to_id.values():
+            raise ValueError(f"Invalid organism ID: {organism}")
+        return organism
+
+    raise TypeError(
+        f"Organism must be a string or integer, not {type(organism).__name__}"
+    )
+
+
+def preprocess_protein_sequence(protein: str) -> str:
+    """
+    Preprocess a protein sequence by cleaning, standardizing, and handling
+    ambiguous amino acids.
+
+    Args:
+        protein (str): The input protein sequence.
+
+    Returns:
+        str: The preprocessed protein sequence.
+
+    Raises:
+        ValueError: If the protein sequence is invalid or if the configuration is invalid.
+    """
+    if not protein:
+        raise ValueError("Protein sequence is empty.")
+
+    # Clean and standardize the protein sequence
+    protein = (
+        protein.upper().strip().replace("\n", "").replace(" ", "").replace("\t", "")
+    )
+
+    # Handle ambiguous amino acids based on the specified behavior
+    config = ProteinConfig()
+    ambiguous_aminoacid_map_override = config.get("ambiguous_aminoacid_map_override")
+    ambiguous_aminoacid_behavior = config.get("ambiguous_aminoacid_behavior")
+    ambiguous_aminoacid_map = AMBIGUOUS_AMINOACID_MAP.copy()
+
+    for aminoacid, standard_aminoacids in ambiguous_aminoacid_map_override.items():
+        ambiguous_aminoacid_map[aminoacid] = standard_aminoacids
+
+    if ambiguous_aminoacid_behavior == "raise_error":
+        if any(aminoacid in ambiguous_aminoacid_map for aminoacid in protein):
+            raise ValueError("Ambiguous amino acids found in protein sequence.")
+    elif ambiguous_aminoacid_behavior == "standardize_deterministic":
+        protein = "".join(
+            ambiguous_aminoacid_map.get(aminoacid, [aminoacid])[0]
+            for aminoacid in protein
+        )
+    elif ambiguous_aminoacid_behavior == "standardize_random":
+        protein = "".join(
+            random.choice(ambiguous_aminoacid_map.get(aminoacid, [aminoacid]))
+            for aminoacid in protein
+        )
+    else:
+        raise ValueError(
+            f"Invalid ambiguous_aminoacid_behavior: {ambiguous_aminoacid_behavior}."
+        )
+
+    # Check for sequence validity
+    if any(aminoacid not in AMINO_ACIDS + STOP_SYMBOLS for aminoacid in protein):
+        raise ValueError("Invalid characters in protein sequence.")
+
+    if protein[-1] not in AMINO_ACIDS + STOP_SYMBOLS:
+        raise ValueError(
+            "Protein sequence must end with `*`, or `_`, or an amino acid."
+        )
+
+    # Replace '*' at the end of protein with STOP_SYMBOL if present
+    if protein[-1] == "*":
+        protein = protein[:-1] + STOP_SYMBOL
+
+    # Add stop symbol to end of protein
+    if protein[-1] != STOP_SYMBOL:
+        protein += STOP_SYMBOL
+
+    return protein
+
+
+def replace_ambiguous_codons(dna: str) -> str:
+    """
+    Replaces ambiguous codons in a DNA sequence with "UNK".
+
+    Args:
+        dna (str): The DNA sequence to process.
+
+    Returns:
+        str: The processed DNA sequence with ambiguous codons replaced by "UNK".
+    """
+    result = []
+    dna = dna.upper()
+
+    # Check codons in DNA sequence
+    for i in range(0, len(dna), 3):
+        codon = dna[i : i + 3]
+
+        if len(codon) == 3 and all(nucleotide in "ATCG" for nucleotide in codon):
+            result.append(codon)
+        else:
+            result.append("UNK")
+
+    return "".join(result)
+
+
+def preprocess_dna_sequence(dna: str) -> str:
+    """
+    Cleans and preprocesses a DNA sequence by standardizing it and replacing
+    ambiguous codons.
+
+    Args:
+        dna (str): The DNA sequence to preprocess.
+
+    Returns:
+        str: The cleaned and preprocessed DNA sequence.
+    """
+    if not dna:
+        return ""
+
+    # Clean and standardize the DNA sequence
+    dna = dna.upper().strip().replace("\n", "").replace(" ", "").replace("\t", "")
+
+    # Replace codons with ambigous nucleotides with "UNK"
+    dna = replace_ambiguous_codons(dna)
+
+    # Add unkown stop codon to end of DNA sequence if not present
+    if dna[-3:] not in STOP_CODONS:
+        dna += "UNK"
+
+    return dna
+
+
+def get_merged_seq(protein: str, dna: str = "", separator: str = "_") -> str:
+    """
+    Return the merged sequence of protein amino acids and DNA codons in the form
+    of tokens separated by space, where each token is composed of an amino acid +
+    separator + codon.
+
+    Args:
+        protein (str): Protein sequence.
+        dna (str): DNA sequence.
+        separator (str): Separator between amino acid and codon.
+
+    Returns:
+        str: Merged sequence.
+
+    Example:
+        >>> get_merged_seq(protein="MAV_", dna="ATGGCTGTGTAA", separator="_")
+        'M_ATG A_GCT V_GTG __TAA'
+
+        >>> get_merged_seq(protein="QHH_", dna="", separator="_")
+        'Q_UNK H_UNK H_UNK __UNK'
+    """
+    merged_seq = ""
+
+    # Prepare protein and dna sequences
+    dna = preprocess_dna_sequence(dna)
+    protein = preprocess_protein_sequence(protein)
+
+    # Check if the length of protein and dna sequences are equal
+    if len(dna) > 0 and len(protein) != len(dna) / 3:
+        raise ValueError(
+            'Length of protein (including stop symbol such as "_") and '
+            "the number of codons in DNA sequence (including stop codon) "
+            "must be equal."
+        )
+
+    # Merge protein and DNA sequences into tokens
+    for i, aminoacid in enumerate(protein):
+        merged_seq += f'{aminoacid}{separator}{dna[i * 3:i * 3 + 3] if dna else "UNK"} '
+
+    return merged_seq.strip()
+
+
+def is_correct_seq(dna: str, protein: str, stop_symbol: str = STOP_SYMBOL) -> bool:
+    """
+    Check if the given DNA and protein pair is correct, that is:
+        1. The length of dna is divisible by 3
+        2. There is an initiator codon in the beginning of dna
+        3. There is only one stop codon in the sequence
+        4. The only stop codon is the last codon
+
+    Note since in Codon Table 3, 'TGA' is interpreted as Triptophan (W),
+    there is a separate check to make sure those sequences are considered correct.
+
+    Args:
+        dna (str): DNA sequence.
+        protein (str): Protein sequence.
+        stop_symbol (str): Stop symbol.
+
+    Returns:
+        bool: True if the sequence is correct, False otherwise.
+    """
+    return (
+        len(dna) % 3 == 0  # Check if DNA length is divisible by 3
+        and dna[:3].upper() in START_CODONS  # Check for initiator codon
+        and protein[-1]
+        == stop_symbol  # Check if the last protein symbol is the stop symbol
+        and protein.count(stop_symbol) == 1  # Check if there is only one stop symbol
+        and len(set(dna))
+        == 4  # Check if DNA consists of 4 unique nucleotides (A, T, C, G)
+    )
+
+
+def get_amino_acid_sequence(
+    dna: str,
+    stop_symbol: str = "_",
+    codon_table: int = 1,
+    return_correct_seq: bool = False,
+) -> Union[str, Tuple[str, bool]]:
+    """
+    Return the translated protein sequence given a DNA sequence and codon table.
+
+    Args:
+        dna (str): DNA sequence.
+        stop_symbol (str): Stop symbol.
+        codon_table (int): Codon table number.
+        return_correct_seq (bool): Whether to return if the sequence is correct.
+
+    Returns:
+        Union[str, Tuple[str, bool]]: Protein sequence and correctness flag if
+        return_correct_seq is True, otherwise just the protein sequence.
+    """
+    dna_seq = Seq(dna).strip()
+
+    # Translate the DNA sequence to a protein sequence
+    protein_seq = str(
+        dna_seq.translate(
+            stop_symbol=stop_symbol,  # Symbol to use for stop codons
+            to_stop=False,  # Translate the entire sequence, including any stop codons
+            cds=False,  # Do not assume the input is a coding sequence
+            table=codon_table,  # Codon table to use for translation
+        )
+    ).strip()
+
+    return (
+        protein_seq
+        if not return_correct_seq
+        else (protein_seq, is_correct_seq(dna_seq, protein_seq, stop_symbol))
+    )
+
+
+def read_fasta_file(
+    input_file: str,
+    save_to_file: Optional[str] = None,
+    organism: str = "",
+    buffer_size: int = 50000,
+) -> pd.DataFrame:
+    """
+    Read a FASTA file of DNA sequences and convert it to a Pandas DataFrame.
+    Optionally, save the DataFrame to a CSV file.
+
+    Args:
+        input_file (str): Path to the input FASTA file.
+        save_to_file (Optional[str]): Path to save the output DataFrame. If None,
+            data is only returned.
+        organism (str): Name of the organism. If empty, it will be extracted from
+            the FASTA description.
+        buffer_size (int): Number of records to process before writing to file.
+
+    Returns:
+        pd.DataFrame: DataFrame containing the DNA sequences if return_dataframe
+        is True, else None.
+
+    Raises:
+        FileNotFoundError: If the input file does not exist.
+    """
+    if not os.path.exists(input_file):
+        raise FileNotFoundError(f"Input file not found: {input_file}")
+
+    buffer = []
+    columns = [
+        "dna",
+        "protein",
+        "correct_seq",
+        "organism",
+        "GeneID",
+        "description",
+        "tokenized",
+    ]
+
+    # Initialize DataFrame to store all data if return_dataframe is True
+    all_data = pd.DataFrame(columns=columns)
+
+    with open(input_file, "r") as fasta_file:
+        for record in tqdm(
+            SeqIO.parse(fasta_file, "fasta"),
+            desc=f"Processing {organism}",
+            unit=" Records",
+        ):
+            dna = str(record.seq).strip().upper()  # Ensure uppercase DNA sequence
+
+            # Determine the organism from the record if not provided
+            current_organism = organism or find_pattern_in_fasta(
+                "organism", record.description
+            )
+            gene_id = find_pattern_in_fasta("GeneID", record.description)
+
+            # Get the appropriate codon table for the organism
+            codon_table = get_codon_table(current_organism)
+
+            # Translate DNA to protein sequence
+            protein, correct_seq = get_amino_acid_sequence(
+                dna,
+                stop_symbol=STOP_SYMBOL,
+                codon_table=codon_table,
+                return_correct_seq=True,
+            )
+            description = record.description.split("[", 1)[0].strip()
+            tokenized = get_merged_seq(protein, dna, separator=STOP_SYMBOL)
+
+            # Create a data row for the current sequence
+            data_row = {
+                "dna": dna,
+                "protein": protein,
+                "correct_seq": correct_seq,
+                "organism": current_organism,
+                "GeneID": gene_id,
+                "description": description,
+                "tokenized": tokenized,
+            }
+            buffer.append(data_row)
+
+            # Write buffer to CSV file when buffer size is reached
+            if save_to_file and len(buffer) >= buffer_size:
+                write_buffer_to_csv(buffer, save_to_file, columns)
+                buffer = []
+
+            all_data = pd.concat(
+                [all_data, pd.DataFrame([data_row])], ignore_index=True
+            )
+
+    # Write remaining buffer to CSV file
+    if save_to_file and buffer:
+        write_buffer_to_csv(buffer, save_to_file, columns)
+
+    return all_data
+
+
+def write_buffer_to_csv(buffer: List[Dict], output_path: str, columns: List[str]):
+    """Helper function to write buffer to CSV file."""
+    buffer_df = pd.DataFrame(buffer, columns=columns)
+    buffer_df.to_csv(
+        output_path,
+        mode="a",
+        header=(not os.path.exists(output_path)),
+        index=True,
+    )
+
+
+def download_codon_frequencies_from_kazusa(
+    taxonomy_id: Optional[int] = None,
+    organism: Optional[str] = None,
+    taxonomy_reference: Optional[str] = None,
+    return_original_format: bool = False,
+) -> AMINO2CODON_TYPE:
+    """
+    Return the codon table of the given taxonomy ID from the Kazusa Database.
+
+    Args:
+        taxonomy_id (Optional[int]): Taxonomy ID.
+        organism (Optional[str]): Name of the organism.
+        taxonomy_reference (Optional[str]): Taxonomy reference.
+        return_original_format (bool): Whether to return in the original format.
+
+    Returns:
+        AMINO2CODON_TYPE: Codon table.
+    """
+    if taxonomy_reference:
+        taxonomy_id = get_taxonomy_id(taxonomy_reference, organism=organism)
+
+    kazusa_amino2codon = pct.get_codons_table(table_name=taxonomy_id)
+
+    if return_original_format:
+        return kazusa_amino2codon
+
+    # Replace "*" with STOP_SYMBOL in the codon table
+    kazusa_amino2codon[STOP_SYMBOL] = kazusa_amino2codon.pop("*")
+
+    # Create amino2codon dictionary
+    amino2codon = {
+        aminoacid: (list(codon2freq.keys()), list(codon2freq.values()))
+        for aminoacid, codon2freq in kazusa_amino2codon.items()
+    }
+
+    return sort_amino2codon_skeleton(amino2codon)
+
+
+def build_amino2codon_skeleton(organism: str) -> AMINO2CODON_TYPE:
+    """
+    Return the empty skeleton of the amino2codon dictionary, needed for
+    get_codon_frequencies.
+
+    Args:
+        organism (str): Name of the organism.
+
+    Returns:
+        AMINO2CODON_TYPE: Empty amino2codon dictionary.
+    """
+    amino2codon = {}
+    possible_codons = [f"{i}{j}{k}" for i in "ACGT" for j in "ACGT" for k in "ACGT"]
+    possible_aminoacids = get_amino_acid_sequence(
+        dna="".join(possible_codons),
+        codon_table=get_codon_table(organism),
+        return_correct_seq=False,
+    )
+
+    # Initialize the amino2codon skeleton with all possible codons and set their
+    # frequencies to 0
+    for i, (codon, amino) in enumerate(zip(possible_codons, possible_aminoacids)):
+        if amino not in amino2codon:
+            amino2codon[amino] = ([], [])
+
+        amino2codon[amino][0].append(codon)
+        amino2codon[amino][1].append(0)
+
+    # Sort the dictionary and each list of codon frequency alphabetically
+    amino2codon = sort_amino2codon_skeleton(amino2codon)
+
+    return amino2codon
+
+
+def get_codon_frequencies(
+    dna_sequences: List[str],
+    protein_sequences: Optional[List[str]] = None,
+    organism: Optional[str] = None,
+) -> AMINO2CODON_TYPE:
+    """
+    Return a dictionary mapping each codon to its respective frequency based on
+    the collection of DNA sequences and protein sequences.
+
+    Args:
+        dna_sequences (List[str]): List of DNA sequences.
+        protein_sequences (Optional[List[str]]): List of protein sequences.
+        organism (Optional[str]): Name of the organism.
+
+    Returns:
+        AMINO2CODON_TYPE: Dictionary mapping each amino acid to a tuple of codons
+        and frequencies.
+    """
+    if organism:
+        codon_table = get_codon_table(organism)
+        protein_sequences = [
+            get_amino_acid_sequence(
+                dna, codon_table=codon_table, return_correct_seq=False
+            )
+            for dna in dna_sequences
+        ]
+
+    amino2codon = build_amino2codon_skeleton(organism)
+
+    # Count the frequencies of each codon for each amino acid
+    for dna, protein in zip(dna_sequences, protein_sequences):
+        for i, amino in enumerate(protein):
+            codon = dna[i * 3 : (i + 1) * 3]
+            codon_loc = amino2codon[amino][0].index(codon)
+            amino2codon[amino][1][codon_loc] += 1
+
+    # Normalize codon frequencies per amino acid so they sum to 1
+    amino2codon = {
+        amino: (codons, [freq / (sum(frequencies) + 1e-100) for freq in frequencies])
+        for amino, (codons, frequencies) in amino2codon.items()
+    }
+
+    return amino2codon
+
+
+def get_organism_to_codon_frequencies(
+    dataset: pd.DataFrame, organisms: List[str]
+) -> Dict[str, AMINO2CODON_TYPE]:
+    """
+    Return a dictionary mapping each organism to their codon frequency distribution.
+
+    Args:
+        dataset (pd.DataFrame): DataFrame containing DNA sequences.
+        organisms (List[str]): List of organisms.
+
+    Returns:
+        Dict[str, AMINO2CODON_TYPE]: Dictionary mapping each organism to its codon
+        frequency distribution.
+    """
+    organism2frequencies = {}
+
+    # Calculate codon frequencies for each organism in the dataset
+    for organism in tqdm(
+        organisms, desc="Calculating Codon Frequencies: ", unit="Organism"
+    ):
+        organism_data = dataset.loc[dataset["organism"] == organism]
+
+        dna_sequences = organism_data["dna"].to_list()
+        protein_sequences = organism_data["protein"].to_list()
+
+        codon_frequencies = get_codon_frequencies(dna_sequences, protein_sequences)
+        organism2frequencies[organism] = codon_frequencies
+
+    return organism2frequencies
+
+
+def get_codon_table(organism: str) -> int:
+    """
+    Return the appropriate NCBI codon table for a given organism.
+
+    Args:
+        organism (str): Name of the organism.
+
+    Returns:
+        int: Codon table number.
+    """
+    # Common codon table (Table 1) for many model organisms
+    if organism in [
+        "Arabidopsis thaliana",
+        "Caenorhabditis elegans",
+        "Chlamydomonas reinhardtii",
+        "Saccharomyces cerevisiae",
+        "Danio rerio",
+        "Drosophila melanogaster",
+        "Homo sapiens",
+        "Mus musculus",
+        "Nicotiana tabacum",
+        "Solanum tuberosum",
+        "Solanum lycopersicum",
+        "Oryza sativa",
+        "Glycine max",
+        "Zea mays",
+    ]:
+        codon_table = 1
+
+    # Chloroplast codon table (Table 11)
+    elif organism in [
+        "Chlamydomonas reinhardtii chloroplast",
+        "Nicotiana tabacum chloroplast",
+    ]:
+        codon_table = 11
+
+    # Default to Table 11 for other bacteria and archaea
+    else:
+        codon_table = 11
+
+    return codon_table

model/CodonTransformer/CodonEvaluation.py

@@ -0,0 +1,278 @@
+"""
+File: CodonEvaluation.py
+---------------------------
+Includes functions to calculate various evaluation metrics along with helper
+functions.
+"""
+
+from typing import Dict, List, Tuple
+
+import pandas as pd
+from CAI import CAI, relative_adaptiveness
+from tqdm import tqdm
+
+
+def get_CSI_weights(sequences: List[str]) -> Dict[str, float]:
+    """
+    Calculate the Codon Similarity Index (CSI) weights for a list of DNA sequences.
+
+    Args:
+        sequences (List[str]): List of DNA sequences.
+
+    Returns:
+        dict: The CSI weights.
+    """
+    return relative_adaptiveness(sequences=sequences)
+
+
+def get_CSI_value(dna: str, weights: Dict[str, float]) -> float:
+    """
+    Calculate the Codon Similarity Index (CSI) for a DNA sequence.
+
+    Args:
+        dna (str): The DNA sequence.
+        weights (dict): The CSI weights from get_CSI_weights.
+
+    Returns:
+        float: The CSI value.
+    """
+    return CAI(dna, weights)
+
+
+def get_organism_to_CSI_weights(
+    dataset: pd.DataFrame, organisms: List[str]
+) -> Dict[str, dict]:
+    """
+    Calculate the Codon Similarity Index (CSI) weights for a list of organisms.
+
+    Args:
+        dataset (pd.DataFrame): Dataset containing organism and DNA sequence info.
+        organisms (List[str]): List of organism names.
+
+    Returns:
+        Dict[str, dict]: A dictionary mapping each organism to its CSI weights.
+    """
+    organism2weights = {}
+
+    # Iterate through each organism to calculate its CSI weights
+    for organism in tqdm(organisms, desc="Calculating CSI Weights: ", unit="Organism"):
+        organism_data = dataset.loc[dataset["organism"] == organism]
+        sequences = organism_data["dna"].to_list()
+        weights = get_CSI_weights(sequences)
+        organism2weights[organism] = weights
+
+    return organism2weights
+
+
+def get_GC_content(dna: str, lower: bool = False) -> float:
+    """
+    Calculate the GC content of a DNA sequence.
+
+    Args:
+        dna (str): The DNA sequence.
+        lower (bool): If True, converts DNA sequence to lowercase before calculation.
+
+    Returns:
+        float: The GC content as a percentage.
+    """
+    if lower:
+        dna = dna.lower()
+    return (dna.count("G") + dna.count("C")) / len(dna) * 100
+
+
+def get_cfd(
+    dna: str,
+    codon_frequencies: Dict[str, Tuple[List[str], List[float]]],
+    threshold: float = 0.3,
+) -> float:
+    """
+    Calculate the codon frequency distribution (CFD) metric for a DNA sequence.
+
+    Args:
+        dna (str): The DNA sequence.
+        codon_frequencies (Dict[str, Tuple[List[str], List[float]]]): Codon
+            frequency distribution per amino acid.
+        threshold (float): Frequency threshold for counting rare codons.
+
+    Returns:
+        float: The CFD metric as a percentage.
+    """
+    # Get a dictionary mapping each codon to its normalized frequency
+    codon2frequency = {
+        codon: freq / max(frequencies)
+        for amino, (codons, frequencies) in codon_frequencies.items()
+        for codon, freq in zip(codons, frequencies)
+    }
+
+    cfd = 0
+
+    # Iterate through the DNA sequence in steps of 3 to process each codon
+    for i in range(0, len(dna), 3):
+        codon = dna[i : i + 3]
+        codon_frequency = codon2frequency[codon]
+
+        if codon_frequency < threshold:
+            cfd += 1
+
+    return cfd / (len(dna) / 3) * 100
+
+
+def get_min_max_percentage(
+    dna: str,
+    codon_frequencies: Dict[str, Tuple[List[str], List[float]]],
+    window_size: int = 18,
+) -> List[float]:
+    """
+    Calculate the %MinMax metric for a DNA sequence.
+
+    Args:
+        dna (str): The DNA sequence.
+        codon_frequencies (Dict[str, Tuple[List[str], List[float]]]): Codon
+            frequency distribution per amino acid.
+        window_size (int): Size of the window to calculate %MinMax.
+
+    Returns:
+        List[float]: List of %MinMax values for the sequence.
+
+    Credit: https://github.com/chowington/minmax
+    """
+    # Get a dictionary mapping each codon to its respective amino acid
+    codon2amino = {
+        codon: amino
+        for amino, (codons, frequencies) in codon_frequencies.items()
+        for codon in codons
+    }
+
+    min_max_values = []
+    codons = [dna[i : i + 3] for i in range(0, len(dna), 3)]  # Split DNA into codons
+
+    # Iterate through the DNA sequence using the specified window size
+    for i in range(len(codons) - window_size + 1):
+        codon_window = codons[i : i + window_size]  # Codons in the current window
+
+        Actual = 0.0  # Average of the actual codon frequencies
+        Max = 0.0  # Average of the min codon frequencies
+        Min = 0.0  # Average of the max codon frequencies
+        Avg = 0.0  # Average of the averages of all frequencies for each amino acid
+
+        # Sum the frequencies for codons in the current window
+        for codon in codon_window:
+            aminoacid = codon2amino[codon]
+            frequencies = codon_frequencies[aminoacid][1]
+            codon_index = codon_frequencies[aminoacid][0].index(codon)
+            codon_frequency = codon_frequencies[aminoacid][1][codon_index]
+
+            Actual += codon_frequency
+            Max += max(frequencies)
+            Min += min(frequencies)
+            Avg += sum(frequencies) / len(frequencies)
+
+        # Divide by the window size to get the averages
+        Actual = Actual / window_size
+        Max = Max / window_size
+        Min = Min / window_size
+        Avg = Avg / window_size
+
+        # Calculate %MinMax
+        percentMax = ((Actual - Avg) / (Max - Avg)) * 100
+        percentMin = ((Avg - Actual) / (Avg - Min)) * 100
+
+        # Append the appropriate %MinMax value
+        if percentMax >= 0:
+            min_max_values.append(percentMax)
+        else:
+            min_max_values.append(-percentMin)
+
+    # Populate the last floor(window_size / 2) entries of min_max_values with None
+    for i in range(int(window_size / 2)):
+        min_max_values.append(None)
+
+    return min_max_values
+
+
+def get_sequence_complexity(dna: str) -> float:
+    """
+    Calculate the sequence complexity score of a DNA sequence.
+
+    Args:
+        dna (str): The DNA sequence.
+
+    Returns:
+        float: The sequence complexity score.
+    """
+
+    def sum_up_to(x):
+        """Recursive function to calculate the sum of integers from 1 to x."""
+        if x <= 1:
+            return 1
+        else:
+            return x + sum_up_to(x - 1)
+
+    def f(x):
+        """Returns 4 if x is greater than or equal to 4, else returns x."""
+        if x >= 4:
+            return 4
+        elif x < 4:
+            return x
+
+    unique_subseq_length = []
+
+    # Calculate unique subsequences lengths
+    for i in range(1, len(dna) + 1):
+        unique_subseq = set()
+        for j in range(len(dna) - (i - 1)):
+            unique_subseq.add(dna[j : (j + i)])
+        unique_subseq_length.append(len(unique_subseq))
+
+    # Calculate complexity score
+    complexity_score = (
+        sum(unique_subseq_length) / (sum_up_to(len(dna) - 1) + f(len(dna)))
+    ) * 100
+
+    return complexity_score
+
+
+def get_sequence_similarity(
+    original: str, predicted: str, truncate: bool = True, window_length: int = 1
+) -> float:
+    """
+    Calculate the sequence similarity between two sequences.
+
+    Args:
+        original (str): The original sequence.
+        predicted (str): The predicted sequence.
+        truncate (bool): If True, truncate the original sequence to match the length
+            of the predicted sequence.
+        window_length (int): Length of the window for comparison (1 for amino acids,
+            3 for codons).
+
+    Returns:
+        float: The sequence similarity as a percentage.
+
+    Preconditions:
+        len(predicted) <= len(original).
+    """
+    if not truncate and len(original) != len(predicted):
+        raise ValueError(
+            "Set truncate to True if the length of sequences do not match."
+        )
+
+    identity = 0.0
+    original = original.strip()
+    predicted = predicted.strip()
+
+    if truncate:
+        original = original[: len(predicted)]
+
+    if window_length == 1:
+        # Simple comparison for amino acid
+        for i in range(len(predicted)):
+            if original[i] == predicted[i]:
+                identity += 1
+    else:
+        # Comparison for substrings based on window_length
+        for i in range(0, len(original) - window_length + 1, window_length):
+            if original[i : i + window_length] == predicted[i : i + window_length]:
+                identity += 1
+
+    return (identity / (len(predicted) / window_length)) * 100

model/CodonTransformer/CodonJupyter.py

@@ -0,0 +1,311 @@
+"""
+File: CodonJupyter.py
+---------------------
+Includes Jupyter-specific functions for displaying interactive widgets.
+"""
+
+from typing import Dict, List, Tuple
+
+import ipywidgets as widgets
+from IPython.display import HTML, display
+
+from CodonTransformer.CodonUtils import (
+    COMMON_ORGANISMS,
+    ID2ORGANISM,
+    ORGANISM2ID,
+    DNASequencePrediction,
+)
+
+
+class UserContainer:
+    """
+    A container class to store user inputs for organism and protein sequence.
+    Attributes:
+        organism (int): The selected organism id.
+        protein (str): The input protein sequence.
+    """
+
+    def __init__(self) -> None:
+        self.organism: int = -1
+        self.protein: str = ""
+
+
+def create_styled_options(
+    organisms: list, organism2id: Dict[str, int], is_fine_tuned: bool = False
+) -> list:
+    """
+    Create styled options for the dropdown widget.
+
+    Args:
+        organisms (list): List of organism names.
+        organism2id (Dict[str, int]): Dictionary mapping organism names to their IDs.
+        is_fine_tuned (bool): Whether these are fine-tuned organisms.
+
+    Returns:
+        list: Styled options for the dropdown widget.
+    """
+    styled_options = []
+    for organism in organisms:
+        organism_id = organism2id[organism]
+        if is_fine_tuned:
+            if organism_id < 10:
+                styled_options.append(f"\u200b{organism_id:>6}.  {organism}")
+            elif organism_id < 100:
+                styled_options.append(f"\u200b{organism_id:>5}.  {organism}")
+            else:
+                styled_options.append(f"\u200b{organism_id:>4}.  {organism}")
+        else:
+            if organism_id < 10:
+                styled_options.append(f"{organism_id:>6}.  {organism}")
+            elif organism_id < 100:
+                styled_options.append(f"{organism_id:>5}.  {organism}")
+            else:
+                styled_options.append(f"{organism_id:>4}.  {organism}")
+    return styled_options
+
+
+def create_dropdown_options(organism2id: Dict[str, int]) -> list:
+    """
+    Create the full list of dropdown options, including section headers.
+
+    Args:
+        organism2id (Dict[str, int]): Dictionary mapping organism names to their IDs.
+
+    Returns:
+        list: Full list of dropdown options.
+    """
+    fine_tuned_organisms = sorted(
+        [org for org in organism2id.keys() if org in COMMON_ORGANISMS]
+    )
+    all_organisms = sorted(organism2id.keys())
+
+    fine_tuned_options = create_styled_options(
+        fine_tuned_organisms, organism2id, is_fine_tuned=True
+    )
+    all_organisms_options = create_styled_options(
+        all_organisms, organism2id, is_fine_tuned=False
+    )
+
+    return (
+        [""]
+        + ["Selected Organisms"]
+        + fine_tuned_options
+        + [""]
+        + ["All Organisms"]
+        + all_organisms_options
+    )
+
+
+def create_organism_dropdown(container: UserContainer) -> widgets.Dropdown:
+    """
+    Create and configure the organism dropdown widget.
+
+    Args:
+        container (UserContainer): Container to store the selected organism.
+
+    Returns:
+        widgets.Dropdown: Configured dropdown widget.
+    """
+    dropdown = widgets.Dropdown(
+        options=create_dropdown_options(ORGANISM2ID),
+        description="",
+        layout=widgets.Layout(width="40%", margin="0 0 10px 0"),
+        style={"description_width": "initial"},
+    )
+
+    def show_organism(change: Dict[str, str]) -> None:
+        """
+        Update the container with the selected organism and print to terminal.
+
+        Args:
+            change (Dict[str, str]): Information about the change in dropdown value.
+        """
+        dropdown_choice = change["new"]
+        if dropdown_choice and dropdown_choice not in [
+            "Selected Organisms",
+            "All Organisms",
+        ]:
+            organism = "".join(filter(str.isdigit, dropdown_choice))
+            organism_id = ID2ORGANISM[int(organism)]
+            container.organism = organism_id
+        else:
+            container.organism = None
+
+    dropdown.observe(show_organism, names="value")
+    return dropdown
+
+
+def get_dropdown_style() -> str:
+    """
+    Return the custom CSS style for the dropdown widget.
+
+    Returns:
+        str: CSS style string.
+    """
+    return """
+    <style>
+        .widget-dropdown > select {
+            font-size: 16px;
+            font-weight: normal;
+            background-color: #f0f0f0;
+            border-radius: 5px;
+            padding: 5px;
+        }
+        .widget-label {
+            font-size: 18px;
+            font-weight: bold;
+        }
+        .custom-container {
+            display: flex;
+            flex-direction: column;
+            align-items: flex-start;
+        }
+        .widget-dropdown option[value^="\u200b"] {
+            font-family: sans-serif;
+            font-weight: bold;
+            font-size: 18px;
+            padding: 510px;
+        }
+        .widget-dropdown option[value*="Selected Organisms"],
+        .widget-dropdown option[value*="All Organisms"] {
+            text-align: center;
+            font-family: Arial, sans-serif;
+            font-weight: bold;
+            font-size: 20px;
+            color: #6900A1;
+            background-color: #00D8A1;
+        }
+    </style>
+    """
+
+
+def display_organism_dropdown(container: UserContainer) -> None:
+    """
+    Display the organism dropdown widget and apply custom styles.
+
+    Args:
+        container (UserContainer): Container to store the selected organism.
+    """
+    dropdown = create_organism_dropdown(container)
+    header = widgets.HTML(
+        '<b style="font-size:20px;">Select Organism:</b>'
+        '<div style="height:10px;"></div>'
+    )
+    container_widget = widgets.VBox(
+        [header, dropdown],
+        layout=widgets.Layout(padding="12px 0 12px 25px"),
+    )
+    display(container_widget)
+    display(HTML(get_dropdown_style()))
+
+
+def display_protein_input(container: UserContainer) -> None:
+    """
+    Display a widget for entering a protein sequence and save it to the container.
+
+    Args:
+        container (UserContainer): A container to store the entered protein sequence.
+    """
+    protein_input = widgets.Textarea(
+        value="",
+        placeholder="Enter here...",
+        description="",
+        layout=widgets.Layout(width="100%", height="100px", margin="0 0 10px 0"),
+        style={"description_width": "initial"},
+    )
+
+    # Custom CSS for the input widget
+    input_style = """
+        <style>
+            .widget-textarea > textarea {
+                font-size: 12px;
+                font-family: Arial, sans-serif;
+                font-weight: normal;
+                background-color: #f0f0f0;
+                border-radius: 5px;
+                padding: 10px;
+            }
+            .widget-label {
+                font-size: 18px;
+                font-weight: bold;
+            }
+            .custom-container {
+                display: flex;
+                flex-direction: column;
+                align-items: flex-start;
+            }
+        </style>
+    """
+
+    # Function to save the input protein sequence to the container
+    def save_protein(change: Dict[str, str]) -> None:
+        """
+        Save the input protein sequence to the container.
+
+        Args:
+            change (Dict[str, str]): A dictionary containing information about
+            the change in textarea value.
+        """
+        container.protein = (
+            change["new"]
+            .upper()
+            .strip()
+            .replace("\n", "")
+            .replace(" ", "")
+            .replace("\t", "")
+        )
+
+    # Attach the function to the input widget
+    protein_input.observe(save_protein, names="value")
+
+    # Display the input widget
+    header = widgets.HTML(
+        '<b style="font-size:20px;">Enter Protein Sequence:</b>'
+        '<div style="height:18px;"></div>'
+    )
+    container_widget = widgets.VBox(
+        [header, protein_input], layout=widgets.Layout(padding="12px 12px 0 25px")
+    )
+
+    display(container_widget)
+    display(widgets.HTML(input_style))
+
+
+def format_model_output(output: DNASequencePrediction) -> str:
+    """
+    Format DNA sequence prediction output in an appealing and easy-to-read manner.
+
+    This function takes the prediction output and formats it into
+    a structured string with clear section headers and separators.
+
+    Args:
+        output (DNASequencePrediction): Object containing the prediction output.
+            Expected attributes:
+            - organism (str): The organism name.
+            - protein (str): The input protein sequence.
+            - processed_input (str): The processed input sequence.
+            - predicted_dna (str): The predicted DNA sequence.
+
+    Returns:
+        str: A formatted string containing the organized output.
+    """
+
+    def format_section(title: str, content: str) -> str:
+        """Helper function to format individual sections."""
+        separator = "-" * 29
+        title_line = f"| {title.center(25)} |"
+        return f"{separator}\n{title_line}\n{separator}\n{content}\n\n"
+
+    sections: List[Tuple[str, str]] = [
+        ("Organism", output.organism),
+        ("Input Protein", output.protein),
+        ("Processed Input", output.processed_input),
+        ("Predicted DNA", output.predicted_dna),
+    ]
+
+    formatted_output = ""
+    for title, content in sections:
+        formatted_output += format_section(title, content)
+
+    # Remove the last newline to avoid extra space at the end
+    return formatted_output.rstrip()

model/CodonTransformer/CodonPrediction.py

@@ -0,0 +1,855 @@
+"""
+File: CodonPrediction.py
+---------------------------
+Includes functions to tokenize input, load models, infer predicted dna sequences and
+helper functions related to processing data for passing to the model.
+"""
+
+import warnings
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import onnxruntime as rt
+import torch
+import transformers
+from transformers import (
+    AutoTokenizer,
+    BatchEncoding,
+    BigBirdConfig,
+    BigBirdForMaskedLM,
+    PreTrainedTokenizerFast,
+)
+
+from CodonTransformer.CodonData import get_merged_seq
+from CodonTransformer.CodonUtils import (
+    AMINO_ACID_TO_INDEX,
+    INDEX2TOKEN,
+    NUM_ORGANISMS,
+    ORGANISM2ID,
+    TOKEN2INDEX,
+    DNASequencePrediction,
+)
+
+
+def predict_dna_sequence(
+    protein: str,
+    organism: Union[int, str],
+    device: torch.device,
+    tokenizer: Union[str, PreTrainedTokenizerFast] = None,
+    model: Union[str, torch.nn.Module] = None,
+    attention_type: str = "original_full",
+    deterministic: bool = True,
+    temperature: float = 0.2,
+    top_p: float = 0.95,
+    num_sequences: int = 1,
+    match_protein: bool = False,
+) -> Union[DNASequencePrediction, List[DNASequencePrediction]]:
+    """
+    Predict the DNA sequence(s) for a given protein using the CodonTransformer model.
+
+    This function takes a protein sequence and an organism (as ID or name) as input
+    and returns the predicted DNA sequence(s) using the CodonTransformer model. It can use
+    either provided tokenizer and model objects or load them from specified paths.
+
+    Args:
+        protein (str): The input protein sequence for which to predict the DNA sequence.
+        organism (Union[int, str]): Either the ID of the organism or its name (e.g.,
+            "Escherichia coli general"). If a string is provided, it will be converted
+            to the corresponding ID using ORGANISM2ID.
+        device (torch.device): The device (CPU or GPU) to run the model on.
+        tokenizer (Union[str, PreTrainedTokenizerFast, None], optional): Either a file
+            path to load the tokenizer from, a pre-loaded tokenizer object, or None. If
+            None, it will be loaded from HuggingFace. Defaults to None.
+        model (Union[str, torch.nn.Module, None], optional): Either a file path to load
+            the model from, a pre-loaded model object, or None. If None, it will be
+            loaded from HuggingFace. Defaults to None.
+        attention_type (str, optional): The type of attention mechanism to use in the
+            model. Can be either 'block_sparse' or 'original_full'. Defaults to
+            "original_full".
+        deterministic (bool, optional): Whether to use deterministic decoding (most
+            likely tokens). If False, samples tokens according to their probabilities
+            adjusted by the temperature. Defaults to True.
+        temperature (float, optional): A value controlling the randomness of predictions
+            during non-deterministic decoding. Lower values (e.g., 0.2) make the model
+            more conservative, while higher values (e.g., 0.8) increase randomness.
+            Using high temperatures may result in prediction of DNA sequences that
+            do not translate to the input protein.
+            Recommended values are:
+                - Low randomness: 0.2
+                - Medium randomness: 0.5
+                - High randomness: 0.8
+            The temperature must be a positive float. Defaults to 0.2.
+        top_p (float, optional): The cumulative probability threshold for nucleus sampling.
+            Tokens with cumulative probability up to top_p are considered for sampling.
+            This parameter helps balance diversity and coherence in the predicted DNA sequences.
+            The value must be a float between 0 and 1. Defaults to 0.95.
+        num_sequences (int, optional): The number of DNA sequences to generate. Only applicable
+            when deterministic is False. Defaults to 1.
+        match_protein (bool, optional): Ensures the predicted DNA sequence is translated
+            to the input protein sequence by sampling from only the respective codons of
+            given amino acids. Defaults to False.
+
+    Returns:
+        Union[DNASequencePrediction, List[DNASequencePrediction]]: An object or list of objects
+        containing the prediction results:
+            - organism (str): Name of the organism used for prediction.
+            - protein (str): Input protein sequence for which DNA sequence is predicted.
+            - processed_input (str): Processed input sequence (merged protein and DNA).
+            - predicted_dna (str): Predicted DNA sequence.
+
+    Raises:
+        ValueError: If the protein sequence is empty, if the organism is invalid,
+            if the temperature is not a positive float, if top_p is not between 0 and 1,
+            or if num_sequences is less than 1 or used with deterministic mode.
+
+    Note:
+        This function uses ORGANISM2ID, INDEX2TOKEN, and AMINO_ACID_TO_INDEX dictionaries
+        imported from CodonTransformer.CodonUtils. ORGANISM2ID maps organism names to their
+        corresponding IDs. INDEX2TOKEN maps model output indices (token IDs) to
+        respective codons. AMINO_ACID_TO_INDEX maps each amino acid and stop symbol to indices
+        of codon tokens that translate to it.
+
+    Example:
+        >>> import torch
+        >>> from transformers import AutoTokenizer, BigBirdForMaskedLM
+        >>> from CodonTransformer.CodonPrediction import predict_dna_sequence
+        >>> from CodonTransformer.CodonJupyter import format_model_output
+        >>>
+        >>> # Set up device
+        >>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        >>>
+        >>> # Load tokenizer and model
+        >>> tokenizer = AutoTokenizer.from_pretrained("adibvafa/CodonTransformer")
+        >>> model = BigBirdForMaskedLM.from_pretrained("adibvafa/CodonTransformer")
+        >>> model = model.to(device)
+        >>>
+        >>> # Define protein sequence and organism
+        >>> protein = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLA"
+        >>> organism = "Escherichia coli general"
+        >>>
+        >>> # Predict DNA sequence with deterministic decoding (single sequence)
+        >>> output = predict_dna_sequence(
+        ...     protein=protein,
+        ...     organism=organism,
+        ...     device=device,
+        ...     tokenizer=tokenizer,
+        ...     model=model,
+        ...     attention_type="original_full",
+        ...     deterministic=True
+        ... )
+        >>>
+        >>> # Predict multiple DNA sequences with low randomness and top_p sampling
+        >>> output_random = predict_dna_sequence(
+        ...     protein=protein,
+        ...     organism=organism,
+        ...     device=device,
+        ...     tokenizer=tokenizer,
+        ...     model=model,
+        ...     attention_type="original_full",
+        ...     deterministic=False,
+        ...     temperature=0.2,
+        ...     top_p=0.95,
+        ...     num_sequences=3
+        ... )
+        >>>
+        >>> print(format_model_output(output))
+        >>> for i, seq in enumerate(output_random, 1):
+        ...     print(f"Sequence {i}:")
+        ...     print(format_model_output(seq))
+        ...     print()
+    """
+    if not protein:
+        raise ValueError("Protein sequence cannot be empty.")
+
+    if not isinstance(temperature, (float, int)) or temperature <= 0:
+        raise ValueError("Temperature must be a positive float.")
+
+    if not isinstance(top_p, (float, int)) or not 0 < top_p <= 1.0:
+        raise ValueError("top_p must be a float between 0 and 1.")
+
+    if not isinstance(num_sequences, int) or num_sequences < 1:
+        raise ValueError("num_sequences must be a positive integer.")
+
+    if deterministic and num_sequences > 1:
+        raise ValueError(
+            "Multiple sequences can only be generated in non-deterministic mode."
+        )
+
+    # Load tokenizer
+    if not isinstance(tokenizer, PreTrainedTokenizerFast):
+        tokenizer = load_tokenizer(tokenizer)
+
+    # Load model
+    if not isinstance(model, torch.nn.Module):
+        model = load_model(model, device=device, attention_type=attention_type)
+    else:
+        model.eval()
+        model.bert.set_attention_type(attention_type)
+        model.to(device)
+
+    # Validate organism and convert to organism_id and organism_name
+    organism_id, organism_name = validate_and_convert_organism(organism)
+
+    # Inference loop
+    with torch.no_grad():
+        # Tokenize the input sequence
+        merged_seq = get_merged_seq(protein=protein, dna="")
+        input_dict = {
+            "idx": 0,  # sample index
+            "codons": merged_seq,
+            "organism": organism_id,
+        }
+        tokenized_input = tokenize([input_dict], tokenizer=tokenizer).to(device)
+
+        # Get the model predictions
+        output_dict = model(**tokenized_input, return_dict=True)
+        logits = output_dict.logits.detach().cpu()
+        logits = logits[:, 1:-1, :]  # Remove [CLS] and [SEP] tokens
+
+        # Mask the logits of codons that do not correspond to the input protein sequence
+        if match_protein:
+            possible_tokens_per_position = [
+                AMINO_ACID_TO_INDEX[token[0]] for token in merged_seq.split(" ")
+            ]
+            mask = torch.full_like(logits, float("-inf"))
+
+            for pos, possible_tokens in enumerate(possible_tokens_per_position):
+                mask[:, pos, possible_tokens] = 0
+
+            logits = mask + logits
+
+        predictions = []
+        for _ in range(num_sequences):
+            # Decode the predicted DNA sequence from the model output
+            if deterministic:
+                predicted_indices = logits.argmax(dim=-1).squeeze().tolist()
+            else:
+                predicted_indices = sample_non_deterministic(
+                    logits=logits, temperature=temperature, top_p=top_p
+                )
+
+            predicted_dna = list(map(INDEX2TOKEN.__getitem__, predicted_indices))
+            predicted_dna = (
+                "".join([token[-3:] for token in predicted_dna]).strip().upper()
+            )
+
+            predictions.append(
+                DNASequencePrediction(
+                    organism=organism_name,
+                    protein=protein,
+                    processed_input=merged_seq,
+                    predicted_dna=predicted_dna,
+                )
+            )
+
+    return predictions[0] if num_sequences == 1 else predictions
+
+
+def sample_non_deterministic(
+    logits: torch.Tensor,
+    temperature: float = 0.2,
+    top_p: float = 0.95,
+) -> List[int]:
+    """
+    Sample token indices from logits using temperature scaling and nucleus (top-p) sampling.
+
+    This function applies temperature scaling to the logits, computes probabilities,
+    and then performs nucleus sampling to select token indices. It is used for
+    non-deterministic decoding in language models to introduce randomness while
+    maintaining coherence in the generated sequences.
+
+    Args:
+        logits (torch.Tensor): The logits output from the model of shape
+            [seq_len, vocab_size] or [batch_size, seq_len, vocab_size].
+        temperature (float, optional): Temperature value for scaling logits.
+            Must be a positive float. Defaults to 1.0.
+        top_p (float, optional): Cumulative probability threshold for nucleus sampling.
+            Must be a float between 0 and 1. Tokens with cumulative probability up to
+            `top_p` are considered for sampling. Defaults to 0.95.
+
+    Returns:
+        List[int]: A list of sampled token indices corresponding to the predicted tokens.
+
+    Raises:
+        ValueError: If `temperature` is not a positive float or if `top_p` is not between 0 and 1.
+
+    Example:
+        >>> logits = model_output.logits  # Assume logits is a tensor of shape [seq_len, vocab_size]
+        >>> predicted_indices = sample_non_deterministic(logits, temperature=0.7, top_p=0.9)
+    """
+    if not isinstance(temperature, (float, int)) or temperature <= 0:
+        raise ValueError("Temperature must be a positive float.")
+
+    if not isinstance(top_p, (float, int)) or not 0 < top_p <= 1.0:
+        raise ValueError("top_p must be a float between 0 and 1.")
+
+    # Compute probabilities using temperature scaling
+    probs = torch.softmax(logits / temperature, dim=-1)
+
+
+    # Remove batch dimension if present
+    if probs.dim() == 3:
+        probs = probs.squeeze(0)  # Shape: [seq_len, vocab_size]
+
+    # Sort probabilities in descending order
+    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
+    probs_sum = torch.cumsum(probs_sort, dim=-1)
+    mask = probs_sum - probs_sort > top_p
+
+    # Zero out probabilities for tokens beyond the top-p threshold
+    probs_sort[mask] = 0.0
+
+    # Renormalize the probabilities
+    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
+    next_token = torch.multinomial(probs_sort, num_samples=1)
+    predicted_indices = torch.gather(probs_idx, -1, next_token).squeeze(-1)
+
+    return predicted_indices.tolist()
+
+
+def load_model(
+    model_path: Optional[str] = None,
+    device: torch.device = None,
+    attention_type: str = "original_full",
+    num_organisms: int = None,
+    remove_prefix: bool = True,
+) -> torch.nn.Module:
+    """
+    Load a BigBirdForMaskedLM model from a model file, checkpoint, or HuggingFace.
+
+    Args:
+        model_path (Optional[str]): Path to the model file or checkpoint. If None,
+            load from HuggingFace.
+        device (torch.device, optional): The device to load the model onto.
+        attention_type (str, optional): The type of attention, 'block_sparse'
+            or 'original_full'.
+        num_organisms (int, optional): Number of organisms, needed if loading from a
+            checkpoint that requires this.
+        remove_prefix (bool, optional): Whether to remove the "model." prefix from the
+            keys in the state dict.
+
+    Returns:
+        torch.nn.Module: The loaded model.
+    """
+    if not model_path:
+        warnings.warn("Model path not provided. Loading from HuggingFace.", UserWarning)
+        model = BigBirdForMaskedLM.from_pretrained("adibvafa/CodonTransformer")
+
+    elif model_path.endswith(".ckpt"):
+        checkpoint = torch.load(model_path)
+        state_dict = checkpoint["state_dict"]
+
+        # Remove the "model." prefix from the keys
+        if remove_prefix:
+            state_dict = {
+                key.replace("model.", ""): value for key, value in state_dict.items()
+            }
+
+        if num_organisms is None:
+            num_organisms = NUM_ORGANISMS
+
+        # Load model configuration and instantiate the model
+        config = load_bigbird_config(num_organisms)
+        model = BigBirdForMaskedLM(config=config)
+        model.load_state_dict(state_dict)
+
+    elif model_path.endswith(".pt"):
+        state_dict = torch.load(model_path)
+        config = state_dict.pop("self.config")
+        model = BigBirdForMaskedLM(config=config)
+        model.load_state_dict(state_dict)
+
+    else:
+        raise ValueError(
+            "Unsupported file type. Please provide a .ckpt or .pt file, "
+            "or None to load from HuggingFace."
+        )
+
+    # Prepare model for evaluation
+    model.bert.set_attention_type(attention_type)
+    model.eval()
+    if device:
+        model.to(device)
+
+    return model
+
+
+def load_bigbird_config(num_organisms: int) -> BigBirdConfig:
+    """
+    Load the config object used to train the BigBird transformer.
+
+    Args:
+        num_organisms (int): The number of organisms.
+
+    Returns:
+        BigBirdConfig: The configuration object for BigBird.
+    """
+    config = transformers.BigBirdConfig(
+        vocab_size=len(TOKEN2INDEX),  # Equal to len(tokenizer)
+        type_vocab_size=num_organisms,
+        sep_token_id=2,
+    )
+    return config
+
+
+def create_model_from_checkpoint(
+    checkpoint_dir: str, output_model_dir: str, num_organisms: int
+) -> None:
+    """
+    Save a model to disk using a previous checkpoint.
+
+    Args:
+        checkpoint_dir (str): Directory where the checkpoint is stored.
+        output_model_dir (str): Directory where the model will be saved.
+        num_organisms (int): Number of organisms.
+    """
+    checkpoint = load_model(model_path=checkpoint_dir, num_organisms=num_organisms)
+    state_dict = checkpoint.state_dict()
+    state_dict["self.config"] = load_bigbird_config(num_organisms=num_organisms)
+
+    # Save the model state dict to the output directory
+    torch.save(state_dict, output_model_dir)
+
+
+def load_tokenizer(tokenizer_path: Optional[str] = None) -> PreTrainedTokenizerFast:
+    """
+    Create and return a tokenizer object from tokenizer path or HuggingFace.
+
+    Args:
+        tokenizer_path (Optional[str]): Path to the tokenizer file. If None,
+        load from HuggingFace.
+
+    Returns:
+        PreTrainedTokenizerFast: The tokenizer object.
+    """
+    if not tokenizer_path:
+        warnings.warn(
+            "Tokenizer path not provided. Loading from HuggingFace.", UserWarning
+        )
+        return AutoTokenizer.from_pretrained("adibvafa/CodonTransformer")
+
+    return transformers.PreTrainedTokenizerFast(
+        tokenizer_file=tokenizer_path,
+        bos_token="[CLS]",
+        eos_token="[SEP]",
+        unk_token="[UNK]",
+        sep_token="[SEP]",
+        pad_token="[PAD]",
+        cls_token="[CLS]",
+        mask_token="[MASK]",
+    )
+
+
+def tokenize(
+    batch: List[Dict[str, Any]],
+    tokenizer: Union[PreTrainedTokenizerFast, str] = None,
+    max_len: int = 2048,
+) -> BatchEncoding:
+    """
+    Return the tokenized sequences given a batch of input data.
+    Each data in the batch is expected to be a dictionary with "codons" and
+    "organism" keys.
+
+    Args:
+        batch (List[Dict[str, Any]]): A list of dictionaries with "codons" and
+            "organism" keys.
+        tokenizer (PreTrainedTokenizerFast, str, optional): The tokenizer object or
+            path to the tokenizer file.
+        max_len (int, optional): Maximum length of the tokenized sequence.
+
+    Returns:
+        BatchEncoding: The tokenized batch.
+    """
+    if not isinstance(tokenizer, PreTrainedTokenizerFast):
+        tokenizer = load_tokenizer(tokenizer)
+
+    tokenized = tokenizer(
+        [data["codons"] for data in batch],
+        return_attention_mask=True,
+        return_token_type_ids=True,
+        truncation=True,
+        padding=True,
+        max_length=max_len,
+        return_tensors="pt",
+    )
+
+    # Add token type IDs for species
+    seq_len = tokenized["input_ids"].shape[-1]
+    species_index = torch.tensor([[data["organism"]] for data in batch])
+    tokenized["token_type_ids"] = species_index.repeat(1, seq_len)
+
+    return tokenized
+
+
+def validate_and_convert_organism(organism: Union[int, str]) -> Tuple[int, str]:
+    """
+    Validate and convert the organism input to both ID and name.
+
+    This function takes either an organism ID or name as input and returns both
+    the ID and name. It performs validation to ensure the input corresponds to
+    a valid organism in the ORGANISM2ID dictionary.
+
+    Args:
+        organism (Union[int, str]): Either the ID of the organism (int) or its
+        name (str).
+
+    Returns:
+        Tuple[int, str]: A tuple containing the organism ID (int) and name (str).
+
+    Raises:
+        ValueError: If the input is neither a string nor an integer, if the
+        organism name is not found in ORGANISM2ID, if the organism ID is not a
+        value in ORGANISM2ID, or if no name is found for a given ID.
+
+    Note:
+        This function relies on the ORGANISM2ID dictionary imported from
+        CodonTransformer.CodonUtils, which maps organism names to their
+        corresponding IDs.
+    """
+    if isinstance(organism, str):
+        if organism not in ORGANISM2ID:
+            raise ValueError(
+                f"Invalid organism name: {organism}. "
+                "Please use a valid organism name or ID."
+            )
+        organism_id = ORGANISM2ID[organism]
+        organism_name = organism
+
+    elif isinstance(organism, int):
+        if organism not in ORGANISM2ID.values():
+            raise ValueError(
+                f"Invalid organism ID: {organism}. "
+                "Please use a valid organism name or ID."
+            )
+
+        organism_id = organism
+        organism_name = next(
+            (name for name, id in ORGANISM2ID.items() if id == organism), None
+        )
+        if organism_name is None:
+            raise ValueError(f"No organism name found for ID: {organism}")
+
+    return organism_id, organism_name
+
+
+def get_high_frequency_choice_sequence(
+    protein: str, codon_frequencies: Dict[str, Tuple[List[str], List[float]]]
+) -> str:
+    """
+    Return the DNA sequence optimized using High Frequency Choice (HFC) approach
+    in which the most frequent codon for a given amino acid is always chosen.
+
+    Args:
+        protein (str): The protein sequence.
+        codon_frequencies (Dict[str, Tuple[List[str], List[float]]]): Codon
+        frequencies for each amino acid.
+
+    Returns:
+        str: The optimized DNA sequence.
+    """
+    # Select the most frequent codon for each amino acid in the protein sequence
+    dna_codons = [
+        codon_frequencies[aminoacid][0][np.argmax(codon_frequencies[aminoacid][1])]
+        for aminoacid in protein
+    ]
+    return "".join(dna_codons)
+
+
+def precompute_most_frequent_codons(
+    codon_frequencies: Dict[str, Tuple[List[str], List[float]]],
+) -> Dict[str, str]:
+    """
+    Precompute the most frequent codon for each amino acid.
+
+    Args:
+        codon_frequencies (Dict[str, Tuple[List[str], List[float]]]): Codon
+        frequencies for each amino acid.
+
+    Returns:
+        Dict[str, str]: The most frequent codon for each amino acid.
+    """
+    # Create a dictionary mapping each amino acid to its most frequent codon
+    return {
+        aminoacid: codons[np.argmax(frequencies)]
+        for aminoacid, (codons, frequencies) in codon_frequencies.items()
+    }
+
+
+def get_high_frequency_choice_sequence_optimized(
+    protein: str, codon_frequencies: Dict[str, Tuple[List[str], List[float]]]
+) -> str:
+    """
+    Efficient implementation of get_high_frequency_choice_sequence that uses
+    vectorized operations and helper functions, achieving up to x10 faster speed.
+
+    Args:
+        protein (str): The protein sequence.
+        codon_frequencies (Dict[str, Tuple[List[str], List[float]]]): Codon
+        frequencies for each amino acid.
+
+    Returns:
+        str: The optimized DNA sequence.
+    """
+    # Precompute the most frequent codons for each amino acid
+    most_frequent_codons = precompute_most_frequent_codons(codon_frequencies)
+
+    return "".join(most_frequent_codons[aminoacid] for aminoacid in protein)
+
+
+def get_background_frequency_choice_sequence(
+    protein: str, codon_frequencies: Dict[str, Tuple[List[str], List[float]]]
+) -> str:
+    """
+    Return the DNA sequence optimized using Background Frequency Choice (BFC)
+    approach in which a random codon for a given amino acid is chosen using
+    the codon frequencies probability distribution.
+
+    Args:
+        protein (str): The protein sequence.
+        codon_frequencies (Dict[str, Tuple[List[str], List[float]]]): Codon
+        frequencies for each amino acid.
+
+    Returns:
+        str: The optimized DNA sequence.
+    """
+    # Select a random codon for each amino acid based on the codon frequencies
+    # probability distribution
+    dna_codons = [
+        np.random.choice(
+            codon_frequencies[aminoacid][0], p=codon_frequencies[aminoacid][1]
+        )
+        for aminoacid in protein
+    ]
+    return "".join(dna_codons)
+
+
+def precompute_cdf(
+    codon_frequencies: Dict[str, Tuple[List[str], List[float]]],
+) -> Dict[str, Tuple[List[str], Any]]:
+    """
+    Precompute the cumulative distribution function (CDF) for each amino acid.
+
+    Args:
+        codon_frequencies (Dict[str, Tuple[List[str], List[float]]]): Codon
+        frequencies for each amino acid.
+
+    Returns:
+        Dict[str, Tuple[List[str], Any]]: CDFs for each amino acid.
+    """
+    cdf = {}
+
+    # Calculate the cumulative distribution function for each amino acid
+    for aminoacid, (codons, frequencies) in codon_frequencies.items():
+        cdf[aminoacid] = (codons, np.cumsum(frequencies))
+
+    return cdf
+
+
+def get_background_frequency_choice_sequence_optimized(
+    protein: str, codon_frequencies: Dict[str, Tuple[List[str], List[float]]]
+) -> str:
+    """
+    Efficient implementation of get_background_frequency_choice_sequence that uses
+    vectorized operations and helper functions, achieving up to x8 faster speed.
+
+    Args:
+        protein (str): The protein sequence.
+        codon_frequencies (Dict[str, Tuple[List[str], List[float]]]): Codon
+        frequencies for each amino acid.
+
+    Returns:
+        str: The optimized DNA sequence.
+    """
+    dna_codons = []
+    cdf = precompute_cdf(codon_frequencies)
+
+    # Select a random codon for each amino acid using the precomputed CDFs
+    for aminoacid in protein:
+        codons, cumulative_prob = cdf[aminoacid]
+        selected_codon_index = np.searchsorted(cumulative_prob, np.random.rand())
+        dna_codons.append(codons[selected_codon_index])
+
+    return "".join(dna_codons)
+
+
+def get_uniform_random_choice_sequence(
+    protein: str, codon_frequencies: Dict[str, Tuple[List[str], List[float]]]
+) -> str:
+    """
+    Return the DNA sequence optimized using Uniform Random Choice (URC) approach
+    in which a random codon for a given amino acid is chosen using a uniform
+    prior.
+
+    Args:
+        protein (str): The protein sequence.
+        codon_frequencies (Dict[str, Tuple[List[str], List[float]]]): Codon
+        frequencies for each amino acid.
+
+    Returns:
+        str: The optimized DNA sequence.
+    """
+    # Select a random codon for each amino acid using a uniform prior distribution
+    dna_codons = []
+    for aminoacid in protein:
+        codons = codon_frequencies[aminoacid][0]
+        random_index = np.random.randint(0, len(codons))
+        dna_codons.append(codons[random_index])
+    return "".join(dna_codons)
+
+
+def get_icor_prediction(input_seq: str, model_path: str, stop_symbol: str) -> str:
+    """
+    Return the optimized codon sequence for the given protein sequence using ICOR.
+
+    Credit: ICOR: improving codon optimization with recurrent neural networks
+            Rishab Jain, Aditya Jain, Elizabeth Mauro, Kevin LeShane, Douglas
+            Densmore
+
+    Args:
+        input_seq (str): The input protein sequence.
+        model_path (str): The path to the ICOR model.
+        stop_symbol (str): The symbol representing stop codons in the sequence.
+
+    Returns:
+        str: The optimized DNA sequence.
+    """
+    input_seq = input_seq.strip().upper()
+    input_seq = input_seq.replace(stop_symbol, "*")
+
+    # Define categorical labels from when model was trained.
+    labels = [
+        "AAA",
+        "AAC",
+        "AAG",
+        "AAT",
+        "ACA",
+        "ACG",
+        "ACT",
+        "AGC",
+        "ATA",
+        "ATC",
+        "ATG",
+        "ATT",
+        "CAA",
+        "CAC",
+        "CAG",
+        "CCG",
+        "CCT",
+        "CTA",
+        "CTC",
+        "CTG",
+        "CTT",
+        "GAA",
+        "GAT",
+        "GCA",
+        "GCC",
+        "GCG",
+        "GCT",
+        "GGA",
+        "GGC",
+        "GTC",
+        "GTG",
+        "GTT",
+        "TAA",
+        "TAT",
+        "TCA",
+        "TCG",
+        "TCT",
+        "TGG",
+        "TGT",
+        "TTA",
+        "TTC",
+        "TTG",
+        "TTT",
+        "ACC",
+        "CAT",
+        "CCA",
+        "CGG",
+        "CGT",
+        "GAC",
+        "GAG",
+        "GGT",
+        "AGT",
+        "GGG",
+        "GTA",
+        "TGC",
+        "CCC",
+        "CGA",
+        "CGC",
+        "TAC",
+        "TAG",
+        "TCC",
+        "AGA",
+        "AGG",
+        "TGA",
+    ]
+
+    # Define aa to integer table
+    def aa2int(seq: str) -> List[int]:
+        _aa2int = {
+            "A": 1,
+            "R": 2,
+            "N": 3,
+            "D": 4,
+            "C": 5,
+            "Q": 6,
+            "E": 7,
+            "G": 8,
+            "H": 9,
+            "I": 10,
+            "L": 11,
+            "K": 12,
+            "M": 13,
+            "F": 14,
+            "P": 15,
+            "S": 16,
+            "T": 17,
+            "W": 18,
+            "Y": 19,
+            "V": 20,
+            "B": 21,
+            "Z": 22,
+            "X": 23,
+            "*": 24,
+            "-": 25,
+            "?": 26,
+        }
+        return [_aa2int[i] for i in seq]
+
+    # Create empty array to fill
+    oh_array = np.zeros(shape=(26, len(input_seq)))
+
+    # Load placements from aa2int
+    aa_placement = aa2int(input_seq)
+
+    # One-hot encode the amino acid sequence:
+
+    # style nit: more pythonic to write for i in range(0, len(aa_placement)):
+    for i in range(0, len(aa_placement)):
+        oh_array[aa_placement[i], i] = 1
+        i += 1
+
+    oh_array = [oh_array]
+    x = np.array(np.transpose(oh_array))
+
+    y = x.astype(np.float32)
+
+    y = np.reshape(y, (y.shape[0], 1, 26))
+
+    # Start ICOR session using model.
+    sess = rt.InferenceSession(model_path)
+    input_name = sess.get_inputs()[0].name
+
+    # Get prediction:
+    pred_onx = sess.run(None, {input_name: y})
+
+    # Get the index of the highest probability from softmax output:
+    pred_indices = []
+    for pred in pred_onx[0]:
+        pred_indices.append(np.argmax(pred))
+
+    out_str = ""
+    for index in pred_indices:
+        out_str += labels[index]
+
+    return out_str

model/CodonTransformer/CodonUtils.py

@@ -0,0 +1,830 @@
+"""
+File: CodonUtils.py
+---------------------
+Includes constants and helper functions used by other Python scripts.
+"""
+
+import itertools
+import gzip
+import json
+import os
+import pickle
+import re
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Dict, Iterator, List, Optional, Tuple
+
+import pandas as pd
+import requests
+import torch
+
+# List of all amino acids
+AMINO_ACIDS: List[str] = [
+    "A",  # Alanine
+    "C",  # Cysteine
+    "D",  # Aspartic acid
+    "E",  # Glutamic acid
+    "F",  # Phenylalanine
+    "G",  # Glycine
+    "H",  # Histidine
+    "I",  # Isoleucine
+    "K",  # Lysine
+    "L",  # Leucine
+    "M",  # Methionine
+    "N",  # Asparagine
+    "P",  # Proline
+    "Q",  # Glutamine
+    "R",  # Arginine
+    "S",  # Serine
+    "T",  # Threonine
+    "V",  # Valine
+    "W",  # Tryptophan
+    "Y",  # Tyrosine
+]
+STOP_SYMBOLS = ["_", "*"]  # Stop codon symbols
+
+# Dictionary ambiguous amino acids to standard amino acids
+AMBIGUOUS_AMINOACID_MAP: Dict[str, list[str]] = {
+    "B": ["N", "D"],  # Asparagine (N) or Aspartic acid (D)
+    "Z": ["Q", "E"],  # Glutamine (Q) or Glutamic acid (E)
+    "X": ["A"],  # Any amino acid (typically replaced with Alanine)
+    "J": ["L", "I"],  # Leucine (L) or Isoleucine (I)
+    "U": ["C"],  # Selenocysteine (typically replaced with Cysteine)
+    "O": ["K"],  # Pyrrolysine (typically replaced with Lysine)
+}
+
+# List of all possible start and stop codons
+START_CODONS: List[str] = ["ATG", "TTG", "CTG", "GTG"]
+STOP_CODONS: List[str] = ["TAA", "TAG", "TGA"]
+
+# Token-to-index mapping for amino acids and special tokens
+TOKEN2INDEX: Dict[str, int] = {
+    "[UNK]": 0,
+    "[CLS]": 1,
+    "[SEP]": 2,
+    "[PAD]": 3,
+    "[MASK]": 4,
+    "a_unk": 5,
+    "c_unk": 6,
+    "d_unk": 7,
+    "e_unk": 8,
+    "f_unk": 9,
+    "g_unk": 10,
+    "h_unk": 11,
+    "i_unk": 12,
+    "k_unk": 13,
+    "l_unk": 14,
+    "m_unk": 15,
+    "n_unk": 16,
+    "p_unk": 17,
+    "q_unk": 18,
+    "r_unk": 19,
+    "s_unk": 20,
+    "t_unk": 21,
+    "v_unk": 22,
+    "w_unk": 23,
+    "y_unk": 24,
+    "__unk": 25,
+    "k_aaa": 26,
+    "n_aac": 27,
+    "k_aag": 28,
+    "n_aat": 29,
+    "t_aca": 30,
+    "t_acc": 31,
+    "t_acg": 32,
+    "t_act": 33,
+    "r_aga": 34,
+    "s_agc": 35,
+    "r_agg": 36,
+    "s_agt": 37,
+    "i_ata": 38,
+    "i_atc": 39,
+    "m_atg": 40,
+    "i_att": 41,
+    "q_caa": 42,
+    "h_cac": 43,
+    "q_cag": 44,
+    "h_cat": 45,
+    "p_cca": 46,
+    "p_ccc": 47,
+    "p_ccg": 48,
+    "p_cct": 49,
+    "r_cga": 50,
+    "r_cgc": 51,
+    "r_cgg": 52,
+    "r_cgt": 53,
+    "l_cta": 54,
+    "l_ctc": 55,
+    "l_ctg": 56,
+    "l_ctt": 57,
+    "e_gaa": 58,
+    "d_gac": 59,
+    "e_gag": 60,
+    "d_gat": 61,
+    "a_gca": 62,
+    "a_gcc": 63,
+    "a_gcg": 64,
+    "a_gct": 65,
+    "g_gga": 66,
+    "g_ggc": 67,
+    "g_ggg": 68,
+    "g_ggt": 69,
+    "v_gta": 70,
+    "v_gtc": 71,
+    "v_gtg": 72,
+    "v_gtt": 73,
+    "__taa": 74,
+    "y_tac": 75,
+    "__tag": 76,
+    "y_tat": 77,
+    "s_tca": 78,
+    "s_tcc": 79,
+    "s_tcg": 80,
+    "s_tct": 81,
+    "__tga": 82,
+    "c_tgc": 83,
+    "w_tgg": 84,
+    "c_tgt": 85,
+    "l_tta": 86,
+    "f_ttc": 87,
+    "l_ttg": 88,
+    "f_ttt": 89,
+}
+
+# Index-to-token mapping, reverse of TOKEN2INDEX
+INDEX2TOKEN: Dict[int, str] = {i: c for c, i in TOKEN2INDEX.items()}
+
+# Dictionary mapping each amino acid and stop symbol to indices of codon tokens that translate to it
+AMINO_ACID_TO_INDEX = {
+    aa: sorted(
+        [i for t, i in TOKEN2INDEX.items() if t[0].upper() == aa and t[-3:] != "unk"]
+    )
+    for aa in (AMINO_ACIDS + STOP_SYMBOLS)
+}
+
+
+# Mask token mapping
+TOKEN2MASK: Dict[int, int] = {
+    0: 0,
+    1: 1,
+    2: 2,
+    3: 3,
+    4: 4,
+    5: 5,
+    6: 6,
+    7: 7,
+    8: 8,
+    9: 9,
+    10: 10,
+    11: 11,
+    12: 12,
+    13: 13,
+    14: 14,
+    15: 15,
+    16: 16,
+    17: 17,
+    18: 18,
+    19: 19,
+    20: 20,
+    21: 21,
+    22: 22,
+    23: 23,
+    24: 24,
+    25: 25,
+    26: 13,
+    27: 16,
+    28: 13,
+    29: 16,
+    30: 21,
+    31: 21,
+    32: 21,
+    33: 21,
+    34: 19,
+    35: 20,
+    36: 19,
+    37: 20,
+    38: 12,
+    39: 12,
+    40: 15,
+    41: 12,
+    42: 18,
+    43: 11,
+    44: 18,
+    45: 11,
+    46: 17,
+    47: 17,
+    48: 17,
+    49: 17,
+    50: 19,
+    51: 19,
+    52: 19,
+    53: 19,
+    54: 14,
+    55: 14,
+    56: 14,
+    57: 14,
+    58: 8,
+    59: 7,
+    60: 8,
+    61: 7,
+    62: 5,
+    63: 5,
+    64: 5,
+    65: 5,
+    66: 10,
+    67: 10,
+    68: 10,
+    69: 10,
+    70: 22,
+    71: 22,
+    72: 22,
+    73: 22,
+    74: 25,
+    75: 24,
+    76: 25,
+    77: 24,
+    78: 20,
+    79: 20,
+    80: 20,
+    81: 20,
+    82: 25,
+    83: 6,
+    84: 23,
+    85: 6,
+    86: 14,
+    87: 9,
+    88: 14,
+    89: 9,
+}
+
+# List of organisms used for fine-tuning
+FINE_TUNE_ORGANISMS: List[str] = [
+    "Arabidopsis thaliana",
+    "Bacillus subtilis",
+    "Caenorhabditis elegans",
+    "Chlamydomonas reinhardtii",
+    "Chlamydomonas reinhardtii chloroplast",
+    "Danio rerio",
+    "Drosophila melanogaster",
+    "Homo sapiens",
+    "Mus musculus",
+    "Nicotiana tabacum",
+    "Nicotiana tabacum chloroplast",
+    "Pseudomonas putida",
+    "Saccharomyces cerevisiae",
+    "Escherichia coli O157-H7 str. Sakai",
+    "Escherichia coli general",
+    "Escherichia coli str. K-12 substr. MG1655",
+    "Thermococcus barophilus MPT",
+]
+
+# List of organisms most commonly used for coodn optimization
+COMMON_ORGANISMS: List[str] = [
+    "Arabidopsis thaliana",
+    "Bacillus subtilis",
+    "Caenorhabditis elegans",
+    "Chlamydomonas reinhardtii",
+    "Danio rerio",
+    "Drosophila melanogaster",
+    "Homo sapiens",
+    "Mus musculus",
+    "Nicotiana tabacum",
+    "Pseudomonas putida",
+    "Saccharomyces cerevisiae",
+    "Escherichia coli general",
+]
+
+# Dictionary mapping each organism name to respective organism id
+ORGANISM2ID: Dict[str, int] = {
+    "Arabidopsis thaliana": 0,
+    "Atlantibacter hermannii": 1,
+    "Bacillus subtilis": 2,
+    "Brenneria goodwinii": 3,
+    "Buchnera aphidicola (Schizaphis graminum)": 4,
+    "Caenorhabditis elegans": 5,
+    "Candidatus Erwinia haradaeae": 6,
+    "Candidatus Hamiltonella defensa 5AT (Acyrthosiphon pisum)": 7,
+    "Chlamydomonas reinhardtii": 8,
+    "Chlamydomonas reinhardtii chloroplast": 9,
+    "Citrobacter amalonaticus": 10,
+    "Citrobacter braakii": 11,
+    "Citrobacter cronae": 12,
+    "Citrobacter europaeus": 13,
+    "Citrobacter farmeri": 14,
+    "Citrobacter freundii": 15,
+    "Citrobacter koseri ATCC BAA-895": 16,
+    "Citrobacter portucalensis": 17,
+    "Citrobacter werkmanii": 18,
+    "Citrobacter youngae": 19,
+    "Cronobacter dublinensis subsp. dublinensis LMG 23823": 20,
+    "Cronobacter malonaticus LMG 23826": 21,
+    "Cronobacter sakazakii": 22,
+    "Cronobacter turicensis": 23,
+    "Danio rerio": 24,
+    "Dickeya dadantii 3937": 25,
+    "Dickeya dianthicola": 26,
+    "Dickeya fangzhongdai": 27,
+    "Dickeya solani": 28,
+    "Dickeya zeae": 29,
+    "Drosophila melanogaster": 30,
+    "Edwardsiella anguillarum ET080813": 31,
+    "Edwardsiella ictaluri": 32,
+    "Edwardsiella piscicida": 33,
+    "Edwardsiella tarda": 34,
+    "Enterobacter asburiae": 35,
+    "Enterobacter bugandensis": 36,
+    "Enterobacter cancerogenus": 37,
+    "Enterobacter chengduensis": 38,
+    "Enterobacter cloacae": 39,
+    "Enterobacter hormaechei": 40,
+    "Enterobacter kobei": 41,
+    "Enterobacter ludwigii": 42,
+    "Enterobacter mori": 43,
+    "Enterobacter quasiroggenkampii": 44,
+    "Enterobacter roggenkampii": 45,
+    "Enterobacter sichuanensis": 46,
+    "Erwinia amylovora CFBP1430": 47,
+    "Erwinia persicina": 48,
+    "Escherichia albertii": 49,
+    "Escherichia coli O157-H7 str. Sakai": 50,
+    "Escherichia coli general": 51,
+    "Escherichia coli str. K-12 substr. MG1655": 52,
+    "Escherichia fergusonii": 53,
+    "Escherichia marmotae": 54,
+    "Escherichia ruysiae": 55,
+    "Ewingella americana": 56,
+    "Hafnia alvei": 57,
+    "Hafnia paralvei": 58,
+    "Homo sapiens": 59,
+    "Kalamiella piersonii": 60,
+    "Klebsiella aerogenes": 61,
+    "Klebsiella grimontii": 62,
+    "Klebsiella michiganensis": 63,
+    "Klebsiella oxytoca": 64,
+    "Klebsiella pasteurii": 65,
+    "Klebsiella pneumoniae subsp. pneumoniae HS11286": 66,
+    "Klebsiella quasipneumoniae": 67,
+    "Klebsiella quasivariicola": 68,
+    "Klebsiella variicola": 69,
+    "Kosakonia cowanii": 70,
+    "Kosakonia radicincitans": 71,
+    "Leclercia adecarboxylata": 72,
+    "Lelliottia amnigena": 73,
+    "Lonsdalea populi": 74,
+    "Moellerella wisconsensis": 75,
+    "Morganella morganii": 76,
+    "Mus musculus": 77,
+    "Nicotiana tabacum": 78,
+    "Nicotiana tabacum chloroplast": 79,
+    "Obesumbacterium proteus": 80,
+    "Pantoea agglomerans": 81,
+    "Pantoea allii": 82,
+    "Pantoea ananatis PA13": 83,
+    "Pantoea dispersa": 84,
+    "Pantoea stewartii": 85,
+    "Pantoea vagans": 86,
+    "Pectobacterium aroidearum": 87,
+    "Pectobacterium atrosepticum": 88,
+    "Pectobacterium brasiliense": 89,
+    "Pectobacterium carotovorum": 90,
+    "Pectobacterium odoriferum": 91,
+    "Pectobacterium parmentieri": 92,
+    "Pectobacterium polaris": 93,
+    "Pectobacterium versatile": 94,
+    "Photorhabdus laumondii subsp. laumondii TTO1": 95,
+    "Plesiomonas shigelloides": 96,
+    "Pluralibacter gergoviae": 97,
+    "Proteus faecis": 98,
+    "Proteus mirabilis HI4320": 99,
+    "Proteus penneri": 100,
+    "Proteus terrae subsp. cibarius": 101,
+    "Proteus vulgaris": 102,
+    "Providencia alcalifaciens": 103,
+    "Providencia heimbachae": 104,
+    "Providencia rettgeri": 105,
+    "Providencia rustigianii": 106,
+    "Providencia stuartii": 107,
+    "Providencia thailandensis": 108,
+    "Pseudomonas putida": 109,
+    "Pyrococcus furiosus": 110,
+    "Pyrococcus horikoshii": 111,
+    "Pyrococcus yayanosii": 112,
+    "Rahnella aquatilis CIP 78.65 = ATCC 33071": 113,
+    "Raoultella ornithinolytica": 114,
+    "Raoultella planticola": 115,
+    "Raoultella terrigena": 116,
+    "Rosenbergiella epipactidis": 117,
+    "Rouxiella badensis": 118,
+    "Saccharolobus solfataricus": 119,
+    "Saccharomyces cerevisiae": 120,
+    "Salmonella bongori N268-08": 121,
+    "Salmonella enterica subsp. enterica serovar Typhimurium str. LT2": 122,
+    "Serratia bockelmannii": 123,
+    "Serratia entomophila": 124,
+    "Serratia ficaria": 125,
+    "Serratia fonticola": 126,
+    "Serratia grimesii": 127,
+    "Serratia liquefaciens": 128,
+    "Serratia marcescens": 129,
+    "Serratia nevei": 130,
+    "Serratia plymuthica AS9": 131,
+    "Serratia proteamaculans": 132,
+    "Serratia quinivorans": 133,
+    "Serratia rubidaea": 134,
+    "Serratia ureilytica": 135,
+    "Shigella boydii": 136,
+    "Shigella dysenteriae": 137,
+    "Shigella flexneri 2a str. 301": 138,
+    "Shigella sonnei": 139,
+    "Thermoccoccus kodakarensis": 140,
+    "Thermococcus barophilus MPT": 141,
+    "Thermococcus chitonophagus": 142,
+    "Thermococcus gammatolerans": 143,
+    "Thermococcus litoralis": 144,
+    "Thermococcus onnurineus": 145,
+    "Thermococcus sibiricus": 146,
+    "Xenorhabdus bovienii str. feltiae Florida": 147,
+    "Yersinia aldovae 670-83": 148,
+    "Yersinia aleksiciae": 149,
+    "Yersinia alsatica": 150,
+    "Yersinia enterocolitica": 151,
+    "Yersinia frederiksenii ATCC 33641": 152,
+    "Yersinia intermedia": 153,
+    "Yersinia kristensenii": 154,
+    "Yersinia massiliensis CCUG 53443": 155,
+    "Yersinia mollaretii ATCC 43969": 156,
+    "Yersinia pestis A1122": 157,
+    "Yersinia proxima": 158,
+    "Yersinia pseudotuberculosis IP 32953": 159,
+    "Yersinia rochesterensis": 160,
+    "Yersinia rohdei": 161,
+    "Yersinia ruckeri": 162,
+    "Yokenella regensburgei": 163,
+}
+
+# Dictionary mapping each organism id to respective organism name
+ID2ORGANISM = {v: k for k, v in ORGANISM2ID.items()}
+
+# Type alias for amino acid to codon mapping
+AMINO2CODON_TYPE = Dict[str, Tuple[List[str], List[float]]]
+
+# Constants for the number of organisms and sequence lengths
+NUM_ORGANISMS = 164
+MAX_LEN = 2048
+MAX_AMINO_ACIDS = MAX_LEN - 2  # Without special tokens [CLS] and [SEP]
+STOP_SYMBOL = "_"
+
+
+@dataclass
+class DNASequencePrediction:
+    """
+    A class to hold the output of the DNA sequence prediction.
+
+    Attributes:
+        organism (str): Name of the organism used for prediction.
+        protein (str): Input protein sequence for which DNA sequence is predicted.
+        processed_input (str): Processed input sequence (merged protein and DNA).
+        predicted_dna (str): Predicted DNA sequence.
+    """
+
+    organism: str
+    protein: str
+    processed_input: str
+    predicted_dna: str
+
+
+class IterableData(torch.utils.data.IterableDataset):
+    """
+    Defines the logic for iterable datasets (working over streams of
+    data) in parallel multi-processing environments, e.g., multi-GPU.
+
+    Args:
+        dist_env (Optional[str]): The distribution environment identifier
+        (e.g., "slurm").
+
+    Credit: Guillaume Filion
+    """
+
+    def __init__(self, dist_env: Optional[str] = None):
+        super().__init__()
+        self.world_size_handle, self.rank_handle = {
+            "slurm": ("SLURM_NTASKS", "SLURM_PROCID")
+        }.get(dist_env, ("WORLD_SIZE", "LOCAL_RANK"))
+
+    @property
+    def iterator(self) -> Iterator:
+        """Define the stream logic for the dataset. Implement in subclasses."""
+        raise NotImplementedError
+
+    def __iter__(self) -> Iterator:
+        """
+        Create an iterator for the dataset, handling multi-processing contexts.
+
+        Returns:
+            Iterator: The iterator for the dataset.
+        """
+        worker_info = torch.utils.data.get_worker_info()
+        if worker_info is None:
+            return self.iterator
+
+        # In multi-processing context, use 'os.environ' to
+        # find global worker rank. Then use 'islice' to allocate
+        # the items of the stream to the workers.
+        world_size = int(os.environ.get(self.world_size_handle))
+        global_rank = int(os.environ.get(self.rank_handle))
+        local_rank = worker_info.id
+        local_num_workers = worker_info.num_workers
+
+        # Assume that each process has the same number of local workers.
+        worker_rk = global_rank * local_num_workers + local_rank
+        worker_nb = world_size * local_num_workers
+        return itertools.islice(self.iterator, worker_rk, None, worker_nb)
+
+
+class IterableJSONData(IterableData):
+    """
+    Iterate over the lines of a JSON file and uncompress if needed.
+
+    Args:
+        data_path (str): The path to the JSON data file.
+        train (bool): Flag indicating if the dataset is for training.
+        **kwargs: Additional keyword arguments for the base class.
+    """
+
+    def __init__(self, data_path: str, train: bool = True, **kwargs):
+        super().__init__(**kwargs)
+        self.data_path = data_path
+        self.train = train
+
+    @property
+    def iterator(self) -> Iterator:
+        open_fn = gzip.open if self.data_path.endswith(".gz") else open
+        with open_fn(self.data_path, "rt") as file:
+            for line in file:
+                line = line.strip()
+                if line:
+                    yield json.loads(line)
+
+
+class ConfigManager(ABC):
+    """
+    Abstract base class for managing configuration settings.
+    """
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_type is not None:
+            print(f"Exception occurred: {exc_type}, {exc_value}, {traceback}")
+        self.reset_config()
+
+    @abstractmethod
+    def reset_config(self) -> None:
+        """Reset the configuration to default values."""
+        pass
+
+    def get(self, key: str) -> Any:
+        """
+        Get the value of a configuration key.
+
+        Args:
+            key (str): The key to retrieve the value for.
+
+        Returns:
+            Any: The value of the configuration key.
+        """
+        return self._config.get(key)
+
+    def set(self, key: str, value: Any) -> None:
+        """
+        Set the value of a configuration key.
+
+        Args:
+            key (str): The key to set the value for.
+            value (Any): The value to set for the key.
+        """
+        self.validate_inputs(key, value)
+        self._config[key] = value
+
+    def update(self, config_dict: dict) -> None:
+        """
+        Update the configuration with a dictionary of key-value pairs after validating them.
+
+        Args:
+            config_dict (dict): A dictionary of key-value pairs to update the configuration.
+        """
+        for key, value in config_dict.items():
+            self.validate_inputs(key, value)
+        self._config.update(config_dict)
+
+    @abstractmethod
+    def validate_inputs(self, key: str, value: Any) -> None:
+        """Validate the inputs for the configuration."""
+        pass
+
+
+class ProteinConfig(ConfigManager):
+    """
+    A class to manage configuration settings for protein sequences.
+
+    This class ensures that the configuration is a singleton.
+    It provides methods to get, set, and update configuration values.
+
+    Attributes:
+        _instance (Optional[ConfigManager]): The singleton instance of the ConfigManager.
+        _config (Dict[str, Any]): The configuration dictionary.
+    """
+
+    _instance = None
+
+    def __new__(cls):
+        """
+        Create a new instance of the ProteinConfig class.
+
+        Returns:
+            ProteinConfig: The singleton instance of the ProteinConfig.
+        """
+        if cls._instance is None:
+            cls._instance = super(ProteinConfig, cls).__new__(cls)
+            cls._instance.reset_config()
+        return cls._instance
+
+    def validate_inputs(self, key: str, value: Any) -> None:
+        """
+        Validate the inputs for the configuration.
+
+        Args:
+            key (str): The key to validate.
+            value (Any): The value to validate.
+
+        Raises:
+            ValueError: If the value is invalid.
+            TypeError: If the value is of the wrong type.
+        """
+        if key == "ambiguous_aminoacid_behavior":
+            if value not in [
+                "raise_error",
+                "standardize_deterministic",
+                "standardize_random",
+            ]:
+                raise ValueError(
+                    f"Invalid value for ambiguous_aminoacid_behavior: {value}."
+                )
+        elif key == "ambiguous_aminoacid_map_override":
+            if not isinstance(value, dict):
+                raise TypeError(
+                    f"Invalid type for ambiguous_aminoacid_map_override: {value}."
+                )
+            for ambiguous_aminoacid, aminoacids in value.items():
+                if not isinstance(aminoacids, list):
+                    raise TypeError(f"Invalid type for aminoacids: {aminoacids}.")
+                if not aminoacids:
+                    raise ValueError(
+                        f"Override for aminoacid '{ambiguous_aminoacid}' cannot be empty list."
+                    )
+                if ambiguous_aminoacid not in AMBIGUOUS_AMINOACID_MAP:
+                    raise ValueError(
+                        f"Invalid amino acid in ambiguous_aminoacid_map_override: {ambiguous_aminoacid}"
+                    )
+        else:
+            raise ValueError(f"Invalid configuration key: {key}")
+
+    def reset_config(self) -> None:
+        """
+        Reset the configuration to the default values.
+        """
+        self._config = {
+            "ambiguous_aminoacid_behavior": "standardize_random",
+            "ambiguous_aminoacid_map_override": {},
+        }
+
+
+def load_python_object_from_disk(file_path: str) -> Any:
+    """
+    Load a Pickle object from disk and return it as a Python object.
+
+    Args:
+        file_path (str): The path to the Pickle file.
+
+    Returns:
+        Any: The loaded Python object.
+    """
+    with open(file_path, "rb") as file:
+        return pickle.load(file)
+
+
+def save_python_object_to_disk(input_object: Any, file_path: str) -> None:
+    """
+    Save a Python object to disk using Pickle.
+
+    Args:
+        input_object (Any): The Python object to save.
+        file_path (str): The path where the object will be saved.
+    """
+    with open(file_path, "wb") as file:
+        pickle.dump(input_object, file)
+
+
+def find_pattern_in_fasta(keyword: str, text: str) -> str:
+    """
+    Find a specific keyword pattern in text. Helpful for identifying parts
+    of a FASTA sequence.
+
+    Args:
+        keyword (str): The keyword pattern to search for.
+        text (str): The text to search within.
+
+    Returns:
+        str: The found pattern or an empty string if not found.
+    """
+    # Search for the keyword pattern in the text using regex
+    result = re.search(keyword + r"=(.*?)]", text)
+    return result.group(1) if result else ""
+
+
+def get_organism2id_dict(organism_reference: str) -> Dict[str, int]:
+    """
+    Return a dictionary mapping each organism in training data to an index
+    used for training.
+
+    Args:
+        organism_reference (str): Path to a CSV file containing a list of
+            all organisms. The format of the CSV file should be as follows:
+
+                0,Escherichia coli
+                1,Homo sapiens
+                2,Mus musculus
+
+    Returns:
+        Dict[str, int]: Dictionary mapping organism names to their respective indices.
+    """
+    # Read the CSV file and create a dictionary mapping organisms to their indices
+    organisms = pd.read_csv(organism_reference, index_col=0, header=None)
+    organism2id = {organisms.iloc[i].values[0]: i for i in organisms.index}
+
+    return organism2id
+
+
+def get_taxonomy_id(
+    taxonomy_reference: str, organism: Optional[str] = None, return_dict: bool = False
+) -> Any:
+    """
+    Return the taxonomy id of a given organism using a reference file.
+    Optionally, return the whole dictionary instead if return_dict is True.
+
+    Args:
+        taxonomy_reference (str): Path to the taxonomy reference file.
+        organism (Optional[str]): The name of the organism to look up.
+        return_dict (bool): Whether to return the entire dictionary.
+
+    Returns:
+        Any: The taxonomy id of the organism or the entire dictionary.
+    """
+    # Load the organism-to-taxonomy mapping from a Pickle file
+    organism2taxonomy = load_python_object_from_disk(taxonomy_reference)
+
+    if return_dict:
+        return dict(sorted(organism2taxonomy.items()))
+
+    return organism2taxonomy[organism]
+
+
+def sort_amino2codon_skeleton(amino2codon: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Sort the amino2codon dictionary alphabetically by amino acid and by codon name.
+
+    Args:
+        amino2codon (Dict[str, Any]): The amino2codon dictionary to sort.
+
+    Returns:
+        Dict[str, Any]: The sorted amino2codon dictionary.
+    """
+    # Sort the dictionary by amino acid and then by codon name
+    amino2codon = dict(sorted(amino2codon.items()))
+    amino2codon = {
+        amino: (
+            [codon for codon, _ in sorted(zip(codons, frequencies))],
+            [freq for _, freq in sorted(zip(codons, frequencies))],
+        )
+        for amino, (codons, frequencies) in amino2codon.items()
+    }
+
+    return amino2codon
+
+
+def load_pkl_from_url(url: str) -> Any:
+    """
+    Download a Pickle file from a URL and return the loaded object.
+
+    Args:
+        url (str): The URL to download the Pickle file from.
+
+    Returns:
+        Any: The loaded Python object from the Pickle file.
+    """
+    response = requests.get(url)
+    response.raise_for_status()  # Ensure the request was successful
+
+    # Load the Pickle object from the response content
+    return pickle.loads(response.content)

model/CodonTransformer/__init__.py

@@ -0,0 +1 @@
+"""CodonTransformer package."""

model/src/CodonTransformerTokenizer.json

@@ -0,0 +1 @@
+{"version": "1.0", "truncation": null, "padding": null, "added_tokens": [{"id": 0, "special": true, "content": "[UNK]", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false}, {"id": 1, "special": true, "content": "[CLS]", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false}, {"id": 2, "special": true, "content": "[SEP]", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false}, {"id": 3, "special": true, "content": "[PAD]", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false}, {"id": 4, "special": true, "content": "[MASK]", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false}], "normalizer": {"type": "Sequence", "normalizers": [{"type": "Lowercase"}]}, "pre_tokenizer": {"type": "Sequence", "pretokenizers": [{"type": "Split", "pattern": {"String": " "}, "behavior": "Isolated", "invert": false}, {"type": "Whitespace"}]}, "post_processor": {"type": "TemplateProcessing", "single": [{"SpecialToken": {"id": "[CLS]", "type_id": 0}}, {"Sequence": {"id": "A", "type_id": 0}}, {"SpecialToken": {"id": "[SEP]", "type_id": 0}}], "pair": [{"SpecialToken": {"id": "[CLS]", "type_id": 0}}, {"Sequence": {"id": "A", "type_id": 0}}, {"SpecialToken": {"id": "[SEP]", "type_id": 0}}, {"Sequence": {"id": "B", "type_id": 1}}, {"SpecialToken": {"id": "[SEP]", "type_id": 1}}], "special_tokens": {"[CLS]": {"id": "[CLS]", "ids": [1], "tokens": ["[CLS]"]}, "[SEP]": {"id": "[SEP]", "ids": [2], "tokens": ["[SEP]"]}}}, "decoder": null, "model": {"type": "WordPiece", "unk_token": "[UNK]", "continuing_subword_prefix": "##", "max_input_chars_per_word": 100, "vocab": {"[UNK]": 0, "[CLS]": 1, "[SEP]": 2, "[PAD]": 3, "[MASK]": 4, "a_unk": 5, "c_unk": 6, "d_unk": 7, "e_unk": 8, "f_unk": 9, "g_unk": 10, "h_unk": 11, "i_unk": 12, "k_unk": 13, "l_unk": 14, "m_unk": 15, "n_unk": 16, "p_unk": 17, "q_unk": 18, "r_unk": 19, "s_unk": 20, "t_unk": 21, "v_unk": 22, "w_unk": 23, "y_unk": 24, "__unk": 25, "k_aaa": 26, "n_aac": 27, "k_aag": 28, "n_aat": 29, "t_aca": 30, "t_acc": 31, "t_acg": 32, "t_act": 33, "r_aga": 34, "s_agc": 35, "r_agg": 36, "s_agt": 37, "i_ata": 38, "i_atc": 39, "m_atg": 40, "i_att": 41, "q_caa": 42, "h_cac": 43, "q_cag": 44, "h_cat": 45, "p_cca": 46, "p_ccc": 47, "p_ccg": 48, "p_cct": 49, "r_cga": 50, "r_cgc": 51, "r_cgg": 52, "r_cgt": 53, "l_cta": 54, "l_ctc": 55, "l_ctg": 56, "l_ctt": 57, "e_gaa": 58, "d_gac": 59, "e_gag": 60, "d_gat": 61, "a_gca": 62, "a_gcc": 63, "a_gcg": 64, "a_gct": 65, "g_gga": 66, "g_ggc": 67, "g_ggg": 68, "g_ggt": 69, "v_gta": 70, "v_gtc": 71, "v_gtg": 72, "v_gtt": 73, "__taa": 74, "y_tac": 75, "__tag": 76, "y_tat": 77, "s_tca": 78, "s_tcc": 79, "s_tcg": 80, "s_tct": 81, "__tga": 82, "c_tgc": 83, "w_tgg": 84, "c_tgt": 85, "l_tta": 86, "f_ttc": 87, "l_ttg": 88, "f_ttt": 89}}}

model/src/__init__.py

@@ -0,0 +1 @@
+"""Model weights, tokenizer, and other resources."""

model/src/organism2id.pkl

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

pyproject.toml

@@ -0,0 +1,60 @@
+[tool.poetry]
+name = "CodonTransformer"
+version = "1.6.7"
+description = "The ultimate tool for codon optimization, transforming protein sequences into optimized DNA sequences specific for your target organisms."
+authors = ["Adibvafa Fallahpour <Adibvafa.fallahpour@mail.utoronto.ca>"]
+license = "Apache-2.0"
+readme = "README.md"
+homepage = "https://github.com/adibvafa/CodonTransformer"
+repository = "https://github.com/adibvafa/CodonTransformer"
+packages = [{ include = "CodonTransformer", from = "model" }]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: Apache Software License",
+    "Operating System :: OS Independent",
+]
+
+[tool.poetry.dependencies]
+python = "^3.9"
+biopython = "^1.83"
+ipywidgets = "^7.0.0"
+numpy = "<2.0.0"
+onnxruntime = "^1.16.3"
+pandas = "^2.0.0"
+python_codon_tables = "^0.1.12"
+pytorch_lightning = "^2.2.1"
+scikit-learn = "^1.2.2"
+scipy = "^1.13.1"
+setuptools = "^70.0.0"
+torch = "^2.0.0"
+tqdm = "^4.66.2"
+transformers = "^4.40.0"
+CAI-PyPI = "^2.0.1"
+
+[tool.poetry.dev-dependencies]
+coverage = {version = "^7.0", extras = ["toml"]}
+
+[build-system]
+requires = ["poetry-core>=1.0.0"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.ruff]
+line-length = 88
+indent-width = 4
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E", "F", "I"]
+ignore = []
+
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+skip-magic-trailing-comma = false
+line-ending = "auto"
+
+[tool.coverage.run]
+omit = [
+    # omit pytorch-generated files in /tmp
+    "/tmp/*",
+]

requirements.txt

@@ -0,0 +1,14 @@
+biopython>=1.83,<2.0
+CAI-PyPI>=2.0.1,<3.0
+ipywidgets>=7.0.0,<10.0
+numpy==1.26.3
+onnxruntime>=1.16.3,<3.0
+pandas>=2.0.0,<3.0
+python_codon_tables>=0.1.12,<1.0
+pytorch_lightning>=2.2.1,<3.0
+scikit-learn>=1.2.2,<2.0
+scipy>=1.13.1,<3.0
+setuptools>=70.0.0
+torch>=2.0.0,<3.0
+tqdm>=4.66.2,<5.0
+transformers>=4.40.0,<5.0

scripts/demo/sample_dataset.csv

@@ -0,0 +1,6 @@
+,protein_sequence,organism,predicted_dna
+0,MSEKYIVTWDMLQIHARKLASRLMPSEQWKGIIAVSRGGLVPGALLARELGIRHVDTVCISSYDHDNQRELKVLKRAEGDGEGFIVIDDLVDTGGTAVAIREMYPKAHFVTIFAKPAGRPLVDDYVVDIPQDTWIEQPWDMGVVFVPPISGR_,Escherichia coli general,
+1,MKNIIRTPETHPLTWRLRDDKQPVWLDEYRSKNGYEGARKALTGLSPDEIVNQVKDAGLKGRGGAGFSTGLKWSLMPKDESMNIRYLLCNADEMEPGTYKDRLLMEQLPHLLVEGMLISAFALKAYRGYIFLRGEYIEAAVNLRRAIAEATEAGLLGKNIMGTGFDFELFVHTGAGRYICGEETALINSLEGRRANPRSKPPFPATSGAWGKPTCVNNVETLCNVPAILANGVEWYQNISKSKDAGTKLMGFSGRVKNPGLWELPFGTTAREILEDYAGGMRDGLKFKAWQPGGAGTDFLTEAHLDLPMEFESIGKAGSRLGTALAMAVDHEINMVSLVRNLEEFFARESCGWCTPCRDGLPWSVKILRALERGEGQPGDIETLEQLCRFLGPGKTFCAHAPGAVEPLQSAIKYFREEFEAGIKQPFSNTHLINGIQPNLLKERW_,Escherichia coli general,
+2,MDALQIAEDTLQTLVPHCPVPSGPRRIFLDANVKESYCPLVPHTMYCLPLWQGINLVLLTRSPSAPLALVLSQLMDGFSMLEKKLKEGPEPGASLRSQPLVGDLRQRMDKFVKNRGAQEIQSTWLEFKAKAFSKSEPGSSWELLQACGKLKRQLCAIYRLNFLTTAPSRGGPHLPQHLQDQVQRLMREKLTDWKDFLLVKSRRNITMVSYLEDFPGLVHFIYVDRTTGQMVAPSLNCSQKTSSELGKGPLAAFVKTKVWSLIQLARRYLQKGYTTLLFQEGDFYCSYFLWFENDMGYKLQMIEVPVLSDDSVPIGMLGGDYYRKLLRYYSKNRPTEAVRCYELLALHLSVIPTDLLVQQAGQLARRLWEASRIPLL_,Homo sapiens,
+3,MAFANFRRILRLSTFEKRKSREYEHVRRDLDPNEVWEIVGELGDGAFGKVYKAKNKETGALAAAKVIETKSEEELEDYIVEIEILATCDHPYIVKLLGAYYHDGKLWIMIEFCPGGAVDAIMLELDRGLTEPQIQVVCRQMLEALNFLHSKRIIHRDLKAGNVLMTLEGDIRLADFGVSAKNLKTLQKRDSFIGTPYWMAPEVVMCETMKDTPYDYKADIWSLGITLIEMAQIEPPHHELNPMRVLLKIAKSDPPTLLTPSKWSVEFRDFLKIALDKNPETRPSAAQLLEHPFVSSITSNKALRELVAEAKAEVMEEIEDGRDEGEEEDAVDAASDPKLYKKTLKRTRKFVVDGVEVSITTSKIISEDEKKDEEMRFLRRQELRELRLLQKEEHRNQTQLSNKHELQLEQMHKRFEQEINAKKKFFDTELENLERQQKQQVEKMEQDHAVRRREEARRIRLEQDRDYTRFQEQLKLMKKEVKNEVEKLPRQQRKESMKQKMEEHTQKKQLLDRDFVAKQKEDLELAMKRLTTDNRREICDKERECLMKKQELLRDREAALWEMEEHQLQERHQLVKQQLKDQYFLQRHELLRKHEKEREQMQRYNQRMIEQLKVRQQQEKARLPKIQRSEGKTRMAMYKKSLHINGGGSAAEQREKIKQFSQQEEKRQKSERLQQQQKHENQMRDMLAQCESNMSELQQLQNEKCHLLVEHETQKLKALDESHNQNLKEWRDKLRPRKKALEEDLNQKKREQEMFFKLSEEAECPNPSTPSKAAKFFPYSSADAS_,Homo sapiens,
+4,MTEKDAGGFNMSTFMNRKFQEPIQQIKTFSWMGFSWTCRKRRKHYQSYLRNGVRISVNDFVYVLAEQHKRLVAYIEDLYEDSKGKKMVVVRWFHKTEEVGSVLSDDDNDREIFFSLNRQDISIECIDYLATVLSPQHYEKFLKVPMHVQTVAFFCQKLYGDDGLKPYDITQLEGYWRQEMLRYLNVSILKSFEGAQAPGTDPGLKAPLVGCVGIRSRKRRRPSPVGTLNVSYAGDMKGDCKSSPDSVLAVTDASIFKGDEDGSSHHIKKGSLIEVLSEDSGIRGCWFKALVLKKHKDKVKVQYQDIQDADDESKKLEEWILTSRVAAGDHLGDLRIKGRKVVRPMLKPSKENDVCVIGVGMPVDVWWCDGWWEGIVVQEVSEEKFEVYLPGEKKMSAFHRNDLRQSREWLDDEWLNIRSRSDIVSSVLSLTKKKEMEVKHDEKSSDVGVCNGRMSPKTEAKRTISLPVATTKKSLPKRPIPDLLKDVLVTSDLKWKKSSRKRNRVVSCCPHDPSLNDGFSSERSLDCENCKFMEDTFGSSDGQHLTGLLMSR_,Arabidopsis thaliana,

scripts/demo/sample_predictions.csv

@@ -0,0 +1,6 @@
+,protein_sequence,organism,predicted_dna
+0,MSEKYIVTWDMLQIHARKLASRLMPSEQWKGIIAVSRGGLVPGALLARELGIRHVDTVCISSYDHDNQRELKVLKRAEGDGEGFIVIDDLVDTGGTAVAIREMYPKAHFVTIFAKPAGRPLVDDYVVDIPQDTWIEQPWDMGVVFVPPISGR_,Escherichia coli general,ATGAGCGAAAAATATATTGTCACCTGGGACATGCTGCAGATCCATGCCCGCAAACTGGCCAGCCGCCTGATGCCGTCAGAACAGTGGAAAGGCATTATTGCCGTCAGCCGCGGCGGCCTGGTGCCGGGTGCGCTGCTGGCGCGTGAGCTGGGTATTCGCCACGTCGACACCGTGTGCATCAGCAGCTATGACCACGACAACCAGCGCGAGCTGAAAGTGCTGAAACGTGCGGAAGGCGATGGCGAAGGCTTTATCGTCATTGATGATCTGGTTGATACCGGCGGCACCGCGGTGGCGATCCGTGAAATGTACCCGAAAGCGCACTTTGTCACCATCTTTGCGAAACCGGCAGGCCGTCCGCTGGTTGATGATTATGTGGTTGATATTCCGCAGGACACCTGGATCGAACAGCCGTGGGACATGGGCGTGGTGTTTGTTCCGCCGATCAGCGGCCGCTAA
+1,MKNIIRTPETHPLTWRLRDDKQPVWLDEYRSKNGYEGARKALTGLSPDEIVNQVKDAGLKGRGGAGFSTGLKWSLMPKDESMNIRYLLCNADEMEPGTYKDRLLMEQLPHLLVEGMLISAFALKAYRGYIFLRGEYIEAAVNLRRAIAEATEAGLLGKNIMGTGFDFELFVHTGAGRYICGEETALINSLEGRRANPRSKPPFPATSGAWGKPTCVNNVETLCNVPAILANGVEWYQNISKSKDAGTKLMGFSGRVKNPGLWELPFGTTAREILEDYAGGMRDGLKFKAWQPGGAGTDFLTEAHLDLPMEFESIGKAGSRLGTALAMAVDHEINMVSLVRNLEEFFARESCGWCTPCRDGLPWSVKILRALERGEGQPGDIETLEQLCRFLGPGKTFCAHAPGAVEPLQSAIKYFREEFEAGIKQPFSNTHLINGIQPNLLKERW_,Escherichia coli general,ATGAAAAATATTATTAGAACACCTGAAACCCATCCGCTGACCTGGCGTCTGCGCGATGACAAACAGCCGGTGTGGCTGGATGAGTACCGCAGCAAAAACGGCTATGAAGGTGCGCGTAAAGCGCTGACCGGTCTGTCTCCGGATGAGATTGTCAATCAGGTCAAAGATGCCGGCCTGAAAGGCCGTGGCGGTGCGGGTTTCTCCACCGGCCTGAAGTGGTCTCTGATGCCGAAAGATGAGAGCATGAACATCCGCTATCTGCTGTGCAATGCCGATGAGATGGAGCCGGGCACCTATAAAGACCGCCTGCTGATGGAGCAGCTGCCGCACCTGCTGGTTGAAGGTATGCTGATCTCTGCATTTGCGCTGAAAGCCTACCGTGGCTACATCTTCCTGCGTGGCGAGTACATCGAAGCGGCGGTGAACCTGCGCCGTGCGATTGCTGAAGCCACTGAAGCAGGTCTGCTGGGTAAAAACATCATGGGTACCGGTTTTGACTTTGAACTGTTCGTCCACACCGGTGCAGGGCGCTACATCTGCGGTGAAGAAACCGCGCTGATCAACAGCCTGGAAGGCCGTCGTGCGAACCCGCGCAGCAAACCGCCGTTCCCGGCAACCTCTGGTGCGTGGGGTAAACCGACCTGTGTTAATAACGTTGAAACCCTGTGCAACGTTCCGGCGATTCTGGCGAACGGTGTGGAATGGTATCAGAACATCTCCAAAAGCAAAGATGCTGGTACCAAGCTGATGGGTTTCTCCGGCCGTGTGAAAAACCCGGGCCTGTGGGAACTGCCGTTTGGTACCACCGCGCGTGAAATCCTGGAAGATTATGCCGGTGGCATGCGTGACGGCCTGAAGTTCAAAGCGTGGCAGCCGGGTGGTGCAGGTACCGATTTCCTGACTGAAGCGCACCTGGATCTGCCGATGGAGTTTGAGTCCATTGGTAAAGCGGGCAGCCGTCTGGGTACAGCGCTGGCGATGGCGGTTGACCATGAGATCAACATGGTGTCGCTGGTGCGTAACCTGGAAGAGTTCTTTGCCCGTGAAAGCTGCGGCTGGTGCACGCCGTGCCGCGACGGGCTGCCGTGGTCAGTGAAAATCCTGCGTGCGCTGGAGCGCGGTGAAGGCCAGCCGGGTGACATCGAAACGCTGGAACAGCTGTGCCGCTTCCTGGGTCCGGGTAAAACCTTCTGTGCACATGCACCCGGTGCGGTGGAACCGCTGCAGTCTGCGATCAAATATTTCCGTGAAGAGTTTGAAGCGGGGATCAAACAGCCGTTCTCCAACACCCATCTGATTAACGGGATTCAGCCGAACCTGCTGAAAGAGCGCTGGTAA
+2,MDALQIAEDTLQTLVPHCPVPSGPRRIFLDANVKESYCPLVPHTMYCLPLWQGINLVLLTRSPSAPLALVLSQLMDGFSMLEKKLKEGPEPGASLRSQPLVGDLRQRMDKFVKNRGAQEIQSTWLEFKAKAFSKSEPGSSWELLQACGKLKRQLCAIYRLNFLTTAPSRGGPHLPQHLQDQVQRLMREKLTDWKDFLLVKSRRNITMVSYLEDFPGLVHFIYVDRTTGQMVAPSLNCSQKTSSELGKGPLAAFVKTKVWSLIQLARRYLQKGYTTLLFQEGDFYCSYFLWFENDMGYKLQMIEVPVLSDDSVPIGMLGGDYYRKLLRYYSKNRPTEAVRCYELLALHLSVIPTDLLVQQAGQLARRLWEASRIPLL_,Homo sapiens,ATGGATGCCCTGCAGATTGCTGAGGACACCCTGCAGACCCTGGTGCCCCACTGCCCTGTGCCCTCTGGGCCCAGGAGGATCTTCCTGGATGCCAATGTGAAGGAGAGCTACTGCCCCCTGGTGCCCCACACCATGTACTGCCTGCCCCTGTGGCAGGGCATCAACCTGGTCCTGCTGACCAGGTCTCCCTCTGCCCCCCTGGCCCTGGTGCTGTCCCAGCTGATGGATGGCTTTTCCATGCTGGAGAAGAAGCTGAAGGAGGGGCCCGAGCCTGGAGCCTCTCTGAGGAGCCAGCCCCTGGTGGGGGACCTGCGGCAGAGAATGGACAAATTTGTGAAGAACCGAGGGGCCCAGGAGATCCAGAGCACCTGGCTGGAATTCAAGGCCAAGGCCTTCTCCAAATCTGAGCCTGGCAGCAGCTGGGAGCTGCTGCAGGCCTGTGGGAAGCTGAAGAGACAGCTGTGTGCCATCTACAGGCTGAACTTCCTGACCACAGCCCCCTCCAGAGGAGGGCCCCACCTGCCCCAGCACCTGCAGGACCAGGTGCAGCGGCTGATGCGGGAGAAGCTGACTGACTGGAAGGACTTCCTGCTGGTGAAGAGCAGGAGGAACATCACCATGGTGTCCTACCTGGAGGACTTCCCTGGCCTGGTGCACTTCATCTATGTGGACAGGACCACTGGGCAGATGGTGGCCCCCAGCCTGAACTGCAGCCAGAAGACCAGCTCTGAGCTGGGCAAGGGGCCCCTGGCTGCCTTTGTGAAGACCAAAGTGTGGAGCCTGATCCAGCTGGCCCGGAGATACCTGCAGAAGGGCTATACCACCCTGCTGTTCCAGGAAGGAGACTTCTACTGCTCCTACTTCCTGTGGTTTGAGAATGATATGGGCTACAAGCTGCAGATGATTGAGGTGCCTGTGCTGTCTGATGACTCTGTCCCCATTGGCATGCTGGGAGGAGACTACTACCGGAAGCTGCTGCGCTATTACAGCAAGAACCGGCCCACTGAGGCTGTGCGCTGCTATGAGCTGCTGGCCCTGCACCTGTCTGTGATCCCCACTGACCTGCTGGTGCAGCAGGCTGGGCAGCTGGCCAGGAGGCTGTGGGAGGCCTCCAGGATCCCCCTGCTGTGA
+3,MAFANFRRILRLSTFEKRKSREYEHVRRDLDPNEVWEIVGELGDGAFGKVYKAKNKETGALAAAKVIETKSEEELEDYIVEIEILATCDHPYIVKLLGAYYHDGKLWIMIEFCPGGAVDAIMLELDRGLTEPQIQVVCRQMLEALNFLHSKRIIHRDLKAGNVLMTLEGDIRLADFGVSAKNLKTLQKRDSFIGTPYWMAPEVVMCETMKDTPYDYKADIWSLGITLIEMAQIEPPHHELNPMRVLLKIAKSDPPTLLTPSKWSVEFRDFLKIALDKNPETRPSAAQLLEHPFVSSITSNKALRELVAEAKAEVMEEIEDGRDEGEEEDAVDAASDPKLYKKTLKRTRKFVVDGVEVSITTSKIISEDEKKDEEMRFLRRQELRELRLLQKEEHRNQTQLSNKHELQLEQMHKRFEQEINAKKKFFDTELENLERQQKQQVEKMEQDHAVRRREEARRIRLEQDRDYTRFQEQLKLMKKEVKNEVEKLPRQQRKESMKQKMEEHTQKKQLLDRDFVAKQKEDLELAMKRLTTDNRREICDKERECLMKKQELLRDREAALWEMEEHQLQERHQLVKQQLKDQYFLQRHELLRKHEKEREQMQRYNQRMIEQLKVRQQQEKARLPKIQRSEGKTRMAMYKKSLHINGGGSAAEQREKIKQFSQQEEKRQKSERLQQQQKHENQMRDMLAQCESNMSELQQLQNEKCHLLVEHETQKLKALDESHNQNLKEWRDKLRPRKKALEEDLNQKKREQEMFFKLSEEAECPNPSTPSKAAKFFPYSSADAS_,Homo sapiens,ATGGCCTTTGCCAACTTCCGGAGAATCCTGCGGCTGTCCACCTTTGAGAAGAGGAAGAGCCGGGAATATGAGCACGTGCGCAGGGACCTGGACCCCAACGAGGTGTGGGAGATTGTGGGGGAGCTGGGGGATGGTGCCTTTGGGAAGGTCTACAAGGCCAAGAACAAGGAGACTGGGGCCTTGGCTGCGGCCAAGGTAATTGAGACCAAATCTGAGGAGGAGCTGGAAGACTACATTGTGGAGATTGAGATTCTGGCCACCTGTGACCACCCCTACATTGTGAAGCTGCTGGGGGCCTACTACCATGATGGCAAGCTGTGGATCATGATCGAGTTCTGCCCTGGAGGGGCTGTGGATGCCATTATGCTGGAGCTGGACCGAGGCCTGACTGAGCCACAGATCCAGGTGGTGTGCAGGCAGATGCTGGAGGCCCTGAACTTCCTGCACAGCAAGAGAATCATTCACAGAGACCTGAAGGCTGGGAACGTGTTGATGACCCTGGAAGGAGACATCCGCTTGGCTGACTTTGGTGTCTCTGCCAAGAACCTGAAGACCTTACAGAAGAGGGACAGCTTCATTGGCACACCCTACTGGATGGCCCCTGAGGTGGTCATGTGTGAGACCATGAAGGACACACCCTATGACTACAAGGCTGACATTTGGAGCCTGGGCATCACCTTGATTGAGATGGCCCAGATTGAGCCACCACACCACGAGCTGAACCCAATGCGCGTGCTGCTGAAGATTGCCAAGTCAGACCCACCCACCTTGTTGACACCAAGCAAGTGGTCTGTGGAGTTCCGGGACTTCCTGAAGATTGCCCTGGACAAGAACCCTGAAACCAGGCCCTCTGCTGCCCAGCTGCTGGAGCACCCCTTTGTCTCCTCCATCACCTCCAACAAGGCCCTGCGGGAGCTGGTGGCTGAGGCCAAGGCTGAGGTGATGGAGGAGATTGAGGATGGCAGAGATGAAGGAGAGGAGGAGGATGCTGTGGATGCCGCCTCTGACCCCAAGCTGTACAAGAAGACCCTGAAGAGGACCCGCAAGTTTGTGGTGGACGGGGTGGAGGTGTCTATCACCACCTCCAAGATCATCTCTGAGGATGAGAAGAAGGATGAAGAGATGAGGTTCCTGCGGCGCCAGGAGCTGCGGGAGCTGCGGCTGCTGCAGAAGGAGGAGCACAGGAACCAGACCCAGCTGTCCAACAAGCATGAGCTGCAGCTGGAGCAGATGCACAAGAGGTTTGAGCAGGAGATCAATGCCAAGAAGAAGTTCTTTGACACTGAGCTGGAGAACCTGGAGAGGCAGCAGAAGCAGCAGGTGGAGAAGATGGAGCAGGACCATGCTGTGCGGAGGCGGGAGGAGGCCCGGAGGATCCGCCTGGAGCAGGACCGGGACTACACCCGCTTCCAGGAGCAGCTGAAGCTGATGAAGAAGGAGGTCAAGAACGAGGTGGAGAAGCTGCCCCGGCAGCAGAGGAAGGAAAGCATGAAGCAGAAGATGGAGGAGCACACCCAGAAGAAGCAGCTGCTGGACAGAGACTTTGTGGCCAAGCAGAAGGAGGACCTGGAGCTGGCTATGAAGAGACTGACCACTGATAACAGGAGGGAGATCTGTGACAAGGAGCGGGAGTGCCTGATGAAGAAGCAGGAGCTGCTGCGGGACCGGGAGGCTGCCCTGTGGGAGATGGAGGAGCACCAGCTGCAGGAGAGGCACCAGCTGGTCAAGCAGCAGCTGAAGGACCAGTACTTCCTGCAGAGGCACGAGCTGCTGAGGAAGCATGAGAAGGAGCGGGAGCAGATGCAGAGGTACAACCAGAGGATGATTGAGCAGCTGAAGGTGCGGCAGCAGCAGGAGAAGGCCAGACTGCCCAAGATCCAGAGATCTGAGGGGAAGACAAGGATGGCCATGTACAAGAAGTCTCTGCACATCAATGGGGGGGGCTCTGCCGCTGAGCAGCGGGAGAAGATCAAGCAGTTCTCCCAGCAGGAGGAGAAGAGGCAGAAGAGTGAGAGGCTGCAGCAGCAGCAGAAGCACGAGAACCAGATGAGGGACATGCTGGCCCAGTGTGAGTCCAACATGTCTGAGCTGCAGCAGCTGCAGAACGAGAAGTGTCACCTGCTGGTGGAGCATGAGACCCAGAAGCTGAAGGCCCTGGATGAGAGCCACAACCAGAATCTGAAGGAGTGGAGGGACAAGCTGAGGCCCAGGAAGAAGGCCCTGGAGGAGGACCTGAACCAGAAGAAGCGGGAGCAGGAGATGTTCTTCAAGCTGTCTGAGGAGGCCGAGTGCCCAAACCCAAGCACTCCATCCAAGGCTGCCAAGTTCTTCCCCTACAGCTCTGCCGACGCCAGCTAA
+4,MTEKDAGGFNMSTFMNRKFQEPIQQIKTFSWMGFSWTCRKRRKHYQSYLRNGVRISVNDFVYVLAEQHKRLVAYIEDLYEDSKGKKMVVVRWFHKTEEVGSVLSDDDNDREIFFSLNRQDISIECIDYLATVLSPQHYEKFLKVPMHVQTVAFFCQKLYGDDGLKPYDITQLEGYWRQEMLRYLNVSILKSFEGAQAPGTDPGLKAPLVGCVGIRSRKRRRPSPVGTLNVSYAGDMKGDCKSSPDSVLAVTDASIFKGDEDGSSHHIKKGSLIEVLSEDSGIRGCWFKALVLKKHKDKVKVQYQDIQDADDESKKLEEWILTSRVAAGDHLGDLRIKGRKVVRPMLKPSKENDVCVIGVGMPVDVWWCDGWWEGIVVQEVSEEKFEVYLPGEKKMSAFHRNDLRQSREWLDDEWLNIRSRSDIVSSVLSLTKKKEMEVKHDEKSSDVGVCNGRMSPKTEAKRTISLPVATTKKSLPKRPIPDLLKDVLVTSDLKWKKSSRKRNRVVSCCPHDPSLNDGFSSERSLDCENCKFMEDTFGSSDGQHLTGLLMSR_,Arabidopsis thaliana,ATGACGGAGAAAGATGCTGGAGGTTTTAATATGTCAACTTTCATGAACAGGAAGTTTCAAGAACCAATTCAACAGATCAAAACTTTCTCCTGGATGGGTTTCTCATGGACTTGTAGGAAGAGGAGGAAACATTATCAATCTTACCTTAGGAATGGAGTGAGGATCTCTGTCAATGATTTTGTTTATGTTCTTGCTGAGCAACACAAGAGGCTTGTTGCTTACATTGAAGATCTTTATGAGGATAGCAAAGGGAAGAAGATGGTTGTTGTTAGGTGGTTCCACAAGACTGAAGAGGTTGGATCTGTTCTTAGCGATGATGACAACGACAGGGAGATCTTCTTCTCTCTCAACAGACAAGACATCAGCATTGAGTGCATTGATTACCTTGCCACTGTTCTCTCTCCTCAACATTACGAGAAGTTTCTCAAGGTTCCTATGCATGTTCAAACTGTTGCTTTCTTCTGCCAGAAGCTCTATGGAGATGATGGTTTGAAACCTTATGACATCACTCAGCTTGAAGGTTACTGGAGACAAGAAATGCTCAGATACCTCAATGTCTCCATTCTCAAGAGCTTTGAAGGAGCTCAAGCTCCTGGAACTGATCCTGGTTTGAAGGCTCCTTTGGTTGGTTGTGTTGGTATCAGAAGCAGGAAGAGGAGGAGACCATCACCGGTTGGAACTCTCAACGTCAGCTACGCTGGAGACATGAAAGGAGACTGCAAAAGCTCTCCTGATTCTGTTTTGGCTGTCACTGATGCTTCGATCTTCAAAGGAGATGAAGATGGATCTTCTCACCACATCAAGAAAGGAAGCTTGATTGAGGTTCTCAGCGAGGACTCTGGGATCCGTGGTTGCTGGTTCAAAGCTTTGGTGTTGAAGAAACACAAGGACAAGGTGAAGGTGCAGTACCAAGACATTCAAGATGCTGATGATGAGAGCAAGAAGCTTGAGGAGTGGATTCTCACTAGCCGTGTTGCTGCTGGAGATCATCTTGGTGATTTGAGGATCAAAGGAAGGAAAGTTGTGAGACCAATGCTCAAACCTTCCAAGGAGAACGATGTTTGTGTGATTGGTGTTGGAATGCCGGTTGATGTTTGGTGGTGTGATGGATGGTGGGAAGGGATTGTGGTTCAAGAGGTTTCTGAGGAGAAGTTTGAGGTTTATCTTCCTGGAGAGAAGAAAATGTCAGCTTTCCACAGAAATGATTTGAGACAAAGCAGAGAGTGGCTTGATGATGAGTGGCTCAACATTAGAAGCAGAAGTGACATTGTTTCTTCTGTTCTTTCTTTGACCAAGAAGAAAGAGATGGAGGTGAAGCATGATGAGAAAAGCAGCGATGTTGGTGTCTGCAATGGAAGAATGTCTCCAAAAACAGAAGCTAAGAGAACAATCTCTCTTCCTGTTGCTACAACCAAGAAATCTCTTCCTAAGAGACCAATTCCTGATCTTCTCAAGGATGTGTTGGTCACTTCTGATTTGAAGTGGAAGAAAAGCTCAAGGAAAAGAAACAGAGTTGTTTCCTGCTGTCCTCATGATCCATCTCTCAATGATGGTTTCTCCTCTGAGAGATCTCTTGATTGTGAGAACTGCAAGTTCATGGAAGACACTTTTGGTTCTTCTGATGGACAACATCTCACTGGTCTTCTTATGTCCAGATAA

scripts/finetune.py

@@ -0,0 +1,274 @@
+"""
+File: finetune.py
+-------------------
+Finetune the CodonTransformer model.
+
+The pretrained model is loaded directly from Hugging Face.
+The dataset is a JSON file. You can use prepare_training_data from CodonData to
+prepare the dataset. The repository README has a guide on how to prepare the
+dataset and use this script.
+"""
+
+import argparse
+import gzip
+import math
+import os
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+MODEL_DIR = PROJECT_ROOT / "model"
+if str(MODEL_DIR) not in sys.path:
+    sys.path.insert(0, str(MODEL_DIR))
+
+import pytorch_lightning as pl
+import torch
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer, BigBirdForMaskedLM
+
+from CodonTransformer.CodonUtils import (
+    MAX_LEN,
+    TOKEN2MASK,
+    IterableJSONData,
+)
+
+
+class MaskedTokenizerCollator:
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    def __call__(self, examples):
+        tokenized = self.tokenizer(
+            [ex["codons"] for ex in examples],
+            return_attention_mask=True,
+            return_token_type_ids=True,
+            truncation=True,
+            padding=True,
+            max_length=MAX_LEN,
+            return_tensors="pt",
+        )
+
+        seq_len = tokenized["input_ids"].shape[-1]
+        species_index = torch.tensor([[ex["organism"]] for ex in examples])
+        tokenized["token_type_ids"] = species_index.repeat(1, seq_len)
+
+        inputs = tokenized["input_ids"]
+        targets = tokenized["input_ids"].clone()
+
+        prob_matrix = torch.full(inputs.shape, 0.15)
+        prob_matrix[torch.where(inputs < 5)] = 0.0
+        selected = torch.bernoulli(prob_matrix).bool()
+
+        # 80% of the time, replace masked input tokens with respective mask tokens
+        replaced = torch.bernoulli(torch.full(selected.shape, 0.8)).bool() & selected
+        inputs[replaced] = torch.tensor(
+            list((map(TOKEN2MASK.__getitem__, inputs[replaced].numpy())))
+        )
+
+        # 10% of the time, we replace masked input tokens with random vector.
+        randomized = (
+            torch.bernoulli(torch.full(selected.shape, 0.1)).bool()
+            & selected
+            & ~replaced
+        )
+        random_idx = torch.randint(26, 90, prob_matrix.shape, dtype=torch.long)
+        inputs[randomized] = random_idx[randomized]
+
+        tokenized["input_ids"] = inputs
+        tokenized["labels"] = torch.where(selected, targets, -100)
+
+        return tokenized
+
+
+class plTrainHarness(pl.LightningModule):
+    def __init__(self, model, learning_rate, warmup_fraction, total_training_steps):
+        super().__init__()
+        self.model = model
+        self.learning_rate = learning_rate
+        self.warmup_fraction = warmup_fraction
+        self.total_training_steps = total_training_steps
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(
+            self.model.parameters(),
+            lr=self.learning_rate,
+        )
+        total_steps = self.total_training_steps or self.trainer.estimated_stepping_batches
+        if total_steps <= 0:
+            raise ValueError(f"Expected positive integer total_steps, but got {total_steps}")
+        lr_scheduler = {
+            "scheduler": torch.optim.lr_scheduler.OneCycleLR(
+                optimizer,
+                max_lr=self.learning_rate,
+                total_steps=total_steps,
+                pct_start=self.warmup_fraction,
+            ),
+            "interval": "step",
+            "frequency": 1,
+        }
+        return [optimizer], [lr_scheduler]
+
+    def training_step(self, batch, batch_idx):
+        self.model.bert.set_attention_type("block_sparse")
+        outputs = self.model(**batch)
+        self.log_dict(
+            dictionary={
+                "loss": outputs.loss,
+                "lr": self.trainer.optimizers[0].param_groups[0]["lr"],
+            },
+            on_step=True,
+            prog_bar=True,
+        )
+        return outputs.loss
+
+
+class DumpStateDict(pl.callbacks.ModelCheckpoint):
+    def __init__(self, checkpoint_dir, checkpoint_filename, every_n_train_steps):
+        super().__init__(
+            dirpath=checkpoint_dir, every_n_train_steps=every_n_train_steps
+        )
+        self.checkpoint_filename = checkpoint_filename
+
+    def on_save_checkpoint(self, trainer, pl_module, checkpoint):
+        model = pl_module.model
+        torch.save(
+            model.state_dict(), os.path.join(self.dirpath, self.checkpoint_filename)
+        )
+
+
+def count_jsonl_records(path):
+    open_fn = gzip.open if path.endswith(".gz") else open
+    with open_fn(path, "rt") as file:
+        return sum(1 for line in file if line.strip())
+
+
+def estimate_training_steps(args):
+    num_records = count_jsonl_records(args.dataset_dir)
+    num_devices = 1 if args.debug else args.num_gpus
+    samples_per_step = max(1, args.batch_size * num_devices)
+    batches_per_epoch = math.ceil(num_records / samples_per_step)
+    optimizer_steps_per_epoch = math.ceil(
+        batches_per_epoch / max(1, args.accumulate_grad_batches)
+    )
+    total_steps = max(1, optimizer_steps_per_epoch * args.max_epochs)
+    print(
+        "Estimated training steps: "
+        f"{total_steps} "
+        f"({num_records} records, batch_size={args.batch_size}, "
+        f"devices={num_devices}, max_epochs={args.max_epochs}, "
+        f"accumulate_grad_batches={args.accumulate_grad_batches})"
+    )
+    return total_steps
+
+
+def main(args):
+    """Finetune the CodonTransformer model."""
+    pl.seed_everything(args.seed)
+    torch.set_float32_matmul_precision("medium")
+    total_training_steps = estimate_training_steps(args)
+
+    # Load the tokenizer and model
+    tokenizer = AutoTokenizer.from_pretrained("adibvafa/CodonTransformer")
+    model = BigBirdForMaskedLM.from_pretrained("adibvafa/CodonTransformer-base")
+    harnessed_model = plTrainHarness(
+        model,
+        args.learning_rate,
+        args.warmup_fraction,
+        total_training_steps,
+    )
+
+    # Load the training data
+    train_data = IterableJSONData(args.dataset_dir, dist_env="slurm")
+    data_loader = DataLoader(
+        dataset=train_data,
+        collate_fn=MaskedTokenizerCollator(tokenizer),
+        batch_size=args.batch_size,
+        num_workers=0 if args.debug else args.num_workers,
+        persistent_workers=False if args.debug else True,
+    )
+
+    # Setup trainer and callbacks
+    save_checkpoint = DumpStateDict(
+        checkpoint_dir=args.checkpoint_dir,
+        checkpoint_filename=args.checkpoint_filename,
+        every_n_train_steps=args.save_every_n_steps,
+    )
+    trainer = pl.Trainer(
+        default_root_dir=args.checkpoint_dir,
+        strategy="ddp_find_unused_parameters_true",
+        accelerator="gpu",
+        devices=1 if args.debug else args.num_gpus,
+        precision="16-mixed",
+        max_epochs=args.max_epochs,
+        deterministic=False,
+        enable_checkpointing=True,
+        callbacks=[save_checkpoint],
+        accumulate_grad_batches=args.accumulate_grad_batches,
+    )
+
+    # Finetune the model
+    trainer.fit(harnessed_model, data_loader)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Finetune the CodonTransformer model.")
+    parser.add_argument(
+        "--dataset_dir",
+        type=str,
+        required=True,
+        help="Directory containing the dataset",
+    )
+    parser.add_argument(
+        "--checkpoint_dir",
+        type=str,
+        required=True,
+        help="Directory where checkpoints will be saved",
+    )
+    parser.add_argument(
+        "--checkpoint_filename",
+        type=str,
+        default="finetune.ckpt",
+        help="Filename for the saved checkpoint",
+    )
+    parser.add_argument(
+        "--batch_size", type=int, default=6, help="Batch size for training"
+    )
+    parser.add_argument(
+        "--max_epochs", type=int, default=15, help="Maximum number of epochs to train"
+    )
+    parser.add_argument(
+        "--num_workers", type=int, default=5, help="Number of workers for data loading"
+    )
+    parser.add_argument(
+        "--accumulate_grad_batches",
+        type=int,
+        default=1,
+        help="Number of batches to accumulate gradients",
+    )
+    parser.add_argument(
+        "--num_gpus", type=int, default=4, help="Number of GPUs to use for training"
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=5e-5,
+        help="Learning rate for the optimizer",
+    )
+    parser.add_argument(
+        "--warmup_fraction",
+        type=float,
+        default=0.1,
+        help="Fraction of total steps to use for warmup",
+    )
+    parser.add_argument(
+        "--save_every_n_steps",
+        type=int,
+        default=512,
+        help="Save checkpoint every N steps",
+    )
+    parser.add_argument(
+        "--seed", type=int, default=123, help="Random seed for reproducibility"
+    )
+    parser.add_argument("--debug", action="store_true", help="Enable debug mode")
+    args = parser.parse_args()
+    main(args)

scripts/notebooks/CodonTransformerDemo.ipynb

@@ -0,0 +1,162 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import warnings\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "import pandas as pd\n",
+    "import torch\n",
+    "from transformers import AutoTokenizer, BigBirdForMaskedLM\n",
+    "\n",
+    "from CodonTransformer.CodonJupyter import (\n",
+    "    UserContainer,\n",
+    "    display_organism_dropdown,\n",
+    "    display_protein_input,\n",
+    "    format_model_output,\n",
+    ")\n",
+    "from CodonTransformer.CodonPrediction import predict_dna_sequence\n",
+    "\n",
+    "warnings.filterwarnings(\"ignore\")\n",
+    "\n",
+    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load model and tokenizer\n",
+    "tokenizer = AutoTokenizer.from_pretrained(\"adibvafa/CodonTransformer\")\n",
+    "model = BigBirdForMaskedLM.from_pretrained(\"adibvafa/CodonTransformer\").to(DEVICE)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "**Optimizing a Single Sequence**\n",
+    "-------------------------------------\n",
+    "1. Run the next code cell and input only your protein sequence and organism\n",
+    "\n",
+    "2. Run the code cell after it to optimize the sequence and display it.\n",
+    "\n",
+    "Protein sequences should end with \"*\" or \"_\" or an amino acid."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Sample: MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKTRREAEDLQVGQVELGG, Homo sapiens\n",
+    "user = UserContainer()\n",
+    "display_protein_input(user)\n",
+    "display_organism_dropdown(user)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "output = predict_dna_sequence(\n",
+    "    protein=user.protein,\n",
+    "    organism=user.organism,\n",
+    "    device=DEVICE,\n",
+    "    tokenizer=tokenizer,\n",
+    "    model=model,\n",
+    "    attention_type=\"original_full\",\n",
+    "    deterministic=True,\n",
+    "    # Can set temperature for non deterministic prediction\n",
+    ")\n",
+    "\n",
+    "print(format_model_output(output))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "**Optimizing Multiple Sequences**\n",
+    "-------------------------------------\n",
+    "1. Create a CSV file that has columns 'protein_sequence' and 'organism'.\n",
+    "   You can have other columns in any order.\n",
+    "\n",
+    "2. Replace the _dataset_path_ below with the actual path to your CSV file.\n",
+    "\n",
+    "3. Run the next code cells to optimize and save the predicted DNA sequences."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Update with the actual path to your dataset\n",
+    "dataset_path = \"scripts/demo/sample_dataset.csv\"\n",
+    "output_path = \"scripts/demo/sample_predictions.csv\"\n",
+    "\n",
+    "dataset = pd.read_csv(dataset_path, index_col=0)\n",
+    "dataset[\"predicted_dna\"] = None\n",
+    "dataset.head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for index, data in tqdm(\n",
+    "    dataset.iterrows(),\n",
+    "    desc=f\"CodonTransformer Predicting\",\n",
+    "    unit=\" Sequences\",\n",
+    "    total=dataset.shape[0],\n",
+    "):\n",
+    "\n",
+    "    outputs = predict_dna_sequence(\n",
+    "        protein=data[\"protein_sequence\"],\n",
+    "        organism=data[\"organism\"],\n",
+    "        device=DEVICE,\n",
+    "        tokenizer_object=tokenizer,\n",
+    "        model_object=model,\n",
+    "    )\n",
+    "    dataset.loc[index, \"predicted_dna\"] = outputs.predicted_dna\n",
+    "\n",
+    "dataset.to_csv(output_path)\n",
+    "dataset.head()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "light",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

scripts/pretrain.py

@@ -0,0 +1,283 @@
+"""
+File: pretrain.py
+-------------------
+Pretrain the CodonTransformer model.
+
+The dataset is a JSON file. You can use prepare_training_data from CodonData to
+prepare the dataset. The repository README has a guide on how to prepare the
+dataset and use this script.
+"""
+
+import argparse
+import gzip
+import math
+import os
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+MODEL_DIR = PROJECT_ROOT / "model"
+if str(MODEL_DIR) not in sys.path:
+    sys.path.insert(0, str(MODEL_DIR))
+
+import pytorch_lightning as pl
+import torch
+from torch.utils.data import DataLoader
+from transformers import BigBirdConfig, BigBirdForMaskedLM, PreTrainedTokenizerFast
+
+from CodonTransformer.CodonUtils import (
+    MAX_LEN,
+    NUM_ORGANISMS,
+    TOKEN2MASK,
+    IterableJSONData,
+)
+
+
+class MaskedTokenizerCollator:
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    def __call__(self, examples):
+        tokenized = self.tokenizer(
+            [ex["codons"] for ex in examples],
+            return_attention_mask=True,
+            return_token_type_ids=True,
+            truncation=True,
+            padding=True,
+            max_length=MAX_LEN,
+            return_tensors="pt",
+        )
+
+        seq_len = tokenized["input_ids"].shape[-1]
+        species_index = torch.tensor([[ex["organism"]] for ex in examples])
+        tokenized["token_type_ids"] = species_index.repeat(1, seq_len)
+
+        inputs = tokenized["input_ids"]
+        targets = inputs.clone()
+
+        prob_matrix = torch.full(inputs.shape, 0.15)
+        prob_matrix[inputs < 5] = 0.0
+        selected = torch.bernoulli(prob_matrix).bool()
+
+        # 80% of the time, replace masked input tokens with respective mask tokens
+        replaced = torch.bernoulli(torch.full(selected.shape, 0.8)).bool() & selected
+        inputs[replaced] = torch.tensor(
+            list((map(TOKEN2MASK.__getitem__, inputs[replaced].numpy())))
+        )
+
+        # 10% of the time, we replace masked input tokens with random vector.
+        randomized = (
+            torch.bernoulli(torch.full(selected.shape, 0.1)).bool()
+            & selected
+            & ~replaced
+        )
+        random_idx = torch.randint(26, 90, inputs.shape, dtype=torch.long)
+        inputs[randomized] = random_idx[randomized]
+
+        tokenized["input_ids"] = inputs
+        tokenized["labels"] = torch.where(selected, targets, -100)
+
+        return tokenized
+
+
+class plTrainHarness(pl.LightningModule):
+    def __init__(self, model, learning_rate, warmup_fraction, total_training_steps):
+        super().__init__()
+        self.model = model
+        self.learning_rate = learning_rate
+        self.warmup_fraction = warmup_fraction
+        self.total_training_steps = total_training_steps
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(
+            self.model.parameters(),
+            lr=self.learning_rate,
+        )
+        total_steps = self.total_training_steps or self.trainer.estimated_stepping_batches
+        if total_steps <= 0:
+            raise ValueError(f"Expected positive integer total_steps, but got {total_steps}")
+        lr_scheduler = {
+            "scheduler": torch.optim.lr_scheduler.OneCycleLR(
+                optimizer,
+                max_lr=self.learning_rate,
+                total_steps=total_steps,
+                pct_start=self.warmup_fraction,
+            ),
+            "interval": "step",
+            "frequency": 1,
+        }
+        return [optimizer], [lr_scheduler]
+
+    def training_step(self, batch, batch_idx):
+        self.model.bert.set_attention_type("block_sparse")
+        outputs = self.model(**batch)
+        self.log_dict(
+            dictionary={
+                "loss": outputs.loss,
+                "lr": self.trainer.optimizers[0].param_groups[0]["lr"],
+            },
+            on_step=True,
+            prog_bar=True,
+        )
+        return outputs.loss
+
+
+class EpochCheckpoint(pl.Callback):
+    def __init__(self, checkpoint_dir, save_interval):
+        super().__init__()
+        self.checkpoint_dir = checkpoint_dir
+        self.save_interval = save_interval
+
+    def on_train_epoch_end(self, trainer, pl_module):
+        current_epoch = trainer.current_epoch
+        if current_epoch % self.save_interval == 0 or current_epoch == 0:
+            checkpoint_path = os.path.join(
+                self.checkpoint_dir, f"epoch_{current_epoch}.ckpt"
+            )
+            trainer.save_checkpoint(checkpoint_path)
+            print(f"\nCheckpoint saved at {checkpoint_path}\n")
+
+
+def count_jsonl_records(path):
+    open_fn = gzip.open if path.endswith(".gz") else open
+    with open_fn(path, "rt") as file:
+        return sum(1 for line in file if line.strip())
+
+
+def estimate_training_steps(args):
+    num_records = count_jsonl_records(args.train_data_path)
+    num_devices = 1 if args.debug else args.num_gpus
+    samples_per_step = max(1, args.batch_size * num_devices)
+    batches_per_epoch = math.ceil(num_records / samples_per_step)
+    optimizer_steps_per_epoch = math.ceil(
+        batches_per_epoch / max(1, args.accumulate_grad_batches)
+    )
+    total_steps = max(1, optimizer_steps_per_epoch * args.max_epochs)
+    print(
+        "Estimated training steps: "
+        f"{total_steps} "
+        f"({num_records} records, batch_size={args.batch_size}, "
+        f"devices={num_devices}, max_epochs={args.max_epochs}, "
+        f"accumulate_grad_batches={args.accumulate_grad_batches})"
+    )
+    return total_steps
+
+
+def main(args):
+    """Pretrain the CodonTransformer model."""
+    pl.seed_everything(args.seed)
+    torch.set_float32_matmul_precision("medium")
+    total_training_steps = estimate_training_steps(args)
+
+    # Load the tokenizer and model
+    tokenizer = PreTrainedTokenizerFast(
+        tokenizer_file=args.tokenizer_path,
+        bos_token="[CLS]",
+        eos_token="[SEP]",
+        unk_token="[UNK]",
+        sep_token="[SEP]",
+        pad_token="[PAD]",
+        cls_token="[CLS]",
+        mask_token="[MASK]",
+    )
+    config = BigBirdConfig(
+        vocab_size=len(tokenizer),
+        type_vocab_size=NUM_ORGANISMS,
+        sep_token_id=2,
+    )
+    model = BigBirdForMaskedLM(config=config)
+    harnessed_model = plTrainHarness(
+        model,
+        args.learning_rate,
+        args.warmup_fraction,
+        total_training_steps,
+    )
+
+    # Load the training data
+    train_data = IterableJSONData(args.train_data_path, dist_env="slurm")
+    data_loader = DataLoader(
+        dataset=train_data,
+        collate_fn=MaskedTokenizerCollator(tokenizer),
+        batch_size=args.batch_size,
+        num_workers=0 if args.debug else args.num_workers,
+        persistent_workers=False if args.debug else True,
+    )
+
+    # Setup trainer and callbacks
+    save_checkpoint = EpochCheckpoint(args.checkpoint_dir, args.save_interval)
+    trainer = pl.Trainer(
+        default_root_dir=args.checkpoint_dir,
+        strategy="ddp_find_unused_parameters_true",
+        accelerator="gpu",
+        devices=1 if args.debug else args.num_gpus,
+        precision="16-mixed",
+        max_epochs=args.max_epochs,
+        deterministic=False,
+        enable_checkpointing=True,
+        callbacks=[save_checkpoint],
+        accumulate_grad_batches=args.accumulate_grad_batches,
+    )
+
+    # Pretrain the model
+    trainer.fit(harnessed_model, data_loader)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Pretrain the CodonTransformer model.")
+    parser.add_argument(
+        "--tokenizer_path",
+        type=str,
+        required=True,
+        help="Path to the tokenizer model file",
+    )
+    parser.add_argument(
+        "--train_data_path",
+        type=str,
+        required=True,
+        help="Path to the training data JSON file",
+    )
+    parser.add_argument(
+        "--checkpoint_dir",
+        type=str,
+        required=True,
+        help="Directory where checkpoints will be saved",
+    )
+    parser.add_argument(
+        "--batch_size", type=int, default=6, help="Batch size for training"
+    )
+    parser.add_argument(
+        "--max_epochs", type=int, default=5, help="Maximum number of epochs to train"
+    )
+    parser.add_argument(
+        "--num_workers", type=int, default=5, help="Number of workers for data loading"
+    )
+    parser.add_argument(
+        "--accumulate_grad_batches",
+        type=int,
+        default=1,
+        help="Number of batches to accumulate gradients",
+    )
+    parser.add_argument(
+        "--num_gpus", type=int, default=16, help="Number of GPUs to use for training"
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=5e-5,
+        help="Learning rate for the optimizer",
+    )
+    parser.add_argument(
+        "--warmup_fraction",
+        type=float,
+        default=0.1,
+        help="Fraction of total steps to use for warmup",
+    )
+    parser.add_argument(
+        "--save_interval", type=int, default=5, help="Save checkpoint every N epochs"
+    )
+    parser.add_argument(
+        "--seed", type=int, default=123, help="Random seed for reproducibility"
+    )
+    parser.add_argument("--debug", action="store_true", help="Enable debug mode")
+    args = parser.parse_args()
+    main(args)

scripts/slurm/export_finetuned_model.sh

@@ -0,0 +1,61 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Export a Lightning .ckpt checkpoint to a plain .pt model file for inference.
+# Run this after scripts/slurm/run_finetune.sh.
+# This script does not request SLURM resources.
+
+PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"
+CONDA_ENV="${CONDA_ENV:-struct-evo}"
+
+CHECKPOINT_DIR="${CHECKPOINT_DIR:-${PROJECT_DIR}/weight/checkpoints/finetune}"
+CHECKPOINT_PATH="${CHECKPOINT_PATH:-}"
+OUTPUT_MODEL_PATH="${OUTPUT_MODEL_PATH:-${CHECKPOINT_DIR}/finetuned_model.pt}"
+NUM_ORGANISMS="${NUM_ORGANISMS:-164}"
+
+cd "${PROJECT_DIR}"
+
+export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
+export CHECKPOINT_DIR
+export CHECKPOINT_PATH
+export OUTPUT_MODEL_PATH
+export NUM_ORGANISMS
+export PYTHONFAULTHANDLER=1
+
+if [[ -n "${CONDA_ENV}" ]] && command -v conda >/dev/null 2>&1; then
+    # shellcheck disable=SC1091
+    source "$(conda info --base)/etc/profile.d/conda.sh"
+    conda activate "${CONDA_ENV}"
+fi
+
+if [[ -z "${CHECKPOINT_PATH}" ]]; then
+    CHECKPOINT_PATH="$(find "${CHECKPOINT_DIR}" -maxdepth 1 -type f -name 'epoch=*-step=*.ckpt' | sort | tail -n 1)"
+fi
+
+if [[ -z "${CHECKPOINT_PATH}" ]] || [[ ! -f "${CHECKPOINT_PATH}" ]]; then
+    echo "Missing CHECKPOINT_PATH. Set it explicitly or check CHECKPOINT_DIR: ${CHECKPOINT_DIR}" >&2
+    exit 1
+fi
+
+mkdir -p "$(dirname "${OUTPUT_MODEL_PATH}")"
+export CHECKPOINT_PATH
+
+python - <<'PY'
+import os
+
+from CodonTransformer.CodonPrediction import create_model_from_checkpoint
+
+checkpoint_path = os.environ["CHECKPOINT_PATH"]
+output_model_path = os.environ["OUTPUT_MODEL_PATH"]
+num_organisms = int(os.environ["NUM_ORGANISMS"])
+
+print(f"Checkpoint: {checkpoint_path}")
+print(f"Output model: {output_model_path}")
+print(f"Num organisms: {num_organisms}")
+
+create_model_from_checkpoint(
+    checkpoint_dir=checkpoint_path,
+    output_model_dir=output_model_path,
+    num_organisms=num_organisms,
+)
+PY

scripts/slurm/finetune.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+#SBATCH --job-name=finetune
+#SBATCH --output=your_output_directory/output_%j.out
+#SBATCH --error=your_error_directory/error_%j.err
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=4
+#SBATCH --gpus-per-node=4
+#SBATCH --cpus-per-task=6
+#SBATCH --time=15:00:00
+#SBATCH --partition=compute_full_node
+
+# Load required modules
+module --ignore_cache load cuda/11.4.4
+module --ignore_cache load anaconda3
+source activate your_environment
+
+# Change to the working directory
+cd your_working_directory
+
+# Set environment variables
+export CUBLAS_WORKSPACE_CONFIG=:4096:2
+export PYTHONPATH="${PWD}/model:${PYTHONPATH:-}"
+export NCCL_DEBUG=INFO
+export PYTHONFAULTHANDLER=1
+
+# Run the Python script with arguments
+stdbuf -oL -eL srun python scripts/finetune.py \
+    --dataset_dir your_dataset_directory \
+    --checkpoint_dir your_checkpoint_directory \
+    --checkpoint_filename finetune.ckpt \
+    --batch_size 6 \
+    --max_epochs 15 \
+    --num_workers 5 \
+    --accumulate_grad_batches 1 \
+    --num_gpus 4 \
+    --learning_rate 0.00005 \
+    --warmup_fraction 0.1 \
+    --save_every_n_steps 512 \
+    --seed 123

scripts/slurm/prepare_finetune_data.sh

@@ -0,0 +1,45 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Convert a user-provided finetuning CSV into the JSONL format required by finetuning.
+# Run this before scripts/slurm/run_finetune.sh.
+# This script does not request SLURM resources.
+
+PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"
+CONDA_ENV="${CONDA_ENV:-struct-evo}"
+
+INPUT_CSV="${INPUT_CSV:-${PROJECT_DIR}/scripts/data/raw/your_data.csv}"
+OUTPUT_JSON="${OUTPUT_JSON:-${PROJECT_DIR}/scripts/data/processed/finetune_data.json}"
+
+cd "${PROJECT_DIR}"
+mkdir -p "$(dirname "${OUTPUT_JSON}")"
+
+export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
+export INPUT_CSV
+export OUTPUT_JSON
+export PYTHONFAULTHANDLER=1
+
+if [[ -n "${CONDA_ENV}" ]] && command -v conda >/dev/null 2>&1; then
+    # shellcheck disable=SC1091
+    source "$(conda info --base)/etc/profile.d/conda.sh"
+    conda activate "${CONDA_ENV}"
+fi
+
+if [[ ! -f "${INPUT_CSV}" ]]; then
+    echo "Missing INPUT_CSV: ${INPUT_CSV}" >&2
+    exit 1
+fi
+
+python - <<'PY'
+import os
+
+from CodonTransformer.CodonData import prepare_training_data
+
+input_csv = os.environ["INPUT_CSV"]
+output_json = os.environ["OUTPUT_JSON"]
+
+print(f"Input CSV: {input_csv}")
+print(f"Output JSONL: {output_json}")
+
+prepare_training_data(input_csv, output_json)
+PY

scripts/slurm/prepare_pretrain_data.sh

@@ -0,0 +1,45 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Convert the full CSV dataset into the JSONL format required by pretraining.
+# Run this before scripts/slurm/run_pretrain.sh.
+# This script does not request SLURM resources.
+
+PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"
+CONDA_ENV="${CONDA_ENV:-struct-evo}"
+
+INPUT_CSV="${INPUT_CSV:-${PROJECT_DIR}/scripts/data/raw/dataset.csv}"
+OUTPUT_JSON="${OUTPUT_JSON:-${PROJECT_DIR}/scripts/data/processed/pretrain_data.json}"
+
+cd "${PROJECT_DIR}"
+mkdir -p "$(dirname "${OUTPUT_JSON}")"
+
+export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
+export INPUT_CSV
+export OUTPUT_JSON
+export PYTHONFAULTHANDLER=1
+
+if [[ -n "${CONDA_ENV}" ]] && command -v conda >/dev/null 2>&1; then
+    # shellcheck disable=SC1091
+    source "$(conda info --base)/etc/profile.d/conda.sh"
+    conda activate "${CONDA_ENV}"
+fi
+
+if [[ ! -f "${INPUT_CSV}" ]]; then
+    echo "Missing INPUT_CSV: ${INPUT_CSV}" >&2
+    exit 1
+fi
+
+python - <<'PY'
+import os
+
+from CodonTransformer.CodonData import prepare_training_data
+
+input_csv = os.environ["INPUT_CSV"]
+output_json = os.environ["OUTPUT_JSON"]
+
+print(f"Input CSV: {input_csv}")
+print(f"Output JSONL: {output_json}")
+
+prepare_training_data(input_csv, output_json)
+PY

scripts/slurm/pretrain.sh

@@ -0,0 +1,38 @@
+#!/bin/bash
+#SBATCH --job-name=pretrain
+#SBATCH --output=your_output_directory/output_%j.out
+#SBATCH --error=your_error_directory/error_%j.err
+#SBATCH --nodes=4
+#SBATCH --gpus-per-node=4
+#SBATCH --ntasks-per-node=4
+#SBATCH --time=23:59:00
+#SBATCH -p compute_full_node
+
+# Load required modules
+module --ignore_cache load cuda/11.4.4
+module --ignore_cache load anaconda3
+source activate your_environment
+
+# Change to the working directory
+cd your_working_directory
+
+# Set environment variables
+export CUBLAS_WORKSPACE_CONFIG=:4096:2
+export PYTHONPATH="${PWD}/model:${PYTHONPATH:-}"
+export NCCL_DEBUG=INFO
+export PYTHONFAULTHANDLER=1
+
+# Run the Python script with arguments
+stdbuf -oL -eL srun python scripts/pretrain.py \
+    --tokenizer_path model/src/CodonTransformerTokenizer.json \
+    --train_data_path your_data_directory/pretrain_dataset.json \
+    --checkpoint_dir your_checkpoint_directory \
+    --batch_size 6 \
+    --max_epochs 5 \
+    --num_workers 5 \
+    --accumulate_grad_batches 1 \
+    --num_gpus 16 \
+    --learning_rate 0.00005 \
+    --warmup_fraction 0.1 \
+    --save_interval 5 \
+    --seed 123

scripts/slurm/run_finetune.sh

@@ -0,0 +1,73 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Run this script inside an allocated/interactive GPU session.
+# It does not request SLURM resources.
+
+PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"
+HF_HOME="${HF_HOME:-/public/home/scnb9biwet/.cache/huggingface}"
+CONDA_ENV="${CONDA_ENV:-struct-evo}"
+
+DATASET_JSON="${DATASET_JSON:-${PROJECT_DIR}/scripts/data/processed/finetune_data.json}"
+CHECKPOINT_DIR="${CHECKPOINT_DIR:-${PROJECT_DIR}/weight/checkpoints/finetune}"
+CHECKPOINT_FILENAME="${CHECKPOINT_FILENAME:-finetune.ckpt}"
+
+BATCH_SIZE="${BATCH_SIZE:-6}"
+MAX_EPOCHS="${MAX_EPOCHS:-15}"
+NUM_WORKERS="${NUM_WORKERS:-5}"
+ACCUMULATE_GRAD_BATCHES="${ACCUMULATE_GRAD_BATCHES:-1}"
+NUM_GPUS="${NUM_GPUS:-4}"
+LEARNING_RATE="${LEARNING_RATE:-0.00005}"
+WARMUP_FRACTION="${WARMUP_FRACTION:-0.1}"
+SAVE_EVERY_N_STEPS="${SAVE_EVERY_N_STEPS:-512}"
+SEED="${SEED:-123}"
+DEBUG="${DEBUG:-0}"
+OFFLINE="${OFFLINE:-1}"
+
+cd "${PROJECT_DIR}"
+mkdir -p "${CHECKPOINT_DIR}"
+
+export HF_HOME
+export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
+export CUBLAS_WORKSPACE_CONFIG=:4096:2
+export NCCL_DEBUG="${NCCL_DEBUG:-INFO}"
+export PYTHONFAULTHANDLER=1
+
+if [[ "${OFFLINE}" == "1" ]]; then
+    export HF_HUB_OFFLINE=1
+    export TRANSFORMERS_OFFLINE=1
+fi
+
+if [[ -n "${CONDA_ENV}" ]] && command -v conda >/dev/null 2>&1; then
+    # shellcheck disable=SC1091
+    source "$(conda info --base)/etc/profile.d/conda.sh"
+    conda activate "${CONDA_ENV}"
+fi
+
+if [[ ! -f "${DATASET_JSON}" ]]; then
+    echo "Missing DATASET_JSON: ${DATASET_JSON}" >&2
+    echo "Create it first with CodonTransformer.CodonData.prepare_training_data." >&2
+    exit 1
+fi
+
+args=(
+    "${PROJECT_DIR}/scripts/finetune.py"
+    --dataset_dir "${DATASET_JSON}"
+    --checkpoint_dir "${CHECKPOINT_DIR}"
+    --checkpoint_filename "${CHECKPOINT_FILENAME}"
+    --batch_size "${BATCH_SIZE}"
+    --max_epochs "${MAX_EPOCHS}"
+    --num_workers "${NUM_WORKERS}"
+    --accumulate_grad_batches "${ACCUMULATE_GRAD_BATCHES}"
+    --num_gpus "${NUM_GPUS}"
+    --learning_rate "${LEARNING_RATE}"
+    --warmup_fraction "${WARMUP_FRACTION}"
+    --save_every_n_steps "${SAVE_EVERY_N_STEPS}"
+    --seed "${SEED}"
+)
+
+if [[ "${DEBUG}" == "1" ]]; then
+    args+=(--debug)
+fi
+
+python "${args[@]}"

scripts/slurm/run_inference_batch.sh

@@ -0,0 +1,89 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Batch deterministic inference from a CSV file.
+# Run this inside an allocated/interactive GPU session. No SLURM resources are requested here.
+
+PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"
+HF_HOME="${HF_HOME:-/public/home/scnb9biwet/.cache/huggingface}"
+CONDA_ENV="${CONDA_ENV:-struct-evo}"
+
+INPUT_CSV="${INPUT_CSV:-${PROJECT_DIR}/scripts/demo/sample_dataset.csv}"
+OUTPUT_CSV="${OUTPUT_CSV:-${PROJECT_DIR}/outputs/sample_predictions.csv}"
+OFFLINE="${OFFLINE:-1}"
+
+cd "${PROJECT_DIR}"
+mkdir -p "$(dirname "${OUTPUT_CSV}")"
+
+export HF_HOME
+export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
+export INPUT_CSV
+export OUTPUT_CSV
+export OFFLINE
+export PYTHONFAULTHANDLER=1
+
+if [[ "${OFFLINE}" == "1" ]]; then
+    export HF_HUB_OFFLINE=1
+    export TRANSFORMERS_OFFLINE=1
+fi
+
+if [[ -n "${CONDA_ENV}" ]] && command -v conda >/dev/null 2>&1; then
+    # shellcheck disable=SC1091
+    source "$(conda info --base)/etc/profile.d/conda.sh"
+    conda activate "${CONDA_ENV}"
+fi
+
+python - <<'PY'
+import os
+
+import pandas as pd
+import torch
+from tqdm import tqdm
+from transformers import AutoTokenizer, BigBirdForMaskedLM
+
+from CodonTransformer.CodonPrediction import predict_dna_sequence
+
+input_csv = os.environ["INPUT_CSV"]
+output_csv = os.environ["OUTPUT_CSV"]
+local_files_only = os.environ.get("OFFLINE", "1") == "1"
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"HF_HOME: {os.environ.get('HF_HOME')}")
+print(f"Device: {device}")
+print(f"Local files only: {local_files_only}")
+print(f"Input CSV: {input_csv}")
+
+tokenizer = AutoTokenizer.from_pretrained(
+    "adibvafa/CodonTransformer",
+    local_files_only=local_files_only,
+)
+model = BigBirdForMaskedLM.from_pretrained(
+    "adibvafa/CodonTransformer",
+    local_files_only=local_files_only,
+).to(device)
+
+dataset = pd.read_csv(input_csv)
+if "Unnamed: 0" in dataset.columns:
+    dataset = dataset.drop(columns=["Unnamed: 0"])
+
+required_columns = {"protein_sequence", "organism"}
+missing = required_columns - set(dataset.columns)
+if missing:
+    raise ValueError(f"Input CSV is missing required columns: {sorted(missing)}")
+
+dataset["predicted_dna"] = ""
+for index, row in tqdm(dataset.iterrows(), total=len(dataset), desc="Predicting"):
+    output = predict_dna_sequence(
+        protein=row["protein_sequence"],
+        organism=row["organism"],
+        device=device,
+        tokenizer=tokenizer,
+        model=model,
+        attention_type="original_full",
+        deterministic=True,
+    )
+    dataset.loc[index, "predicted_dna"] = output.predicted_dna
+
+dataset.to_csv(output_csv, index=False)
+print(f"Saved predictions to {output_csv}")
+PY

scripts/slurm/run_inference_finetuned.sh

@@ -0,0 +1,70 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Single deterministic inference with an exported finetuned .pt model.
+# Run this after scripts/slurm/export_finetuned_model.sh.
+# This script does not request SLURM resources.
+
+PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"
+CONDA_ENV="${CONDA_ENV:-struct-evo}"
+
+TOKENIZER_PATH="${TOKENIZER_PATH:-${PROJECT_DIR}/model/src/CodonTransformerTokenizer.json}"
+MODEL_PATH="${MODEL_PATH:-${PROJECT_DIR}/weight/checkpoints/finetune/finetuned_model.pt}"
+PROTEIN="${PROTEIN:-MFWY}"
+ORGANISM="${ORGANISM:-Escherichia coli general}"
+MATCH_PROTEIN="${MATCH_PROTEIN:-1}"
+
+cd "${PROJECT_DIR}"
+
+export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
+export TOKENIZER_PATH
+export MODEL_PATH
+export PROTEIN
+export ORGANISM
+export MATCH_PROTEIN
+export PYTHONFAULTHANDLER=1
+
+if [[ -n "${CONDA_ENV}" ]] && command -v conda >/dev/null 2>&1; then
+    # shellcheck disable=SC1091
+    source "$(conda info --base)/etc/profile.d/conda.sh"
+    conda activate "${CONDA_ENV}"
+fi
+
+if [[ ! -f "${TOKENIZER_PATH}" ]]; then
+    echo "Missing TOKENIZER_PATH: ${TOKENIZER_PATH}" >&2
+    exit 1
+fi
+
+if [[ ! -f "${MODEL_PATH}" ]]; then
+    echo "Missing MODEL_PATH: ${MODEL_PATH}" >&2
+    exit 1
+fi
+
+python - <<'PY'
+import os
+
+import torch
+
+from CodonTransformer.CodonJupyter import format_model_output
+from CodonTransformer.CodonPrediction import predict_dna_sequence
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+match_protein = os.environ.get("MATCH_PROTEIN", "1") == "1"
+
+print(f"Device: {device}")
+print(f"Tokenizer: {os.environ['TOKENIZER_PATH']}")
+print(f"Model: {os.environ['MODEL_PATH']}")
+
+output = predict_dna_sequence(
+    protein=os.environ["PROTEIN"],
+    organism=os.environ["ORGANISM"],
+    device=device,
+    tokenizer=os.environ["TOKENIZER_PATH"],
+    model=os.environ["MODEL_PATH"],
+    attention_type="original_full",
+    deterministic=True,
+    match_protein=match_protein,
+)
+
+print(format_model_output(output))
+PY

scripts/slurm/run_inference_multiple.sh

@@ -0,0 +1,113 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Multiple variable inference for one protein sequence.
+# Run this inside an allocated/interactive GPU session. No SLURM resources are requested here.
+
+PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"
+HF_HOME="${HF_HOME:-/public/home/scnb9biwet/.cache/huggingface}"
+CONDA_ENV="${CONDA_ENV:-struct-evo}"
+
+PROTEIN="${PROTEIN:-MFWY}"
+ORGANISM="${ORGANISM:-Escherichia coli general}"
+OUTPUT_CSV="${OUTPUT_CSV:-${PROJECT_DIR}/outputs/multiple_predictions.csv}"
+NUM_SEQUENCES="${NUM_SEQUENCES:-5}"
+TEMPERATURE="${TEMPERATURE:-0.5}"
+TOP_P="${TOP_P:-0.95}"
+MATCH_PROTEIN="${MATCH_PROTEIN:-1}"
+OFFLINE="${OFFLINE:-1}"
+
+cd "${PROJECT_DIR}"
+mkdir -p "$(dirname "${OUTPUT_CSV}")"
+
+export HF_HOME
+export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
+export PROTEIN
+export ORGANISM
+export OUTPUT_CSV
+export NUM_SEQUENCES
+export TEMPERATURE
+export TOP_P
+export MATCH_PROTEIN
+export OFFLINE
+export PYTHONFAULTHANDLER=1
+
+if [[ "${OFFLINE}" == "1" ]]; then
+    export HF_HUB_OFFLINE=1
+    export TRANSFORMERS_OFFLINE=1
+fi
+
+if [[ -n "${CONDA_ENV}" ]] && command -v conda >/dev/null 2>&1; then
+    # shellcheck disable=SC1091
+    source "$(conda info --base)/etc/profile.d/conda.sh"
+    conda activate "${CONDA_ENV}"
+fi
+
+python - <<'PY'
+import os
+
+import pandas as pd
+import torch
+from transformers import AutoTokenizer, BigBirdForMaskedLM
+
+from CodonTransformer.CodonJupyter import format_model_output
+from CodonTransformer.CodonPrediction import predict_dna_sequence
+
+protein = os.environ["PROTEIN"]
+organism = os.environ["ORGANISM"]
+output_csv = os.environ["OUTPUT_CSV"]
+num_sequences = int(os.environ["NUM_SEQUENCES"])
+temperature = float(os.environ["TEMPERATURE"])
+top_p = float(os.environ["TOP_P"])
+match_protein = os.environ.get("MATCH_PROTEIN", "1") == "1"
+local_files_only = os.environ.get("OFFLINE", "1") == "1"
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"HF_HOME: {os.environ.get('HF_HOME')}")
+print(f"Device: {device}")
+print(f"Local files only: {local_files_only}")
+print(
+    "Sampling: "
+    f"num_sequences={num_sequences}, temperature={temperature}, "
+    f"top_p={top_p}, match_protein={match_protein}"
+)
+
+tokenizer = AutoTokenizer.from_pretrained(
+    "adibvafa/CodonTransformer",
+    local_files_only=local_files_only,
+)
+model = BigBirdForMaskedLM.from_pretrained(
+    "adibvafa/CodonTransformer",
+    local_files_only=local_files_only,
+).to(device)
+
+outputs = predict_dna_sequence(
+    protein=protein,
+    organism=organism,
+    device=device,
+    tokenizer=tokenizer,
+    model=model,
+    attention_type="original_full",
+    deterministic=False,
+    temperature=temperature,
+    top_p=top_p,
+    num_sequences=num_sequences,
+    match_protein=match_protein,
+)
+
+rows = []
+for index, output in enumerate(outputs, start=1):
+    print(f"\nSequence {index}")
+    print(format_model_output(output))
+    rows.append(
+        {
+            "sequence_index": index,
+            "protein_sequence": output.protein,
+            "organism": output.organism,
+            "predicted_dna": output.predicted_dna,
+        }
+    )
+
+pd.DataFrame(rows).to_csv(output_csv, index=False)
+print(f"Saved predictions to {output_csv}")
+PY

scripts/slurm/run_inference_single.sh

@@ -0,0 +1,73 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Single deterministic inference for one protein sequence.
+# Run this inside an allocated/interactive GPU session. No SLURM resources are requested here.
+
+PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"
+HF_HOME="${HF_HOME:-/public/home/scnb9biwet/.cache/huggingface}"
+CONDA_ENV="${CONDA_ENV:-struct-evo}"
+
+PROTEIN="${PROTEIN:-MFWY}"
+ORGANISM="${ORGANISM:-Escherichia coli general}"
+OFFLINE="${OFFLINE:-1}"
+
+cd "${PROJECT_DIR}"
+
+export HF_HOME
+export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
+export PROTEIN
+export ORGANISM
+export OFFLINE
+export PYTHONFAULTHANDLER=1
+
+if [[ "${OFFLINE}" == "1" ]]; then
+    export HF_HUB_OFFLINE=1
+    export TRANSFORMERS_OFFLINE=1
+fi
+
+if [[ -n "${CONDA_ENV}" ]] && command -v conda >/dev/null 2>&1; then
+    # shellcheck disable=SC1091
+    source "$(conda info --base)/etc/profile.d/conda.sh"
+    conda activate "${CONDA_ENV}"
+fi
+
+python - <<'PY'
+import os
+
+import torch
+from transformers import AutoTokenizer, BigBirdForMaskedLM
+
+from CodonTransformer.CodonJupyter import format_model_output
+from CodonTransformer.CodonPrediction import predict_dna_sequence
+
+protein = os.environ["PROTEIN"]
+organism = os.environ["ORGANISM"]
+local_files_only = os.environ.get("OFFLINE", "1") == "1"
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"HF_HOME: {os.environ.get('HF_HOME')}")
+print(f"Device: {device}")
+print(f"Local files only: {local_files_only}")
+
+tokenizer = AutoTokenizer.from_pretrained(
+    "adibvafa/CodonTransformer",
+    local_files_only=local_files_only,
+)
+model = BigBirdForMaskedLM.from_pretrained(
+    "adibvafa/CodonTransformer",
+    local_files_only=local_files_only,
+).to(device)
+
+output = predict_dna_sequence(
+    protein=protein,
+    organism=organism,
+    device=device,
+    tokenizer=tokenizer,
+    model=model,
+    attention_type="original_full",
+    deterministic=True,
+)
+
+print(format_model_output(output))
+PY

scripts/slurm/run_pretrain.sh

@@ -0,0 +1,78 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Run this script inside an allocated/interactive GPU session.
+# It does not request SLURM resources.
+
+PROJECT_DIR="${PROJECT_DIR:-/public/home/scnb9biwet/jiangqq/CodonTransformer-main}"
+HF_HOME="${HF_HOME:-/public/home/scnb9biwet/.cache/huggingface}"
+CONDA_ENV="${CONDA_ENV:-struct-evo}"
+
+TOKENIZER_PATH="${TOKENIZER_PATH:-${PROJECT_DIR}/model/src/CodonTransformerTokenizer.json}"
+TRAIN_DATA_PATH="${TRAIN_DATA_PATH:-${PROJECT_DIR}/scripts/data/processed/pretrain_data.json}"
+CHECKPOINT_DIR="${CHECKPOINT_DIR:-${PROJECT_DIR}/weight/checkpoints/pretrain}"
+
+BATCH_SIZE="${BATCH_SIZE:-6}"
+MAX_EPOCHS="${MAX_EPOCHS:-5}"
+NUM_WORKERS="${NUM_WORKERS:-5}"
+ACCUMULATE_GRAD_BATCHES="${ACCUMULATE_GRAD_BATCHES:-1}"
+NUM_GPUS="${NUM_GPUS:-16}"
+LEARNING_RATE="${LEARNING_RATE:-0.00005}"
+WARMUP_FRACTION="${WARMUP_FRACTION:-0.1}"
+SAVE_INTERVAL="${SAVE_INTERVAL:-5}"
+SEED="${SEED:-123}"
+DEBUG="${DEBUG:-0}"
+OFFLINE="${OFFLINE:-1}"
+
+cd "${PROJECT_DIR}"
+mkdir -p "${CHECKPOINT_DIR}"
+
+export HF_HOME
+export PYTHONPATH="${PROJECT_DIR}/model:${PYTHONPATH:-}"
+export CUBLAS_WORKSPACE_CONFIG=:4096:2
+export NCCL_DEBUG="${NCCL_DEBUG:-INFO}"
+export PYTHONFAULTHANDLER=1
+
+if [[ "${OFFLINE}" == "1" ]]; then
+    export HF_HUB_OFFLINE=1
+    export TRANSFORMERS_OFFLINE=1
+fi
+
+if [[ -n "${CONDA_ENV}" ]] && command -v conda >/dev/null 2>&1; then
+    # shellcheck disable=SC1091
+    source "$(conda info --base)/etc/profile.d/conda.sh"
+    conda activate "${CONDA_ENV}"
+fi
+
+if [[ ! -f "${TOKENIZER_PATH}" ]]; then
+    echo "Missing TOKENIZER_PATH: ${TOKENIZER_PATH}" >&2
+    exit 1
+fi
+
+if [[ ! -f "${TRAIN_DATA_PATH}" ]]; then
+    echo "Missing TRAIN_DATA_PATH: ${TRAIN_DATA_PATH}" >&2
+    echo "Create it first with CodonTransformer.CodonData.prepare_training_data." >&2
+    exit 1
+fi
+
+args=(
+    "${PROJECT_DIR}/scripts/pretrain.py"
+    --tokenizer_path "${TOKENIZER_PATH}"
+    --train_data_path "${TRAIN_DATA_PATH}"
+    --checkpoint_dir "${CHECKPOINT_DIR}"
+    --batch_size "${BATCH_SIZE}"
+    --max_epochs "${MAX_EPOCHS}"
+    --num_workers "${NUM_WORKERS}"
+    --accumulate_grad_batches "${ACCUMULATE_GRAD_BATCHES}"
+    --num_gpus "${NUM_GPUS}"
+    --learning_rate "${LEARNING_RATE}"
+    --warmup_fraction "${WARMUP_FRACTION}"
+    --save_interval "${SAVE_INTERVAL}"
+    --seed "${SEED}"
+)
+
+if [[ "${DEBUG}" == "1" ]]; then
+    args+=(--debug)
+fi
+
+python "${args[@]}"

scripts/tests/test_CodonData.py

@@ -0,0 +1,111 @@
+import tempfile
+import unittest
+
+import pandas as pd
+from Bio.Data.CodonTable import TranslationError
+
+from CodonTransformer.CodonData import (
+    build_amino2codon_skeleton,
+    get_amino_acid_sequence,
+    is_correct_seq,
+    preprocess_protein_sequence,
+    read_fasta_file,
+)
+from CodonTransformer.CodonUtils import ProteinConfig
+
+
+class TestCodonData(unittest.TestCase):
+    def test_preprocess_protein_sequence(self):
+        with ProteinConfig() as config:
+            config.set("ambiguous_aminoacid_behavior", "raise_error")
+            protein = "Z_"
+            try:
+                preprocess_protein_sequence(protein)
+                self.fail("Expected ValueError")
+            except ValueError:
+                pass
+            config.set("ambiguous_aminoacid_behavior", "standardize_deterministic")
+            for _ in range(10):
+                preprocessed_protein = preprocess_protein_sequence(protein)
+                self.assertEqual(preprocessed_protein, "Q_")
+            config.set("ambiguous_aminoacid_behavior", "standardize_random")
+            random_results = set()
+            # The probability of getting the same result 30 times in a row is
+            # 1 in 1.073741824*10^9 if there are only two possible results.
+            for _ in range(30):
+                preprocessed_protein = preprocess_protein_sequence(protein)
+                random_results.add(preprocessed_protein)
+            self.assertGreater(len(random_results), 1)
+
+    def test_read_fasta_file(self):
+        fasta_content = ">sequence1\n" "ATGATGATGATGATG\n" ">sequence2\n" "TGATGATGATGA"
+
+        with tempfile.NamedTemporaryFile(
+            mode="w", delete=False, suffix=".fasta"
+        ) as temp_file:
+            temp_file.write(fasta_content)
+            temp_file_name = temp_file.name
+
+        try:
+            sequences = read_fasta_file(temp_file_name, save_to_file=None)
+            self.assertIsInstance(sequences, pd.DataFrame)
+            self.assertEqual(len(sequences), 2)
+            self.assertEqual(sequences.iloc[0]["dna"], "ATGATGATGATGATG")
+            self.assertEqual(sequences.iloc[1]["dna"], "TGATGATGATGA")
+        finally:
+            import os
+
+            os.unlink(temp_file_name)
+
+    def test_build_amino2codon_skeleton(self):
+        organism = "Homo sapiens"
+        codon_skeleton = build_amino2codon_skeleton(organism)
+
+        expected_amino_acids = "ARNDCQEGHILKMFPSTWYV_"
+
+        for amino_acid in expected_amino_acids:
+            self.assertIn(amino_acid, codon_skeleton)
+            codons, frequencies = codon_skeleton[amino_acid]
+            self.assertIsInstance(codons, list)
+            self.assertIsInstance(frequencies, list)
+            self.assertEqual(len(codons), len(frequencies))
+            self.assertTrue(all(isinstance(codon, str) for codon in codons))
+            self.assertTrue(all(freq == 0 for freq in frequencies))
+
+        all_codons = set(
+            codon for codons, _ in codon_skeleton.values() for codon in codons
+        )
+        self.assertEqual(len(all_codons), 64)  # There should be 64 unique codons
+
+    def test_get_amino_acid_sequence(self):
+        dna = "ATGGCCTGA"
+        protein, is_correct = get_amino_acid_sequence(dna, return_correct_seq=True)
+        self.assertEqual(protein, "MA_")
+        self.assertTrue(is_correct)
+
+    def test_is_correct_seq(self):
+        dna = "ATGGCCTGA"
+        protein = "MA_"
+        self.assertTrue(is_correct_seq(dna, protein))
+
+    def test_read_fasta_file_raises_exception_for_non_dna(self):
+        non_dna_content = ">sequence1\nATGATGATGXYZATG\n>sequence2\nTGATGATGATGA"
+
+        with tempfile.NamedTemporaryFile(
+            mode="w", delete=False, suffix=".fasta"
+        ) as temp_file:
+            temp_file.write(non_dna_content)
+            temp_file_name = temp_file.name
+
+        try:
+            with self.assertRaises(TranslationError) as context:
+                read_fasta_file(temp_file_name)
+            self.assertIn("Codon 'XYZ' is invalid", str(context.exception))
+        finally:
+            import os
+
+            os.unlink(temp_file_name)
+
+
+if __name__ == "__main__":
+    unittest.main()

scripts/tests/test_CodonJupyter.py

@@ -0,0 +1,118 @@
+import unittest
+
+import ipywidgets
+
+from CodonTransformer.CodonJupyter import (
+    DNASequencePrediction,
+    UserContainer,
+    create_dropdown_options,
+    create_organism_dropdown,
+    display_organism_dropdown,
+    display_protein_input,
+    format_model_output,
+)
+from CodonTransformer.CodonUtils import ORGANISM2ID
+
+
+class TestCodonJupyter(unittest.TestCase):
+    def test_UserContainer(self):
+        user_container = UserContainer()
+        self.assertEqual(user_container.organism, -1)
+        self.assertEqual(user_container.protein, "")
+
+    def test_create_organism_dropdown(self):
+        container = UserContainer()
+        dropdown = create_organism_dropdown(container)
+
+        self.assertIsInstance(dropdown, ipywidgets.Dropdown)
+        self.assertGreater(len(dropdown.options), 0)
+        self.assertEqual(dropdown.description, "")
+        self.assertEqual(dropdown.layout.width, "40%")
+        self.assertEqual(dropdown.layout.margin, "0 0 10px 0")
+        self.assertEqual(dropdown.style.description_width, "initial")
+
+        # Test the dropdown options
+        options = dropdown.options
+        self.assertIn("", options)
+        self.assertIn("Selected Organisms", options)
+        self.assertIn("All Organisms", options)
+
+    def test_create_dropdown_options(self):
+        options = create_dropdown_options(ORGANISM2ID)
+        self.assertIsInstance(options, list)
+        self.assertGreater(len(options), 0)
+
+    def test_display_organism_dropdown(self):
+        container = UserContainer()
+        with unittest.mock.patch(
+            "CodonTransformer.CodonJupyter.display"
+        ) as mock_display:
+            display_organism_dropdown(container)
+
+        # Check that display was called twice (for container_widget and HTML)
+        self.assertEqual(mock_display.call_count, 2)
+
+        # Check that the first call to display was with a VBox widget
+        self.assertIsInstance(mock_display.call_args_list[0][0][0], ipywidgets.VBox)
+
+        # Check that the VBox contains a Dropdown
+        dropdown = mock_display.call_args_list[0][0][0].children[1]
+        self.assertIsInstance(dropdown, ipywidgets.Dropdown)
+        self.assertGreater(len(dropdown.options), 0)
+
+    def test_display_protein_input(self):
+        container = UserContainer()
+        with unittest.mock.patch(
+            "CodonTransformer.CodonJupyter.display"
+        ) as mock_display:
+            display_protein_input(container)
+
+        # Check that display was called twice (for container_widget and HTML)
+        self.assertEqual(mock_display.call_count, 2)
+
+        # Check that the first call to display was with a VBox widget
+        self.assertIsInstance(mock_display.call_args_list[0][0][0], ipywidgets.VBox)
+
+        # Check that the VBox contains a Textarea
+        textarea = mock_display.call_args_list[0][0][0].children[1]
+        self.assertIsInstance(textarea, ipywidgets.Textarea)
+
+        # Verify the properties of the Textarea
+        self.assertEqual(textarea.value, "")
+        self.assertEqual(textarea.placeholder, "Enter here...")
+        self.assertEqual(textarea.description, "")
+        self.assertEqual(textarea.layout.width, "100%")
+        self.assertEqual(textarea.layout.height, "100px")
+        self.assertEqual(textarea.layout.margin, "0 0 10px 0")
+        self.assertEqual(textarea.style.description_width, "initial")
+
+    def test_format_model_output(self):
+        output = DNASequencePrediction(
+            organism="Escherichia coli",
+            protein="MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG",
+            processed_input="MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG",
+            predicted_dna="ATGAAAACTGTTCGTCAGGAACGTCTGAAATCTATTGTTCGTATTCTGGAACGTTCTAAAGAACCGGTTTCTGGTGCTCAACTGGCTGAAGAACTGTCTGTTTCTCGTCAGGTTATTGTTCAGGACATTGCTTACCTGCGTTCTCTGGGTTATAA",
+        )
+        formatted_output = format_model_output(output)
+        self.assertIsInstance(formatted_output, str)
+        self.assertIn("Organism", formatted_output)
+        self.assertIn("Escherichia coli", formatted_output)
+        self.assertIn("Input Protein", formatted_output)
+        self.assertIn(
+            "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG",
+            formatted_output,
+        )
+        self.assertIn("Processed Input", formatted_output)
+        self.assertIn(
+            "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG",
+            formatted_output,
+        )
+        self.assertIn("Predicted DNA", formatted_output)
+        self.assertIn(
+            "ATGAAAACTGTTCGTCAGGAACGTCTGAAATCTATTGTTCGTATTCTGGAACGTTCTAAAGAACCGGTTTCTGGTGCTCAACTGGCTGAAGAACTGTCTGTTTCTCGTCAGGTTATTGTTCAGGACATTGCTTACCTGCGTTCTCTGGGTTATAA",
+            formatted_output,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

scripts/tests/test_CodonPrediction.py

@@ -0,0 +1,594 @@
+import random
+import unittest
+import warnings
+
+import torch
+
+from CodonTransformer.CodonData import get_amino_acid_sequence
+from CodonTransformer.CodonPrediction import (
+    load_model,
+    load_tokenizer,
+    predict_dna_sequence,
+)
+from CodonTransformer.CodonUtils import (
+    AMINO_ACIDS,
+    ORGANISM2ID,
+    STOP_SYMBOLS,
+    DNASequencePrediction,
+)
+
+
+class TestCodonPrediction(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        # Suppress warnings about loading from HuggingFace
+        for message in [
+            "Tokenizer path not provided. Loading from HuggingFace.",
+            "Model path not provided. Loading from HuggingFace.",
+        ]:
+            warnings.filterwarnings("ignore", message=message)
+
+        cls.model = load_model(device=cls.device)
+        cls.tokenizer = load_tokenizer()
+
+    def test_predict_dna_sequence_valid_input(self):
+        protein_sequence = "MWWMW"
+        organism = "Escherichia coli general"
+        result = predict_dna_sequence(
+            protein_sequence,
+            organism,
+            device=self.device,
+            tokenizer=self.tokenizer,
+            model=self.model,
+        )
+        self.assertIsInstance(result.predicted_dna, str)
+        self.assertTrue(
+            all(nucleotide in "ATCG" for nucleotide in result.predicted_dna)
+        )
+        self.assertEqual(result.predicted_dna, "ATGTGGTGGATGTGGTGA")
+
+    def test_predict_dna_sequence_non_deterministic(self):
+        protein_sequence = "MFWY"
+        organism = "Escherichia coli general"
+        num_iterations = 100
+        temperatures = [0.2, 0.5, 0.8]
+        possible_outputs = set()
+        possible_encodings_wo_stop = {
+            "ATGTTTTGGTAT",
+            "ATGTTCTGGTAT",
+            "ATGTTTTGGTAC",
+            "ATGTTCTGGTAC",
+        }
+        for _ in range(num_iterations):
+            for temperature in temperatures:
+                result = predict_dna_sequence(
+                    protein=protein_sequence,
+                    organism=organism,
+                    device=self.device,
+                    tokenizer=self.tokenizer,
+                    model=self.model,
+                    deterministic=False,
+                    temperature=temperature,
+                )
+                possible_outputs.add(result.predicted_dna[:-3])  # Remove stop codon
+
+        self.assertEqual(possible_outputs, possible_encodings_wo_stop)
+
+    def test_predict_dna_sequence_invalid_inputs(self):
+        test_cases = [
+            ("MKTZZFVLLL?", "Escherichia coli general", "invalid protein sequence"),
+            ("MKTFFVLLL", "Alien $%#@!", "invalid organism code"),
+            ("", "Escherichia coli general", "empty protein sequence"),
+        ]
+
+        for protein_sequence, organism, error_type in test_cases:
+            with self.subTest(error_type=error_type):
+                with self.assertRaises(ValueError):
+                    predict_dna_sequence(
+                        protein_sequence,
+                        organism,
+                        device=self.device,
+                        tokenizer=self.tokenizer,
+                        model=self.model,
+                    )
+
+    def test_predict_dna_sequence_top_p_effect(self):
+        """Test that changing top_p affects the diversity of outputs."""
+        protein_sequence = "MFWY"
+        organism = "Escherichia coli general"
+        num_iterations = 50
+        temperature = 0.5
+        top_p_values = [0.8, 0.95]
+        outputs_by_top_p = {top_p: set() for top_p in top_p_values}
+
+        for top_p in top_p_values:
+            for _ in range(num_iterations):
+                result = predict_dna_sequence(
+                    protein=protein_sequence,
+                    organism=organism,
+                    device=self.device,
+                    tokenizer=self.tokenizer,
+                    model=self.model,
+                    deterministic=False,
+                    temperature=temperature,
+                    top_p=top_p,
+                )
+                outputs_by_top_p[top_p].add(
+                    result.predicted_dna[:-3]
+                )  # Remove stop codon
+
+        # Assert that higher top_p results in more diverse outputs
+        diversity_lower_top_p = len(outputs_by_top_p[0.8])
+        diversity_higher_top_p = len(outputs_by_top_p[0.95])
+        self.assertGreaterEqual(
+            diversity_higher_top_p,
+            diversity_lower_top_p,
+            "Higher top_p should result in more diverse outputs",
+        )
+
+    def test_predict_dna_sequence_invalid_temperature_and_top_p(self):
+        """Test that invalid temperature and top_p values raise ValueError."""
+        protein_sequence = "MWWMW"
+        organism = "Escherichia coli general"
+        invalid_params = [
+            {"temperature": -0.1, "top_p": 0.95},
+            {"temperature": 0, "top_p": 0.95},
+            {"temperature": 0.5, "top_p": -0.1},
+            {"temperature": 0.5, "top_p": 1.1},
+        ]
+
+        for params in invalid_params:
+            with self.subTest(params=params):
+                with self.assertRaises(ValueError):
+                    predict_dna_sequence(
+                        protein=protein_sequence,
+                        organism=organism,
+                        device=self.device,
+                        tokenizer=self.tokenizer,
+                        model=self.model,
+                        deterministic=False,
+                        temperature=params["temperature"],
+                        top_p=params["top_p"],
+                    )
+
+    def test_predict_dna_sequence_translation_consistency(self):
+        """Test that the predicted DNA translates back to the original protein."""
+        protein_sequence = "MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVE"
+        organism = "Escherichia coli general"
+        result = predict_dna_sequence(
+            protein=protein_sequence,
+            organism=organism,
+            device=self.device,
+            tokenizer=self.tokenizer,
+            model=self.model,
+            deterministic=True,
+        )
+
+        # Translate predicted DNA back to protein
+        translated_protein = get_amino_acid_sequence(result.predicted_dna[:-3])
+
+        self.assertEqual(
+            translated_protein,
+            protein_sequence,
+            "Translated protein does not match the original protein sequence",
+        )
+
+    def test_predict_dna_sequence_long_protein_sequence(self):
+        """Test the function with a very long protein sequence to check performance and correctness."""
+        protein_sequence = (
+            "M"
+            + "MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKTRREAEDLQVGQVELGG"
+            * 20
+            + STOP_SYMBOLS[0]
+        )
+        organism = "Escherichia coli general"
+        result = predict_dna_sequence(
+            protein=protein_sequence,
+            organism=organism,
+            device=self.device,
+            tokenizer=self.tokenizer,
+            model=self.model,
+            deterministic=True,
+        )
+
+        # Check that the predicted DNA translates back to the original protein
+        dna_sequence = result.predicted_dna[:-3]
+        translated_protein = get_amino_acid_sequence(dna_sequence)
+        self.assertEqual(
+            translated_protein,
+            protein_sequence[:-1],
+            "Translated protein does not match the original long protein sequence",
+        )
+
+    def test_predict_dna_sequence_edge_case_organisms(self):
+        """Test the function with organism IDs at the boundaries of the mapping."""
+        protein_sequence = "MWWMW"
+        # Assuming ORGANISM2ID has IDs starting from 0 to N
+        min_organism_id = min(ORGANISM2ID.values())
+        max_organism_id = max(ORGANISM2ID.values())
+        organisms = [min_organism_id, max_organism_id]
+
+        for organism_id in organisms:
+            with self.subTest(organism_id=organism_id):
+                result = predict_dna_sequence(
+                    protein=protein_sequence,
+                    organism=organism_id,
+                    device=self.device,
+                    tokenizer=self.tokenizer,
+                    model=self.model,
+                    deterministic=True,
+                )
+                self.assertIsInstance(result.predicted_dna, str)
+                self.assertTrue(
+                    all(nucleotide in "ATCG" for nucleotide in result.predicted_dna)
+                )
+
+    def test_predict_dna_sequence_concurrent_calls(self):
+        """Test the function's behavior under concurrent execution."""
+        import threading
+
+        protein_sequence = "MWWMW"
+        organism = "Escherichia coli general"
+        results = []
+
+        def call_predict():
+            result = predict_dna_sequence(
+                protein=protein_sequence,
+                organism=organism,
+                device=self.device,
+                tokenizer=self.tokenizer,
+                model=self.model,
+                deterministic=True,
+            )
+            results.append(result.predicted_dna)
+
+        threads = [threading.Thread(target=call_predict) for _ in range(10)]
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join()
+
+        self.assertEqual(len(results), 10)
+        self.assertTrue(all(dna == results[0] for dna in results))
+
+    def test_predict_dna_sequence_random_seed_consistency(self):
+        """Test that setting a random seed results in consistent outputs in non-deterministic mode."""
+        protein_sequence = "MFWY"
+        organism = "Escherichia coli general"
+        temperature = 0.5
+        top_p = 0.95
+        torch.manual_seed(42)
+
+        result1 = predict_dna_sequence(
+            protein=protein_sequence,
+            organism=organism,
+            device=self.device,
+            tokenizer=self.tokenizer,
+            model=self.model,
+            deterministic=False,
+            temperature=temperature,
+            top_p=top_p,
+        )
+
+        torch.manual_seed(42)
+
+        result2 = predict_dna_sequence(
+            protein=protein_sequence,
+            organism=organism,
+            device=self.device,
+            tokenizer=self.tokenizer,
+            model=self.model,
+            deterministic=False,
+            temperature=temperature,
+            top_p=top_p,
+        )
+
+        self.assertEqual(
+            result1.predicted_dna,
+            result2.predicted_dna,
+            "Outputs should be consistent when random seed is set",
+        )
+
+    def test_predict_dna_sequence_invalid_tokenizer_and_model(self):
+        """Test that providing invalid tokenizer or model raises appropriate exceptions."""
+        protein_sequence = "MWWMW"
+        organism = "Escherichia coli general"
+
+        with self.subTest("Invalid tokenizer"):
+            with self.assertRaises(Exception):
+                predict_dna_sequence(
+                    protein=protein_sequence,
+                    organism=organism,
+                    device=self.device,
+                    tokenizer="invalid_tokenizer_path",
+                    model=self.model,
+                )
+
+        with self.subTest("Invalid model"):
+            with self.assertRaises(Exception):
+                predict_dna_sequence(
+                    protein=protein_sequence,
+                    organism=organism,
+                    device=self.device,
+                    tokenizer=self.tokenizer,
+                    model="invalid_model_path",
+                )
+
+    def test_predict_dna_sequence_stop_codon_handling(self):
+        """Test the function's handling of protein sequences ending with a non '_' or '*' stop symbol."""
+        protein_sequence = "MWW/"
+        organism = "Escherichia coli general"
+
+        with self.assertRaises(ValueError):
+            predict_dna_sequence(
+                protein=protein_sequence,
+                organism=organism,
+                device=self.device,
+                tokenizer=self.tokenizer,
+                model=self.model,
+            )
+
+    def test_predict_dna_sequence_device_compatibility(self):
+        """Test that the function works correctly on both CPU and GPU devices."""
+        protein_sequence = "MWWMW"
+        organism = "Escherichia coli general"
+
+        devices = [torch.device("cpu")]
+        if torch.cuda.is_available():
+            devices.append(torch.device("cuda"))
+
+        for device in devices:
+            with self.subTest(device=device):
+                result = predict_dna_sequence(
+                    protein=protein_sequence,
+                    organism=organism,
+                    device=device,
+                    tokenizer=self.tokenizer,
+                    model=self.model,
+                    deterministic=True,
+                )
+                self.assertIsInstance(result.predicted_dna, str)
+                self.assertTrue(
+                    all(nucleotide in "ATCG" for nucleotide in result.predicted_dna)
+                )
+
+    def test_predict_dna_sequence_random_proteins(self):
+        """Test random proteins to ensure translated DNA matches the original protein."""
+        organism = "Escherichia coli general"
+        num_tests = 200
+
+        for _ in range(num_tests):
+            # Generate a random protein sequence of random length between 10 and 50
+            protein_length = random.randint(10, 500)
+            protein_sequence = "M" + "".join(
+                random.choices(AMINO_ACIDS, k=protein_length - 1)
+            )
+            protein_sequence += random.choice(STOP_SYMBOLS)
+
+            result = predict_dna_sequence(
+                protein=protein_sequence,
+                organism=organism,
+                device=self.device,
+                tokenizer=self.tokenizer,
+                model=self.model,
+                deterministic=True,
+            )
+
+            # Remove stop codon from predicted DNA
+            dna_sequence = result.predicted_dna[:-3]
+
+            # Translate predicted DNA back to protein
+            translated_protein = get_amino_acid_sequence(dna_sequence)
+            self.assertEqual(
+                translated_protein,
+                protein_sequence[:-1],  # Remove stop symbol
+                f"Translated protein does not match the original protein sequence for protein: {protein_sequence}",
+            )
+
+    def test_predict_dna_sequence_long_protein_over_max_length(self):
+        """Test that the model handles protein sequences longer than 2048 amino acids."""
+        # Create a protein sequence longer than 2048 amino acids
+        base_sequence = (
+            "MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKTRREAEDLQVGQVELGG"
+        )
+        protein_sequence = base_sequence * 100  # Length > 2048 amino acids
+        organism = "Escherichia coli general"
+
+        result = predict_dna_sequence(
+            protein=protein_sequence,
+            organism=organism,
+            device=self.device,
+            tokenizer=self.tokenizer,
+            model=self.model,
+            deterministic=True,
+        )
+
+        # Remove stop codon from predicted DNA
+        dna_sequence = result.predicted_dna[:-3]
+        translated_protein = get_amino_acid_sequence(dna_sequence)
+
+        # Due to potential model limitations, compare up to the model's max supported length
+        max_length = len(translated_protein)
+        self.assertEqual(
+            translated_protein[:max_length],
+            protein_sequence[:max_length],
+            "Translated protein does not match the original protein sequence up to the maximum length supported.",
+        )
+
+    def test_predict_dna_sequence_multi_output(self):
+        """Test that the function returns multiple sequences when num_sequences > 1."""
+        protein_sequence = "MFQLLAPWY"
+        organism = "Escherichia coli general"
+        num_sequences = 20
+
+        result = predict_dna_sequence(
+            protein=protein_sequence,
+            organism=organism,
+            device=self.device,
+            tokenizer=self.tokenizer,
+            model=self.model,
+            deterministic=False,
+            num_sequences=num_sequences,
+        )
+
+        self.assertIsInstance(result, list)
+        self.assertEqual(len(result), num_sequences)
+
+        for prediction in result:
+            self.assertIsInstance(prediction, DNASequencePrediction)
+            self.assertTrue(
+                all(nucleotide in "ATCG" for nucleotide in prediction.predicted_dna)
+            )
+
+            # Check that all predicted DNA sequences translate back to the original protein
+            translated_protein = get_amino_acid_sequence(prediction.predicted_dna[:-3])
+            self.assertEqual(translated_protein, protein_sequence)
+
+    def test_predict_dna_sequence_deterministic_multi_raises_error(self):
+        """Test that requesting multiple sequences in deterministic mode raises an error."""
+        protein_sequence = "MFWY"
+        organism = "Escherichia coli general"
+
+        with self.assertRaises(ValueError):
+            predict_dna_sequence(
+                protein=protein_sequence,
+                organism=organism,
+                device=self.device,
+                tokenizer=self.tokenizer,
+                model=self.model,
+                deterministic=True,
+                num_sequences=3,
+            )
+
+    def test_predict_dna_sequence_multi_diversity(self):
+        """Test that multiple sequences generated are diverse."""
+        protein_sequence = "MFWYMFWY"
+        organism = "Escherichia coli general"
+        num_sequences = 10
+
+        result = predict_dna_sequence(
+            protein=protein_sequence,
+            organism=organism,
+            device=self.device,
+            tokenizer=self.tokenizer,
+            model=self.model,
+            deterministic=False,
+            num_sequences=num_sequences,
+            temperature=0.8,
+        )
+
+        unique_sequences = set(prediction.predicted_dna for prediction in result)
+
+        self.assertGreater(
+            len(unique_sequences),
+            2,
+            "Multiple sequence generation should produce diverse results",
+        )
+
+        # Check that all sequences are valid translations of the input protein
+        for prediction in result:
+            translated_protein = get_amino_acid_sequence(prediction.predicted_dna[:-3])
+            self.assertEqual(translated_protein, protein_sequence)
+
+    def test_predict_dna_sequence_match_protein_repetitive(self):
+        """Test that match_protein=True correctly handles highly repetitive and unconventional sequences."""
+        test_sequences = (
+            "QQQQQQQQQQQQQQQQ_",
+            "KRKRKRKRKRKRKRKR_",
+            "PGPGPGPGPGPGPGPG_",
+            "DEDEDEDEDEDEDEDEDE_",
+            "M_M_M_M_M_",
+            "MMMMMMMMMM_",
+            "WWWWWWWWWW_",
+            "CCCCCCCCCC_",
+            "MWCHMWCHMWCH_",
+            "Q_QQ_QQQ_QQQQ_",
+            "MWMWMWMWMWMW_",
+            "CCCHHHMMMWWW_",
+            "_",
+            "M_",
+            "MGWC_",
+        )
+
+        organism = "Homo sapiens"
+
+        for protein_sequence in test_sequences:
+            # Generate sequence with match_protein=True
+            result = predict_dna_sequence(
+                protein=protein_sequence,
+                organism=organism,
+                device=self.device,
+                tokenizer=self.tokenizer,
+                model=self.model,
+                deterministic=False,
+                temperature=20,  # High temperature to test protein matching
+                match_protein=True,
+            )
+
+            dna_sequence = result.predicted_dna
+            translated_protein = get_amino_acid_sequence(dna_sequence)
+
+            self.assertEqual(
+                translated_protein,
+                protein_sequence,
+                f"Translated protein must match original when match_protein=True. Failed for sequence: {protein_sequence}",
+            )
+
+    def test_predict_dna_sequence_match_protein_rare_amino_acids(self):
+        """Test match_protein with rare amino acids that have limited codon options."""
+        # Methionine (M) and Tryptophan (W) have only one codon each
+        # While Leucine (L) has 6 codons - testing contrast
+        protein_sequence = "MWLLLMWLLL"
+        organism = "Escherichia coli general"
+
+        # Run multiple predictions
+        results = []
+        num_iterations = 10
+
+        for _ in range(num_iterations):
+            result = predict_dna_sequence(
+                protein=protein_sequence,
+                organism=organism,
+                device=self.device,
+                tokenizer=self.tokenizer,
+                model=self.model,
+                deterministic=False,
+                temperature=20,  # High temperature to test protein matching
+                match_protein=True,
+            )
+            results.append(result.predicted_dna)
+
+        # Check all sequences
+        for dna_sequence in results:
+            # Verify M always uses ATG
+            m_positions = [0, 5]  # Known positions of M in sequence
+            for pos in m_positions:
+                self.assertEqual(
+                    dna_sequence[pos * 3 : (pos + 1) * 3],
+                    "ATG",
+                    "Methionine must use ATG codon.",
+                )
+
+            # Verify W always uses TGG
+            w_positions = [1, 6]  # Known positions of W in sequence
+            for pos in w_positions:
+                self.assertEqual(
+                    dna_sequence[pos * 3 : (pos + 1) * 3],
+                    "TGG",
+                    "Tryptophan must use TGG codon.",
+                )
+
+            # Verify all L codons are valid
+            l_positions = [2, 3, 4, 7, 8, 9]  # Known positions of L in sequence
+            l_codons = [dna_sequence[pos * 3 : (pos + 1) * 3] for pos in l_positions]
+            valid_l_codons = {"TTA", "TTG", "CTT", "CTC", "CTA", "CTG"}
+            self.assertTrue(
+                all(codon in valid_l_codons for codon in l_codons),
+                "All Leucine codons must be valid",
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

scripts/tests/test_CodonUtils.py

@@ -0,0 +1,125 @@
+import os
+import pickle
+import tempfile
+import unittest
+
+from CodonTransformer.CodonUtils import (
+    ProteinConfig,
+    find_pattern_in_fasta,
+    get_organism2id_dict,
+    get_taxonomy_id,
+    load_pkl_from_url,
+    load_python_object_from_disk,
+    save_python_object_to_disk,
+    sort_amino2codon_skeleton,
+)
+
+
+class TestCodonUtils(unittest.TestCase):
+    def test_config_manager(self):
+        with ProteinConfig() as config:
+            config.set("ambiguous_aminoacid_behavior", "standardize_deterministic")
+            self.assertEqual(
+                config.get("ambiguous_aminoacid_behavior"), "standardize_deterministic"
+            )
+            config.set("ambiguous_aminoacid_map_override", {"X": ["A", "G"]})
+            self.assertEqual(
+                config.get("ambiguous_aminoacid_map_override"), {"X": ["A", "G"]}
+            )
+            config.update(
+                {
+                    "ambiguous_aminoacid_behavior": "raise_error",
+                    "ambiguous_aminoacid_map_override": {"X": ["A", "G"]},
+                }
+            )
+            self.assertEqual(config.get("ambiguous_aminoacid_behavior"), "raise_error")
+            self.assertEqual(
+                config.get("ambiguous_aminoacid_map_override"), {"X": ["A", "G"]}
+            )
+            try:
+                config.set("invalid_key", "invalid_value")
+                self.fail("Expected ValueError")
+            except ValueError:
+                pass
+        with ProteinConfig() as config:
+            self.assertEqual(
+                config.get("ambiguous_aminoacid_behavior"), "standardize_random"
+            )
+            self.assertEqual(config.get("ambiguous_aminoacid_map_override"), {})
+
+    def test_load_python_object_from_disk(self):
+        test_obj = {"key1": "value1", "key2": 2}
+        with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as temp_file:
+            temp_file_name = temp_file.name
+            save_python_object_to_disk(test_obj, temp_file_name)
+        loaded_obj = load_python_object_from_disk(temp_file_name)
+        self.assertEqual(test_obj, loaded_obj)
+        os.remove(temp_file_name)
+
+    def test_save_python_object_to_disk(self):
+        test_obj = [1, 2, 3, 4, 5]
+        with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as temp_file:
+            temp_file_name = temp_file.name
+            save_python_object_to_disk(test_obj, temp_file_name)
+            self.assertTrue(os.path.exists(temp_file_name))
+        os.remove(temp_file_name)
+
+    def test_find_pattern_in_fasta(self):
+        text = (
+            ">seq1 [keyword=value1]\nATGCGTACGTAGCTAG\n"
+            ">seq2 [keyword=value2]\nGGTACGATCGATCGAT"
+        )
+        self.assertEqual(find_pattern_in_fasta("keyword", text), "value1")
+        self.assertEqual(find_pattern_in_fasta("nonexistent", text), "")
+
+    def test_get_organism2id_dict(self):
+        with tempfile.NamedTemporaryFile(
+            mode="w", delete=True, suffix=".csv"
+        ) as temp_file:
+            temp_file.write("0,Escherichia coli\n1,Homo sapiens\n2,Mus musculus")
+            temp_file.flush()
+            organism2id = get_organism2id_dict(temp_file.name)
+            self.assertEqual(
+                organism2id,
+                {"Escherichia coli": 0, "Homo sapiens": 1, "Mus musculus": 2},
+            )
+
+    def test_get_taxonomy_id(self):
+        taxonomy_dict = {
+            "Escherichia coli": 562,
+            "Homo sapiens": 9606,
+            "Mus musculus": 10090,
+        }
+        with tempfile.NamedTemporaryFile(suffix=".pkl", delete=True) as temp_file:
+            temp_file_name = temp_file.name
+            save_python_object_to_disk(taxonomy_dict, temp_file_name)
+            self.assertEqual(get_taxonomy_id(temp_file_name, "Escherichia coli"), 562)
+            self.assertEqual(
+                get_taxonomy_id(temp_file_name, return_dict=True), taxonomy_dict
+            )
+
+    def test_sort_amino2codon_skeleton(self):
+        amino2codon = {
+            "A": (["GCT", "GCC", "GCA", "GCG"], [0.0, 0.0, 0.0, 0.0]),
+            "C": (["TGT", "TGC"], [0.0, 0.0]),
+        }
+        sorted_amino2codon = sort_amino2codon_skeleton(amino2codon)
+        self.assertEqual(
+            sorted_amino2codon,
+            {
+                "A": (["GCA", "GCC", "GCG", "GCT"], [0.0, 0.0, 0.0, 0.0]),
+                "C": (["TGC", "TGT"], [0.0, 0.0]),
+            },
+        )
+
+    def test_load_pkl_from_url(self):
+        url = "https://example.com/test.pkl"
+        expected_obj = {"key": "value"}
+        with unittest.mock.patch("requests.get") as mock_get:
+            mock_get.return_value.content = pickle.dumps(expected_obj)
+            loaded_obj = load_pkl_from_url(url)
+        self.assertEqual(loaded_obj, expected_obj)
+
+
+if __name__ == "__main__":
+    unittest.main()

setup.py

@@ -0,0 +1,42 @@
+# setup.py
+import os
+
+from setuptools import find_packages, setup
+
+
+def read_requirements():
+    with open("requirements.txt") as f:
+        return [line.strip() for line in f if line.strip() and not line.startswith("#")]
+
+
+def read_readme():
+    here = os.path.abspath(os.path.dirname(__file__))
+    readme_path = os.path.join(here, "README.md")
+
+    with open(readme_path, "r", encoding="utf-8") as f:
+        return f.read()
+
+
+setup(
+    name="CodonTransformer",
+    version="1.6.7",
+    packages=find_packages(where="model"),
+    package_dir={"": "model"},
+    install_requires=read_requirements(),
+    author="Adibvafa Fallahpour",
+    author_email="Adibvafa.fallahpour@mail.utoronto.ca",
+    description=(
+        "The ultimate tool for codon optimization, "
+        "transforming protein sequences into optimized DNA sequences "
+        "specific for your target organisms."
+    ),
+    long_description=read_readme(),
+    long_description_content_type="text/markdown",
+    url="https://github.com/adibvafa/CodonTransformer",
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: Apache Software License",
+        "Operating System :: OS Independent",
+    ],
+    python_requires=">=3.9",
+)

weight/checkpoints/README.md

@@ -0,0 +1 @@
+#放置CodonTransformer 模型权重