Zeroshot Tutorial

This tutorial demonstrates how to use the scLinguist model for zeroshot prediction of protein expression from RNA data without the need for finetuning. The model is pre-trained and can be directly applied to new datasets.

Import necessary packages and define paths for checkpoints and save directory.

from pathlib import Path
from torch.utils.data import DataLoader
import sys
sys.path.append('../../')
from scLinguist.data_loaders.data_loader import scMultiDataset
from scLinguist.model.configuration_hyena import HyenaConfig
from scLinguist.model.model import scTrans
import importlib, sys
sys.modules['model'] = importlib.import_module('scLinguist.model')

ENCODER_CKPT = Path("../../pretrained_model/encoder.ckpt")
DECODER_CKPT = Path("../../pretrained_model/decoder.ckpt")
FINETUNE_CKPT = Path("../../pretrained_model/finetune.ckpt")
SAVE_DIR = Path("../../docs/tutorials/zeroshot_output")
SAVE_DIR.mkdir(exist_ok=True)

Since this is a zeroshot task, we dont need to prepare our dataloaders. We configure the model with the appropriate encoder and decoder checkpoints, and set the mode to “RNA-protein”. The HyenaConfig class is used to define the model configuration parameters such as d_model, emb_dim, max_seq_len, vocab_len, and n_layer.

enc_cfg = HyenaConfig(
    d_model        = 128,
    emb_dim        = 5,
    max_seq_len    = 19202,
    vocab_len      = 19202,
    n_layer        = 1,
    output_hidden_states=False,
)
dec_cfg = HyenaConfig(
    d_model        = 128,
    emb_dim        = 5,
    max_seq_len    = 6427,
    vocab_len      = 6427,
    n_layer        = 1,
    output_hidden_states=False,
)
model = scTrans.load_from_checkpoint(checkpoint_path=FINETUNE_CKPT)
model.encoder_ckpt_path = ENCODER_CKPT
model.decoder_ckpt_path = DECODER_CKPT
model.mode = "RNA-protein"

Zeroshot Inference with the pre-trained model on the test dataset. Use RNA data to predict proteins in ../../docs/tutorials/protein_names.txt

import scanpy as sc
import torch

# only use 10 cells for example
zeroshot_adata = sc.read_h5ad("../../data/test_sample_rna.h5ad")[:10]
zeroshot_rna_tensor = torch.tensor(zeroshot_adata.X.todense(), dtype=torch.float32).cuda()

model.eval().cuda()

with torch.no_grad():
    _, _, protein_pred = model(zeroshot_rna_tensor)

# predict given proteins
target_proteins = [line.strip() for line in open("../../docs/tutorials/protein_names.txt")]

import pandas as pd
prot_map = pd.read_csv("../../docs/tutorials/protein_index_map.csv")
name_to_idx = dict(zip(prot_map["name"], prot_map["index"]))

idx = [name_to_idx[p] for p in target_proteins if p in name_to_idx]

pred_df = pd.DataFrame(
    protein_pred[:, idx].cpu().numpy(),
    columns = target_proteins,
    index   = zeroshot_adata.obs_names,
)
pred_df.to_csv(SAVE_DIR/"predicted_protein_expression.csv")