Views
No views yet
1from transformers import AutoModelForMaskedLM, AutoTokenizer
2import torch
3
4model = AutoModelForMaskedLM.from_pretrained("kojima-lab/molcrawl-genome-sequence-dnabert2-medium")
5tokenizer = AutoTokenizer.from_pretrained("kojima-lab/molcrawl-genome-sequence-dnabert2-medium")
6
7# Predict masked DNA token
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 = "ATCGATCG{MASK}ATCGATCG".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, skip_special_tokens=True)
22result = prompt.replace(tokenizer.mask_token, predicted_token)
23print(f"Predicted: {result}")
241@misc{molcrawl_genome_sequence_dnabert2_medium,
2 title={molcrawl-genome-sequence-dnabert2-medium},
3 author={{RIKEN}},
4 year={2026},
5 publisher={{Hugging Face}},
6 url={{https://huggingface.co/kojima-lab/molcrawl-genome-sequence-dnabert2-medium}}
7}