plmtogo – protein generation

plmtogo models protein sequences conditioned on Gene Ontology (GO) functional classes. The project ships a lightweight tokenizer, dataset tooling for FASTA files, and single-GPU training and sampling scripts that run comfortably on workstations such as an RTX 6000 Ada.

Quick start

Launch ./minirun.sh help to see the main workflows. The script wraps common commands for base training, fine-tuning, and generation so you can reproduce the end-to-end flow with a single entry point:

./minirun.sh train               # train the base model (writes to checkpoints/)
./minirun.sh generate-base       # sample sequences from the base checkpoint
./minirun.sh fine-tune           # fine-tune on curated data under fine-tuning_data/
./minirun.sh generate-finetuned  # sample from the fine-tuned checkpoint

Feel free to copy these commands or adapt the underlying Python invocations to suit your experiments.

Dataset layout

• Place GO-specific FASTA files under training_data/. Filenames should start with the GO identifier (e.g., GO0000001_description.faa). Each file can carry one or more protein records; headers are preserved for traceability.
• During preprocessing we normalise sequences to uppercase amino acids and replace uncommon symbols with X. See plmtogo/protein_data.py for parsing rules.

Tokenisation & prompts

• plmtogo/protein_tokenizer.py implements a fixed vocabulary: 20 canonical amino acids, common ambiguous residues, and special markers (<bos>, <seq_start>, <seq_end>, <sep>).
• Every GO term receives a dedicated token such as <CLASS_GO0000001>. Training sequences use the template:
  <bos><CLASS_GO0000001><sep><seq_start>...amino acids...<seq_end><eos>.
• At inference time you only need to provide the GO identifier; helper scripts translate it into the matching special token.

Training on a single GPU

python scripts/protein_train.py \
  --data-dir training_data \
  --device cuda:0 \  # RTX 6000 Ada
  --epochs 50 \
  --batch-size 32 \
  --max-length 512 \
  --n-layer 8 \
  --n-head 8 \
  --n-embd 512
# add --run-type fine-tuning to rename the output curve to fine-tuning_curve.png when resuming a model

• The script handles train/validation splits, computes loss-per-token and perplexity, and samples preview sequences each epoch.
• --max-length should exceed the longest tokenised sequence (class token + amino acids + delimiters). Increase the value for longer proteins or reduce it to drop excessive examples.
• Checkpoints are saved to checkpoints/ and contain the model weights, configuration, tokenizer vocabulary, and chosen context length.
• A training loss curve (training_curve.png) is written alongside the checkpoint for quick visualization of convergence (requires matplotlib; the plot step is skipped if the library is missing). Pass --run-type fine-tuning to label the plot as fine-tuning_curve.png during fine-tune runs.

Sampling from a checkpoint

python scripts/protein_generate.py \
  --checkpoint checkpoints/protein-gpt.pt \
  --go-term GO:0000001 \
  --num-samples 5 \
  --temperature 0.7 \
  --top-k 25

• When --go-term is omitted, the script generates sequences for every known GO term.
• Adjust --max-tokens to cap the generation horizon; decoding stops automatically once <seq_end> is emitted.

Project structure

• plmtogo/protein_data.py – FASTA loading, cleaning, and stratified splits.
• plmtogo/protein_tokenizer.py – fixed amino-acid vocabulary with dynamic class tokens.
• plmtogo/protein_dataset.py – PyTorch dataset plus padding/label collation.
• scripts/protein_train.py – single-GPU trainer with AdamW, mixed precision (when available), and preview sampling.
• scripts/protein_generate.py – inference CLI for GO-conditioned sequence synthesis.
• Legacy text workflows from the earlier NanoChat project have been retired to keep the focus on proteins.

Roadmap

• Expand the dataset to cover additional functional classes (GO, PFAM, EC numbers).
• Explore protein-aware architectures (longer context, relative position encodings, structured decoders).
• Integrate biological evaluation metrics (motif coverage, disorder prediction, homology filtering) and fine-tuning recipes.