Views
No views yet
1import torch
2
3# ========== Set device
4device = "cuda:0"
5
6# ========== Prepare Data
7data = [
8 ("Protein", "MASSDKQTSPKPPPSPSPLRNSKFCQSNMRILIS"),
9 ("RNA", "ATGGCGTCTAGTGATAAACAAACAAGCCCAAAGCCTCCTCCTTCACCGTCTCCTCTCCGTAATT")
10]
11
12# ========== BiooBang Model
13from model.modeling_UniBioseq import UniBioseqForEmbedding
14from model.tokenization_UniBioseq import UBSLMTokenizer
15model = UniBioseqForEmbedding.from_pretrained("lonelycrab88/BiooBang-1.0")
16tokenizer = UBSLMTokenizer.from_pretrained("lonelycrab88/BiooBang-1.0")
17model.eval()
18model.to(device)
19
20# ========== get Embeddings
21embeddings = {}
22hidden_states = {}
23for name,input_seq in data:
24 input_ids = tokenizer(input_seq, return_tensors="pt")['input_ids'].to(device)
25 with torch.no_grad():
26 # get sequence embedding
27 embeddings[name] = model(input_ids).logits
28 # get last hidden states (token embeddings)
29 hidden_states[name] = model(input_ids).hidden_states[:,1:-1,:]
30
31# ========== generate CDS
32from transformers.generation.logits_process import LogitsProcessorList
33from model.UBL_utils import CodonLogitsProcessor
34from model.modeling_UniBioseq import UniBioseqForCausalLM
35tokenizer = UBSLMTokenizer.from_pretrained("lonelycrab88/BiooBang-1.0")
36model = UniBioseqForCausalLM.from_pretrained("lonelycrab88/BiooBang-1.0", device_map='auto')
37
38protein_prompt = "MASSDKQTSPKPPPSPSPLRNSKFCQSNMRILIS"
39input_ids = torch.tensor([tokenizer.encode(input_protein)+[36]]).to(model.device)
40max_length = 4*len(input_protein)+6
41
42logits_processor = LogitsProcessorList()
43logits_processor.append(CodonLogitsProcessor(input_protein, tokenizer, len(input_protein)))
44result = model.generate(input_ids, max_length = max_length, num_beams = 10, logits_processor=logits_processor, low_memory=True, num_return_sequences=1)
45result_CDS_tok = tokenizer.decode(result[0][len(input_protein)+3:].tolist()).replace(" ","").upper()1@article {Zhao2024.10.24.620004,
2 author = {Zhao, Heng-Rui and Cheng, Meng-Ting and Zhu, Jinhua and Wang, Hao and Yang, Xiang-Rui and Wang, Bo and Sun, Yuan-Xin and Fang, Ming-Hao and Chen, Enhong and Li, Houqiang and Han, Shu-Jing and Chen, Yuxing and Zhou, Cong-Zhao},
3 title = {Integration of protein and coding sequences enables mutual augmentation of the language model},
4 elocation-id = {2024.10.24.620004},
5 year = {2024},
6 doi = {10.1101/2024.10.24.620004},
7 publisher = {Cold Spring Harbor Laboratory},
8 URL = {https://www.biorxiv.org/content/early/2024/10/29/2024.10.24.620004},
9 eprint = {https://www.biorxiv.org/content/early/2024/10/29/2024.10.24.620004.full.pdf},
10 journal = {bioRxiv}
11}