Create utils.py

utils.py

@@ -0,0 +1,199 @@
+import numpy as np
+import random
+import math
+
+from sklearn.metrics import *
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+import pickle
+
+    
+def word2idx(word, words):
+    if word in words.keys():
+        return int(words[word])
+    
+    return 0
+
+def pad_seq(dataset, max_len):
+    output = []
+    for seq in dataset:
+        pad = np.zeros(max_len)
+        pad[:len(seq)] = seq
+        output.append(pad)
+        
+    return np.array(output)
+
+def str2bool(seq):
+    out = []
+    for s in seq:
+        if s == "positive": 
+            out.append(1)
+        elif s == "negative": 
+            out.append(0)
+            
+    return np.array(out)
+    
+class API_Dataset(Dataset):
+    def __init__(self, apta, esm_prot, y, apta_attn_mask, prot_attn_mask):
+        super(Dataset, self).__init__()
+        
+        self.apta = np.array(apta, dtype=np.int64)
+        self.esm_prot = np.array(esm_prot, dtype=np.int64)
+        self.y = np.array(y, dtype=np.int64)
+        self.apta_attn_mask = np.array(apta_attn_mask)
+        self.prot_attn_mask = np.array(prot_attn_mask)
+        self.len = len(self.apta)
+
+    def __len__(self):
+        return self.len
+
+    def __getitem__(self, index):
+        return torch.tensor(self.apta[index], dtype=torch.int64), torch.tensor(self.esm_prot[index], dtype=torch.int64), torch.tensor(self.y[index], dtype=torch.int64), torch.tensor(self.apta_attn_mask[index], dtype=torch.int64), torch.tensor(self.prot_attn_mask[index], dtype=torch.int64)
+
+def find_opt_threshold(target, pred):
+    result = 0
+    best = 0
+    
+    for i in range(0, 1000):
+        pred_threshold = np.where(pred > i/1000, 1, 0)
+        now = f1_score(target, pred_threshold)
+        if now > best:
+            result = i/1000
+            best = now
+            
+    return result
+
+def argument_seqset(seqset):
+    arg_seqset = []
+    for s, ss in seqset:
+        arg_seqset.append([s, ss]) 
+
+        arg_seqset.append([s[::-1], ss[::-1]])
+
+    return arg_seqset
+
+def augment_apis(apta, prot, ys):
+    aug_apta = []
+    aug_prot = []
+    aug_y = []
+    for a, p, y in zip(apta, prot, ys):
+        aug_apta.append(a) 
+        aug_prot.append(p)
+        aug_y.append(y)
+
+        aug_apta.append(a[::-1]) 
+        aug_prot.append(p)
+        aug_y.append(y)
+
+    return np.array(aug_apta), np.array(aug_prot), np.array(aug_y)
+
+
+
+def load_data_source(filepath):
+    with open(filepath,"rb") as fr:
+        dataset = pickle.load(fr)
+        dataset_train = np.array(dataset[dataset["dataset"]=="training dataset"])
+        dataset_test = np.array(dataset[dataset["dataset"]=="test dataset"])
+        dataset_bench = np.array(dataset[dataset['dataset']=='benchmark dataset'])
+
+        return dataset_train, dataset_test, dataset_bench
+
+
+def get_dataset(filepath, prot_max_len, n_prot_vocabs, prot_words):
+    dataset_train, dataset_test, dataset_bench = load_data_source(filepath)
+
+    
+    arg_apta, arg_prot, arg_y = augment_apis(dataset_train[:, 0], dataset_train[:, 1], dataset_train[:, 2])
+    datasets_train = [rna2vec(arg_apta), tokenize_sequences(arg_prot, prot_max_len, n_prot_vocabs, prot_words), str2bool(arg_y)]
+    datasets_test = [rna2vec(dataset_test[:, 0]), tokenize_sequences(dataset_test[:, 1], prot_max_len, n_prot_vocabs, prot_words), str2bool(dataset_test[:, 2])]
+    datasets_bench = [rna2vec(dataset_bench[:, 0]), tokenize_sequences(dataset_bench[:, 1], prot_max_len, n_prot_vocabs, prot_words), str2bool(dataset_bench[:, 2])]
+
+    return datasets_train, datasets_test, datasets_bench
+
+
+def get_esm_dataset(filepath, batch_converter, alphabet):
+    dataset_train, dataset_test, dataset_bench = load_data_source(filepath)
+
+    # arg_apta, arg_prot, arg_y = augment_apis(dataset_train[:, 0], dataset_train[:, 1], dataset_train[:, 2])
+    # arg_prot is a np.array of strings (4640,) -> convert this to np.array of size (2x4640) where first row is a label
+
+    arg_apta, arg_prot, arg_y = dataset_train[:, 0], dataset_train[:, 1], dataset_train[:, 2]
+    arg_apta, arg_prot, arg_y = augment_apis(arg_apta, arg_prot, arg_y)
+
+    train_inputs = [(i, j) for i, j in zip(arg_y, arg_prot)]
+    _, _, prot_tokens = batch_converter(train_inputs)
+    datasets_train = [rna2vec(arg_apta), prot_tokens, str2bool(arg_y)]
+
+    test_inputs = [(i, j) for i, j in enumerate(dataset_test[:, 1])]
+    _, _, test_prot_tokens = batch_converter(test_inputs)
+    datasets_test = [rna2vec(dataset_test[:, 0]), test_prot_tokens, str2bool(dataset_test[:, 2])]
+
+    bench_inputs = [(i, j) for i, j in enumerate(dataset_bench[:, 1])]
+    _, _, bench_prot_tokens = batch_converter(bench_inputs)
+    # truncating
+    bench_prot_tokenized = bench_prot_tokens[:, :1678]
+    # padding
+    prot_ex = torch.ones((bench_prot_tokenized.shape[0], 1678), dtype=torch.int64)*alphabet.padding_idx
+    prot_ex[:, :bench_prot_tokenized.shape[1]] = bench_prot_tokenized
+    datasets_bench = [rna2vec(dataset_bench[:, 0]), prot_ex, str2bool(dataset_bench[:, 2])]
+
+    return datasets_train, datasets_test, datasets_bench
+
+def get_nt_esm_dataset(filepath, nt_tokenizer, batch_converter, alphabet):
+    dataset_train, dataset_test, dataset_bench = load_data_source(filepath)
+
+    arg_apta, arg_prot, arg_y = augment_apis(dataset_train[:, 0], dataset_train[:, 1], dataset_train[:, 2])
+    # arg_prot is a np.array of strings (4640,) -> convert this to np.array of size (2x4640) where first row is a label
+    max_length = 275#nt_tokenizer.model_max_length
+
+    train_inputs = [(i, j) for i, j in zip(arg_y, arg_prot)]
+    _, _, prot_tokens = batch_converter(train_inputs)
+    apta_toks = nt_tokenizer.batch_encode_plus(arg_apta, return_tensors='pt', padding='max_length', max_length=max_length)['input_ids']
+    apta_attention_mask = apta_toks != nt_tokenizer.pad_token_id
+    prot_attention_mask = prot_tokens != alphabet.padding_idx
+    # datasets_train = [apta_toks, prot_tokens, str2bool(arg_y)]
+    datasets_train = [apta_toks, prot_tokens, str2bool(arg_y), apta_attention_mask, prot_attention_mask]
+
+    test_inputs = [(i, j) for i, j in enumerate(dataset_test[:, 1])]
+    _, _, test_prot_tokens = batch_converter(test_inputs)
+    prot_ex = torch.ones((test_prot_tokens.shape[0], 1680), dtype=torch.int64)*alphabet.padding_idx
+    prot_ex[:, :test_prot_tokens.shape[1]] = test_prot_tokens
+    apta_toks = nt_tokenizer.batch_encode_plus(dataset_test[:, 0], return_tensors='pt', padding='max_length', max_length=max_length)['input_ids']
+    apta_attention_mask = apta_toks != nt_tokenizer.pad_token_id
+    prot_attention_mask = prot_ex != alphabet.padding_idx
+    datasets_test = [apta_toks, prot_ex, str2bool(dataset_test[:, 2]), apta_attention_mask, prot_attention_mask]
+
+    bench_inputs = [(i, j) for i, j in enumerate(dataset_bench[:, 1])]
+    _, _, bench_prot_tokens = batch_converter(bench_inputs)
+    # padding
+    prot_ex = torch.ones((bench_prot_tokens.shape[0], 1680), dtype=torch.int64)*alphabet.padding_idx
+    prot_ex[:, :bench_prot_tokens.shape[1]] = bench_prot_tokens
+    apta_toks = nt_tokenizer.batch_encode_plus(dataset_bench[:, 0], return_tensors='pt', padding='max_length', max_length=max_length)['input_ids']
+    apta_attention_mask = apta_toks != nt_tokenizer.pad_token_id
+    prot_attention_mask = prot_ex != alphabet.padding_idx
+    datasets_bench = [apta_toks, prot_ex, str2bool(dataset_bench[:, 2]), apta_attention_mask, prot_attention_mask]
+
+    return datasets_train, datasets_test, datasets_bench
+
+def get_scores(target, pred):
+    threshold = find_opt_threshold(target, pred)
+    pred_threshold = np.where(pred > threshold, 1, 0)
+    acc = accuracy_score(target, pred_threshold)
+    roc_auc = roc_auc_score(target, pred)
+    mcc = matthews_corrcoef(target, pred_threshold)
+    f1 = f1_score(target, pred_threshold)
+    pr_auc = average_precision_score(target, pred)
+    cls_report = classification_report(target, pred_threshold)
+    scores = {
+        'threshold': threshold,
+        'acc': acc,
+        'roc_auc': roc_auc,
+        'mcc': mcc, 
+        'f1': f1, 
+        'pr_auc': pr_auc,
+        'cls_report': cls_report
+    }
+    return scores