Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModel
3from transformers.models.bert.configuration_bert import BertConfig
4
5config = BertConfig.from_pretrained("YYLY66/mRNABERT")
6tokenizer = AutoTokenizer.from_pretrained("YYLY66/mRNABERT")
7model = AutoModel.from_pretrained("YYLY66/mRNABERT", trust_remote_code=True, config=config)1seq = ["A T C G G A GGG CCC TTT",
2 "A T C G",
3 "TTT CCC GAC ATG"] #Separate the sequences with spaces.
4
5encoding = tokenizer.batch_encode_plus(seq, add_special_tokens=True, padding='longest', return_tensors="pt")
6
7input_ids = encoding['input_ids']
8attention_mask = encoding['attention_mask']
9
10output = model(input_ids=input_ids, attention_mask=attention_mask)
11last_hidden_state = output[0]
12
13attention_mask = attention_mask.unsqueeze(-1).expand_as(last_hidden_state) # Shape : [batch_size, seq_length, hidden_size]
14
15# Sum embeddings along the batch dimension
16sum_embeddings = torch.sum(last_hidden_state * attention_mask, dim=1)
17
18# Also sum the masks along the batch dimension
19sum_masks = attention_mask.sum(1)
20
21# Compute mean embedding.
22mean_embedding = sum_embeddings / sum_masks #Shape:[batch_size, hidden_size]
23@article{xiong2025mrnabert,
title={mRNABERT: advancing mRNA sequence design with a universal language model and comprehensive dataset},
author={Xiong, Ying and Wang, Aowen and Kang, Yu and Shen, Chao and Hsieh, Chang-Yu and Hou, Tingjun},
journal={Nature Communications},
volume={16},
number={1},
pages={10371},
year={2025},
publisher={Nature Publishing Group UK London},
}