| import numba |
| import tempfile |
| numba.core.config.CACHE_DIR = tempfile.mkdtemp() |
| import json |
| import numpy as np |
| import pandas as pd |
| import torch |
| import seqpro |
| import streamlit as st |
| from model.module import CNN_Encoder_Model |
|
|
|
|
| def device_detect() -> tuple[torch.device, type]: |
| """ |
| Detect whether the device is a CPU or GPU. |
| Returns a tuple of (device, Tensor type) |
| """ |
| if torch.cuda.is_available(): |
| device = torch.device("cuda") |
| Tensor = torch.cuda.FloatTensor |
| else: |
| device = torch.device("cpu") |
| Tensor = torch.FloatTensor |
| return device, Tensor |
|
|
|
|
| def model_loading(config_local: str, epoch_local: str) -> CNN_Encoder_Model: |
| """ |
| Load model parameters and return the model in evaluation mode. |
| """ |
| device, _ = device_detect() |
| with open(config_local, "r") as file: |
| model_config = json.load(file) |
| INPUTSHAPE = { |
| "gene": 1, |
| "reg": (4, 5000), |
| "cds": (4, 4500), |
| "trans": (1024,), |
| "rna": (1,), |
| "expression": [1], |
| } |
| model = CNN_Encoder_Model(model_config=model_config, input_shapes=INPUTSHAPE) |
| model.to(device) |
| checkpoint = torch.load(epoch_local, map_location=device) |
| model.load_state_dict(checkpoint["model_state_dict"]) |
| model.eval() |
| return model |
|
|
|
|
| def format_sequence(utr_5, cds, utr_3) -> tuple[np.ndarray, np.ndarray]: |
| utr_formatted = list() |
| cds_formatted = list() |
| for i in range(len(utr_5)): |
| cds_formatted.append(cds[i][:4500].ljust(4500, "N")) |
| utr_formatted.append( |
| utr_5[i][-1000:].rjust(1000, "N") + utr_3[i][-4000:].ljust(4000, "N") |
| ) |
| return seqpro.ohe( |
| np.array(utr_formatted), alphabet=seqpro.alphabets.DNA |
| ), seqpro.ohe(np.array(cds_formatted), alphabet=seqpro.alphabets.DNA) |
|
|
|
|
| def format_mrna_env(envs_local: str): |
| mrna_env = pd.read_csv(envs_local, header=None, names=["value"])["value"].to_numpy() |
| return mrna_env |
|
|
|
|
| def predict( |
| model, |
| utr: np.ndarray, |
| cds: np.ndarray, |
| env: np.ndarray, |
| rna_count=6.3, |
| batch_size=256, |
| ) -> np.ndarray: |
| """ |
| Predicate by model. |
| """ |
| device, Tensor = device_detect() |
| res = [] |
| with torch.no_grad(): |
| for line in range(0, len(utr), batch_size): |
| count = utr[line : line + batch_size].shape[0] |
| out = ( |
| model( |
| Tensor(utr[line : line + batch_size]) |
| .to(device) |
| .transpose(2, 1) |
| .float(), |
| Tensor(cds[line : line + batch_size]) |
| .to(device) |
| .transpose(2, 1) |
| .float(), |
| Tensor(env).to(device).repeat(count, 1), |
| Tensor([rna_count]).to(device).reshape(1, 1).repeat(count, 1), |
| )[:, 0] |
| .to("cpu") |
| .detach() |
| .numpy() |
| ) |
| res.append(out) |
| out = np.concatenate(res).reshape(-1) |
| |
| return out |
|
|
