CodonTranslator

CodonTranslator is a protein-conditioned codon sequence generation model trained on the representative-only data_v3 release.

This repository is the public model and training-code release. It contains:

• final_model/: inference-ready weights
• src/, train.py, sampling.py: training and inference code
• resplit_data_v3.py: the data_v3 reconstruction pipeline
• slurm/: the single-node H200 training and data rebuild submission scripts
• CodonTranslator/ and pyproject.toml: a lightweight packaged inference wrapper

Training configuration

• Architecture: hidden=750, layers=20, heads=15, mlp_ratio=3.2
• Attention: mha
• Precision: bf16
• Parallelism: FSDP full shard
• Effective global batch: 1536
• Weight decay: 1e-4
• Dataset: alegendaryfish/CodonTranslator-data

Dataset release

The corresponding public dataset and species embedding release is:

• alegendaryfish/CodonTranslator-data

That dataset repo contains:

• final representative-only train/, val/, test/ parquet shards
• embeddings_v2/
• split audit files and reconstruction metadata

Quick start

Install

git clone https://huggingface.co/alegendaryfish/CodonTranslator
cd CodonTranslator
conda env create -f environment.yml
conda activate codontranslator
pip install -r requirements.txt
pip install -e .

Both import styles are supported:

from CodonTranslator import CodonTranslator

from codontranslator import CodonTranslator

Train

python train.py \
  --train_data /path/to/train \
  --val_data /path/to/val \
  --embeddings_dir /path/to/embeddings_v2 \
  --output_dir outputs \
  --fsdp \
  --bf16 \
  --attn mha \
  --hidden 750 \
  --layers 20 \
  --heads 15 \
  --mlp_ratio 3.2 \
  --batch_size 48 \
  --grad_accum 4 \
  --epochs 3 \
  --lr 7e-5 \
  --weight_decay 1e-4

The included Slurm launchers use the same training flags as the local single-node H200 workflow:

• slurm/train_v3_h200_8x_single.sbatch
• slurm/submit_train_v3_h200_8x_chain.sh

Sample

python sampling.py \
  --model_path final_model \
  --embeddings_dir /path/to/embeddings_v2 \
  --species "Panicum hallii" \
  --protein_sequence "MSEQUENCE" \
  --strict_species_lookup

Notes

• Training uses precomputed embeddings_v2 for species conditioning.
• The data split is built in protein space with MMseqs clustering and binomial-species test holdout.
• final_model/ is the published inference entrypoint.
• For compatibility, released model directories contain both trainer_config.json and config.json.

Sampling arguments

• enforce_mapping: when True, each generated codon is constrained to encode the provided amino acid at that position.
• temperature: softmax temperature. Lower values are more deterministic; 0 selects argmax greedily.
• top_k: keep only the k highest-logit codon candidates before sampling.
• top_p: nucleus sampling threshold; keep the smallest probability mass whose cumulative sum reaches p.

Inference update (2026-10-30)

The command-line sampler and packaged API now use the corrected inference backbone. For batches with different sequence lengths, cached decoding excludes padded keys and assigns rotary positions from each sequence's valid history. Completed rows retain a valid padding path when synonymous-codon constraints are enabled.

Update an existing checkout with git pull and reinstall with pip install -e .. This code update uses the existing model weights.

Run the CPU cache and mixed-length sampling regression checks with:

python -m unittest discover -s tests