Views
No views yet
1# ========== generate CDS
2import torch
3from model.tokenization_UniBioseq import UBSLMTokenizer
4from model.modeling_UniBioseq import UniBioseqForCausalLM
5from model.UBL_utils import CodonLogitsProcessor
6
7from transformers.generation.logits_process import LogitsProcessorList
8
9tokenizer = UBSLMTokenizer.from_pretrained("lonelycrab88/BiooBang-1.0-HEK293T")
10model = UniBioseqForCausalLM.from_pretrained("lonelycrab88/BiooBang-1.0-HEK293T", device_map='auto')
11
12protein_prompt = "MASSDKQTSPKPPPSPSPLRNSKFCQSNMRILIS"
13input_ids = torch.tensor([tokenizer.encode(input_protein)+[36]]).to(model.device)
14max_length = 4*len(input_protein)+6
15
16logits_processor = LogitsProcessorList()
17logits_processor.append(CodonLogitsProcessor(input_protein, tokenizer, len(input_protein)))
18result = model.generate(input_ids, max_length = max_length, num_beams = 10, logits_processor=logits_processor, low_memory=True, num_return_sequences=1)
19result_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}