JoyXiangLab/rnaseek-full

RNASeek Full

This repository contains code, model checkpoints, notebooks, and data artifacts for RNA sequence generation and regression workflows used in the RNASeek project.

Repository Layout

• efficiency_figure2/ - sequence efficiency regression training scripts, input JSON/JSONL files, TensorBoard logs, and checkpoints.
• regression_stability_functionalviral/ - functional viral element stability regression data, alignment/count summaries, notebooks, training scripts, and checkpoints.
• ribozymegen-figure7/ - supervised ribozyme causal language model fine-tuning inputs and scripts.
• ribozymegen-rl-figure7/ - ribozyme reinforcement learning scripts using reward model and directive scoring HTTP APIs.
• utrgen/ - supervised UTR causal language model fine-tuning scripts, train/validation text files, and checkpoints.
• utrgen-rl.zip - archived UTR reinforcement learning workflow. The raw utrgen-rl/ directory is intentionally distributed as this zip artifact because of its size.
• requirements.txt - frozen Python package list from the environment used to prepare this release.

Environment

Create and activate a Python environment, then install the frozen dependencies:

pip install -r requirements.txt

If you need FlashAttention, use a CUDA/PyTorch-compatible build. Prebuilt wheel instructions are available at https://github.com/mjun0812/flash-attention-prebuild-wheels.

The portable runtime is distributed as portable_runtime.tar.gz. Extract it before using commands that reference portable_runtime/env/bin/python:

tar -xzf portable_runtime.tar.gz

Most training scripts expect CUDA-enabled PyTorch and access to local model checkpoints referenced by each script's path constants. Review the path variables at the top of each script before launching a run.

Common Entry Points

Supervised language-model fine-tuning:

cd utrgen
python3 train.py

cd ribozymegen-figure7
python3 train.py

Regression model training:

cd efficiency_figure2
python3 betterTrain.py

cd regression_stability_functionalviral
python3 betterTrain.py

Packed Regression Model Validation

Model weight files are stored in ZIP archives on the Hub as *.safetensors.zip. Restore them before loading checkpoints:

python3 scripts/unzip_safetensors.py --overwrite

or:

./scripts/unzip_safetensors.sh

The efficiency and stability regression checkpoints include the regression head inside the model checkpoint directory as regression_head.safetensors. The checkpoint index maps the packed head tensors under regression_head.*, while the original regression_head.pt remains as a fallback for older scripts.

Run validation from the repository root. To force CPU-only evaluation, clear CUDA_VISIBLE_DEVICES and pass --device cpu:

CUDA_VISIBLE_DEVICES= MPLCONFIGDIR=/tmp/matplotlib-rnaseek \
  portable_runtime/env/bin/python \
  efficiency_figure2/validate_packed_regression_model.py \
  --device cpu \
  --batch-size 16

This writes:

• efficiency_figure2/packed_validation/valid_predictions.tsv
• efficiency_figure2/packed_validation/valid_metrics.json
• efficiency_figure2/packed_validation/valid_scatter.png

The efficiency validator defaults to:

• model: efficiency_figure2/qwen_regression_ckpt/clean_cosine_restart_besthp_preview_fixed-wd-0.9_reproduce/checkpoint-304419
• tokenizer: the same checkpoint directory, unless --tokenizer-dir is provided
• validation data: efficiency_figure2/evenBetterDataFolded-vl.json

Run the stability packed model validation the same way:

CUDA_VISIBLE_DEVICES= MPLCONFIGDIR=/tmp/matplotlib-rnaseek \
  portable_runtime/env/bin/python \
  regression_stability_functionalviral/validate_packed_regression_model.py \
  --device cpu \
  --batch-size 4

This writes:

• regression_stability_functionalviral/packed_validation/valid_predictions.tsv
• regression_stability_functionalviral/packed_validation/valid_metrics.json
• regression_stability_functionalviral/packed_validation/valid_scatter.png

The stability validator defaults to:

• model: regression_stability_functionalviral/checkpoint-44040_best
• tokenizer: regression_stability_functionalviral/checkpoint-388560_base
• validation data: regression_stability_functionalviral/training_seq_score_extreme_weighted.tsv with split == "val"

To repack a regression head after replacing regression_head.pt, run:

portable_runtime/env/bin/python scripts/pack_regression_head_into_safetensors_index.py \
  path/to/checkpoint

The packer writes regression_head.safetensors and updates model.safetensors.index.json so packed-model loaders can discover the head.

UTR GRPO reinforcement learning uses a separate reward-model API process:

cd utrgen-rl
CUDA_VISIBLE_DEVICES=1 python3 rm_evaluator_api.py --device cuda:0 --port 8008
CUDA_VISIBLE_DEVICES=0 python3 train_grpo.py --rm-api-url http://127.0.0.1:8008

Ribozyme GRPO reinforcement learning expects reward and directive scoring services matching the endpoints configured in ribozymegen-rl-figure7/testgrpo_retry-linearblend_prior.py.

Notes

• Scripts are research-oriented and expose run settings as constants near the top of each file.
• Notebook checkpoint folders, TensorBoard logs, and model checkpoints are included where present in the release artifact.
• The raw utrgen-rl/ directory is not required for Hub upload when utrgen-rl.zip is present.