from transformers import AutoModelForMaskedLM, AutoTokenizer
import torch

model = AutoModelForMaskedLM.from_pretrained("kojima-lab/molcrawl-protein-sequence-bert-large")
tokenizer = AutoTokenizer.from_pretrained("kojima-lab/molcrawl-protein-sequence-bert-large")

# Predict masked amino acid
# Use tokenizer.mask_token instead of hardcoded "[MASK]":
# BERT-style tokenizers vary ("[MASK]", "<mask>", etc.)
if tokenizer.mask_token is None:
    raise ValueError("This tokenizer has no mask_token; masked LM inference is not supported.")
prompt = "MKTAYIAK{MASK}RQISFVKSHFSRQ".replace("{MASK}", tokenizer.mask_token)
inputs = tokenizer(prompt, return_tensors="pt")
mask_index = (inputs["input_ids"] == tokenizer.mask_token_id).nonzero(as_tuple=True)[1]

with torch.no_grad():
    outputs = model(**inputs)
logits = outputs.logits

predicted_token_id = logits[0, mask_index].argmax(dim=-1)
predicted_token = tokenizer.decode(predicted_token_id)
result = prompt.replace(tokenizer.mask_token, predicted_token)
print(f"Predicted: {result}")

Source Code

Training pipeline, configuration files, and data preparation scripts are available in the MolCrawl GitHub repository: https://github.com/mmai-framework-lab/MolCrawl

License

This model is released under the APACHE-2.0 license.

Citation

If you use this model, please cite:

bibtex

@misc{molcrawl_protein_sequence_bert_large,
  title={molcrawl-protein-sequence-bert-large},
  author={{RIKEN}},
  year={2026},
  publisher={{Hugging Face}},
  url={{https://huggingface.co/kojima-lab/molcrawl-protein-sequence-bert-large}}
}

Example Output

End-to-end inference test (downloaded the model from this repo on CPU).

python

import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

REPO_ID = "kojima-lab/molcrawl-protein-sequence-bert-large"
tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
model = AutoModelForMaskedLM.from_pretrained(REPO_ID)
model.eval()

sequence = "MKTAYIAKQRQISFVK<mask>SHFSRQ"
inputs = tokenizer(sequence, return_tensors="pt")
mask_index = (inputs["input_ids"][0] == tokenizer.mask_token_id).nonzero(as_tuple=True)[0]

with torch.no_grad():
    outputs = model(**inputs)

predicted_id = outputs.logits[0, mask_index].argmax(dim=-1)
predicted_aa = tokenizer.convert_ids_to_tokens(predicted_id.tolist())[0]
print(f"Predicted amino acid at mask: {predicted_aa}")
# => Predicted amino acid at mask: S
# (top-5 candidates: ['S', 'P', 'E', 'L', 'A'])

Note on checkpoint selection: the full 60,000-step pretrain run overfit after roughly step 22,000 — eval_loss reached its minimum (2.5938) around step 21,900 and slowly degraded to 2.7526 by the end. The uploaded weights come from checkpoint-24000 (best per trainer_state.best_metric = 2.597), not from the final checkpoint.