Source code for TSARC.process

import os
import torch
import random
import numpy as np
from .tsarp import CnnGru


def seed_torch(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


def oneHot(stringInput):
    matrix = np.zeros((4, len(stringInput)))
    seq = stringInput.upper()
    for pos in range(len(seq)):
        if seq[pos] == 'A':
            matrix[0, pos] = 1
        elif seq[pos] == 'C':
            matrix[1, pos] = 1
        elif seq[pos] == 'G':
            matrix[2, pos] = 1
        elif seq[pos] == 'T':
            matrix[3, pos] = 1
        else:
            continue
    return matrix


def predict(stringInput):
    seed_torch(seed=12)
    model = CnnGru(256, 128)
    data = torch.Tensor(oneHot(stringInput))
    data = torch.unsqueeze(data, 0)
    model.load_state_dict(torch.load('TSARC/saved_model.h5', map_location ='cpu'))
    model.eval()
    with torch.no_grad():
        target = model(data).squeeze(1)
        prob = torch.sigmoid(target)
        pred = target.gt(0.5).float()
    label = pred.numpy()[0]
    prob = prob.numpy()[0]
    p = str(round(prob, 4))
    if label > 0.5:
        return f'TSAR signal, probability: {p}'
    else:
        return f'Not TSAR signal, probability: {p}'