Views
No views yet
1import sys
2import numpy as np
3import torch
4from transformers import AutoTokenizer, AutoModelForSequenceClassification
5
6__authors__ = ["Kazuki Nakamae"]
7__version__ = "1.0.0"
8
9def pred_rna_offtarget(dna, model_dir):
10 try:
11 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
12 tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
13 model = AutoModelForSequenceClassification.from_pretrained(model_dir, trust_remote_code=True).to(device)
14 except Exception as e:
15 print(f"Error loading model from {model_dir}: {e}")
16 sys.exit(1)
17
18 inputs = tokenizer(dna, return_tensors='pt')
19 model.eval()
20 with torch.no_grad():
21 outputs = model(
22 inputs["input_ids"].to(device),
23 inputs["attention_mask"].to(device),
24 )
25 print("[Negative, Positive]")
26 print(outputs.logits)
27 y_preds = np.argmax(outputs.logits.to('cpu').detach().numpy().copy(), axis=1)
28
29 def id2label(x):
30 return model.config.id2label[x]
31 y_dash = [id2label(x) for x in y_preds]
32 print("Result:")
33 print(y_dash)
34 # LABEL_0: Not RNA-offtarget / LABEL_1: RNA-offtarget
35 return (dna, y_dash)
36
37def print_usage():
38 print(f"Usage: {sys.argv[0]} <input DNA sequence> <DNABERT-2 model directory>")
39 print("Options:")
40 print(" -h, --help Show this help message and exit")
41 print(" -v, --version Show version information and exit")
42
43def print_version():
44 print(f"{sys.argv[0]} version {__version__}")
45 print("Authors:", ", ".join(__authors__))
46
47if __name__ == "__main__":
48 if len(sys.argv) != 3:
49 if len(sys.argv) == 2 and sys.argv[1] in ("-h", "--help"):
50 print_usage()
51 sys.exit(0)
52 elif len(sys.argv) == 2 and sys.argv[1] in ("-v", "--version"):
53 print_version()
54 sys.exit(0)
55 else:
56 print_usage()
57 sys.exit(1)
58
59 dna = sys.argv[1]
60 model_dir = sys.argv[2]
61
62 pred_rna_offtarget(dna, model_dir)1$ python pred_rna_offtarget.py GGCAGGGCTGGGGAAGCTTACTGTGTCCAAGAGCCTGCTG KazukiNakamae/STLmodel;
2[Negative, Positive]
3tensor([[-1.6383, 1.4502]])
4Result:
5['LABEL_1']
6$ python pred_rna_offtarget.py GTCATCTAACAAAAATATTCCGTTGCAGGAAAAGCAAGCT KazukiNakamae/STLmodel;
7[Negative, Positive]
8tensor([[ 0.5446, -0.5105]])
9Result:
10['LABEL_0']