example.py

"""
Minimal example: embed an RNA sequence with RCLM.

Usage:
    python example.py [cpu|cuda]
"""
import os
import sys

import numpy as np
import torch

from rclm import Model

device = sys.argv[1] if len(sys.argv) > 1 else "cpu"
expdir = os.path.dirname(os.path.abspath(__file__))

# Same alphabet convention used in DRfold2/cfg_95/EvoMSA2XYZ.py:rnalm_alphas
NT2IDX = {"A": 0, "G": 1, "C": 2, "U": 3,
          "a": 0, "g": 1, "c": 2, "u": 3,
          "T": 3, "t": 3, "-": 4}

lmcfg = {
    "s_in_dim": 5,
    "z_in_dim": 2,
    "s_dim": 512,
    "z_dim": 128,
    "N_elayers": 18,
}


def load_model(device="cpu"):
    model = Model.RNA2nd(lmcfg)
    weight_path = os.path.join(expdir, "weight", "epoch_67000")
    model.load_state_dict(torch.load(weight_path, map_location="cpu"), strict=False)
    model.to(device)
    model.eval()
    return model


def embed(model, seq, device="cpu"):
    L = len(seq)
    aa = np.eye(5)[[NT2IDX.get(c, NT2IDX["-"]) for c in seq]]
    # seq_idx is 1-based, matching DRfold2/cfg_95/test_modeldir.py:data_collect
    seq_idx = np.arange(L) + 1

    in_dict = {
        "aa": torch.FloatTensor(aa).to(device),
        "idx": torch.LongTensor(seq_idx).to(device),
        "mask": torch.zeros(L).to(device),
    }
    with torch.no_grad():
        s, z = model.embedding(in_dict)
    return s, z


if __name__ == "__main__":
    seq = "GGGCUAUUAGCUCAGUUGGUAGAGCCCUGGAUUGUGAUUCCAGUUGUCGUGGGUUCGAAUCCCAUUAGCCCCA"

    rnalm = load_model(device)
    s, z = embed(rnalm, seq, device)

    print(f"sequence length: {len(seq)}")
    print(f"s (per-residue) shape: {tuple(s.shape)}")   # L x 512
    print(f"z (pairwise)     shape: {tuple(z.shape)}")   # L x L x 128