ProRiboGen — Inference Package

Inference-only package for a website / API backend. No training or test CSVs.

ProRiboGen has two modules:

1. Generation — protein amino-acid sequence → RNA FASTA + motifs (+ optional logos)
2. Classification — protein + RNA → binding probability binding_prob ∈ (0, 1)

Protein conditioning uses VESM-3B (weights live on your GPU server; not shipped here).

Package size ≈ 2.8 GB (three checkpoints).

---

Layout

ProRiboGen_code/
├── README.md                 # this file
├── requirements.txt
├── configs/
│   ├── paths.env.example     # copy → paths.env and edit VESM paths
│   ├── sample_fixed80bp.json
│   └── model.json            # classifier backbone / data config
├── checkpoints/
│   ├── generator.pt          # ~1.2G  RNA generation module
│   ├── backbone.pt           # ~1.2G  MLM backbone for classification
│   └── classifier.pt         # ~477M  binding classifier head + fusion
├── generator/                # Generation module (sample.py + src/ + tokenizer/)
├── classifier/               # Classification module (RnaRealismClassifierV3)
├── protein_encoder/          # AA sequence → VESM H5
├── motif/                    # MEME, HOMER conversion, logo PNGs
├── api/                      # CLI + FastAPI
├── examples/                 # demo protein FASTA
├── scripts/                  # smoke tests
└── workspace/                # runtime outputs (created on demand)

Checkpoint roles

File	Role
generator.pt	Generation only — MaskGIT+ sampling
backbone.pt	Classification — build ESM MLM, then load classifier weights
classifier.pt	Classification — binding / realism score

Do not use backbone.pt for generation. Do not use generator.pt as the classifier backbone (the classifier was trained on the backbone run).

---

Setup (another server)

1. Dependencies

cd ProRiboGen_code
python3 -m venv .venv && source .venv/bin/activate   # optional
pip install -r requirements.txt
# Motif discovery needs MEME Suite on PATH (`meme` command).
# Logos need: pip install logomaker matplotlib pandas

2. Configure VESM

cp configs/paths.env.example configs/paths.env
# edit configs/paths.env

Either:

# A) Standard 17-VESM3 tree
VESM_ROOT=/path/to/17-VESM3
# expects:
#   $VESM_ROOT/models/base/facebook_esm2_t36_3B_UR50D
#   $VESM_ROOT/models/weights/VESM_3B.pth

or:

# B) Explicit paths (overrides A)
VESM_BASE_MODEL_DIR=/path/to/facebook_esm2_t36_3B_UR50D
VESM_WEIGHTS=/path/to/VESM_3B.pth
DEVICE=cuda:0

Checkpoint / tokenizer relative paths in paths.env usually need no change.

3. Smoke tests

bash scripts/demo_generate.sh
bash scripts/demo_score.sh

# With a real protein FASTA from this package:
python api/generate_rna_and_motif.py \
  --protein @examples/HS90A_HUMAN_P07900.fasta \
  --num-sequences 32 --length-bp 80

python api/score_binding.py \
  --protein @examples/HS90A_HUMAN_P07900.fasta \
  --rna AUGCAUGCAUGCAUGCAUGCAUGCAUGCAUGC

---

Module 1 — Generation (RNA + motifs)

python api/generate_rna_and_motif.py \
  --protein "MSKSLYVR..." \
  --p-id DemoRBP \
  --num-sequences 64 \
  --length-bp 80

Or:

python api/generate_rna_and_motif.py --protein @/path/to/protein.fasta --num-sequences 64

Outputs under workspace/job_<p_id>/:

Path	Meaning
*_vesm3b.h5	VESM token embeddings
generated/*.fasta	Generated RNA
motifs/*/meme.txt	MEME raw output
homer_motif_matrix/*.motif	HOMER PWMs
homer_motif_logo/*.png	Sequence logos (ACGU colors below)

Skip MEME/logos: --skip-motif.

Logo colors (same as internal create_motif_logo.ipynb):

Base	Hex
A	#65a455
C	#2e45a4
G	#fda562
U	#d54f3f

Standalone logos:

python motif/plot_homer_logos.py workspace/job_DemoRBP/homer_motif_matrix workspace/job_DemoRBP/homer_motif_logo

Website tip: use num_sequences=64–256 for latency; paper-style runs use 2048.

---

Module 2 — Classification (binding probability)

python api/score_binding.py \
  --protein "MSKSLYVR..." \
  --rna "UGCAUGCGAU..." \
  --p-id DemoRBP

Example JSON:

{
  "p_id": "DemoRBP",
  "rna": "UGCAUGCGAU...",
  "rna_len": 80,
  "logit": 1.23,
  "binding_prob": 0.77
}

Default decision threshold: binding_prob >= 0.5 (same as training eval). Reuse a cached H5 with --protein-h5 to skip re-encoding.

---

HTTP API

uvicorn api.app:app --host 0.0.0.0 --port 8000

curl http://127.0.0.1:8000/health

curl -X POST http://127.0.0.1:8000/v1/generate \
  -H 'Content-Type: application/json' \
  -d '{"protein":"MSK...","p_id":"Demo","num_sequences":32,"run_motif":true}'

curl -X POST http://127.0.0.1:8000/v1/score \
  -H 'Content-Type: application/json' \
  -d '{"protein":"MSK...","rna":"UGCA...","p_id":"Demo"}'

Frontend only needs these two endpoints; GPU work stays on this service.

---

Pipeline

AA sequence
    │
    ├─ VESM-3B ─► H5 [L, 2560]
    │
    ├─ generator.pt (generation) ─► FASTA ─► MEME ─► HOMER ─► PNG logos
    │
    └─ backbone.pt + classifier.pt + RNA ─► binding_prob (classification)

No pretrained 337-protein H5 is required; each request encodes VESM on the fly.

---

Copy to another machine

rsync -avP ProRiboGen_code/ user@host:/path/ProRiboGen_code/
# or
tar -cf ProRiboGen_code.tar ProRiboGen_code

On the new host: install deps → edit configs/paths.env → run the smoke scripts.

---

Hardware / notes

1. Prefer GPU ≥ 24 GB (VESM-3B + 12-layer ProRiboGen). Encode then generate/score can be split across steps.
2. Without MEME, use --skip-motif (sequences only).
3. paths.env is gitignored; ship paths.env.example only.
4. Python ≥ 3.10 recommended (list[str] typing in scripts).