Views
No views yet
1from transformers import AutoModelForMaskedLM, AutoTokenizer
2import torch
3
4model = AutoModelForMaskedLM.from_pretrained("kojima-lab/molcrawl-protein-sequence-bert-small")
5tokenizer = AutoTokenizer.from_pretrained("kojima-lab/molcrawl-protein-sequence-bert-small")
6
7# Predict masked amino acid
8# Use tokenizer.mask_token instead of hardcoded "[MASK]":
9# BERT-style tokenizers vary ("[MASK]", "<mask>", etc.)
10if tokenizer.mask_token is None:
11 raise ValueError("This tokenizer has no mask_token; masked LM inference is not supported.")
12prompt = "MKTAYIAK{MASK}RQISFVKSHFSRQ".replace("{MASK}", tokenizer.mask_token)
13inputs = tokenizer(prompt, return_tensors="pt")
14mask_index = (inputs["input_ids"] == tokenizer.mask_token_id).nonzero(as_tuple=True)[1]
15
16with torch.no_grad():
17 outputs = model(**inputs)
18logits = outputs.logits
19
20predicted_token_id = logits[0, mask_index].argmax(dim=-1)
21predicted_token = tokenizer.decode(predicted_token_id)
22result = prompt.replace(tokenizer.mask_token, predicted_token)
23print(f"Predicted: {result}")
241@misc{molcrawl_protein_sequence_bert_small,
2 title={molcrawl-protein-sequence-bert-small},
3 author={{RIKEN}},
4 year={2026},
5 publisher={{Hugging Face}},
6 url={{https://huggingface.co/kojima-lab/molcrawl-protein-sequence-bert-small}}
7}