Views
No views yet
1from transformers import AutoTokenizer, BertForPreTraining
2import torch
3
4repo_name = "MM-YY-WW/GoBERT"
5tokenizer = AutoTokenizer.from_pretrained(repo_name, use_fast=False, trust_remote_code=True)
6model = BertForPreTraining.from_pretrained(repo_name)
7
8# Obtain function-level GoBERT Embedding:
9input_sequences = 'GO:0005739 GO:0005783 GO:0005829 GO:0006914 GO:0006915 GO:0006979 GO:0031966 GO:0051560'
10tokenized_input = tokenizer(input_sequences)
11input_tensor = torch.tensor(tokenized_input['input_ids']).unsqueeze(0)
12attention_mask = torch.tensor(tokenized_input['attention_mask']).unsqueeze(0)
13
14model.eval()
15with torch.no_grad():
16 outputs = model(input_ids=input_tensor, attention_mask=attention_mask, output_hidden_states=True)
17 embedding = outputs.hidden_states[-1].squeeze(0).cpu().numpy() 1@inproceedings{miao2025gobert,
2 title={GoBERT: Gene Ontology Graph Informed BERT for Universal Gene Function Prediction},
3 author={Miao, Yuwei and Guo, Yuzhi and Ma, Hehuan and Yan, Jingquan and Jiang, Feng and Liao, Rui and Huang, Junzhou},
4 booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
5 volume={39},
6 number={1},
7 pages={622--630},
8 year={2025},
9 doi={10.1609/aaai.v39i1.32043}
10}