molcrawl-rna-bert-small

⚠️ Input format: this model uses a Geneformer-style WordLevel tokenizer whose vocabulary is ENSEMBL gene IDs only (e.g. ENSG00000000003). Plain nucleotide strings (AUGCAUGC...) or free text will NOT tokenize correctly — they collapse to one [UNK] and the MLM head will return an ENSG-id at that mask position by design. See the Example Output below for the correct list-of-tokens calling convention.

molcrawl-rna-bert-small

Model Description

BERT-small foundation model (~107M parameters, 12-layer / 768-hidden / 12-head topology, sized to match GPT-2-small) pre-trained on RNA gene expression sequences from the MolCrawl dataset.

• Model Type: bert
• Data Type: RNA
• Training Date: 2026-04-24

Usage

This tokenizer is a WordLevel vocab of ENSEMBL gene IDs only (no text-mode pre-tokenizer). Feed the input as a Python list of gene-ID tokens (one position may be tokenizer.mask_token) and resolve to IDs through convert_tokens_to_ids — calling tokenizer("AUGC...") on raw nucleotide strings collapses to [UNK] and the MLM head returns a uniform-prior prediction.

import torch
from transformers import AutoModelForMaskedLM, AutoTokenizer

REPO_ID = "kojima-lab/molcrawl-rna-bert-small"
tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
model = AutoModelForMaskedLM.from_pretrained(REPO_ID)
model.eval()

# Gene-expression masked prediction (ENSEMBL gene IDs in rank order).
genes = [
    "ENSG00000000003",
    "ENSG00000000005",
    tokenizer.mask_token,   # mask one position to predict
    "ENSG00000001167",
    "ENSG00000002586",
]
ids = tokenizer.convert_tokens_to_ids(genes)
input_ids = torch.tensor([ids])
mask_index = (input_ids == tokenizer.mask_token_id).nonzero(as_tuple=True)[1]

with torch.no_grad():
    outputs = model(input_ids=input_ids)

predicted_id = outputs.logits[0, mask_index].argmax(dim=-1)
predicted_gene = tokenizer.convert_ids_to_tokens(predicted_id.tolist())[0]
print(f"Predicted gene at mask: {predicted_gene}")

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:

@misc{molcrawl_rna_bert_small,
  title={molcrawl-rna-bert-small},
  author={{RIKEN}},
  year={2026},
  publisher={{Hugging Face}},
  url={{https://huggingface.co/kojima-lab/molcrawl-rna-bert-small}}
}

Example Output

Inference test (CPU). Important: this tokenizer uses a WordLevel gene-ID vocabulary without a text-mode pre_tokenizer, so feed input as a list of ENSEMBL IDs (and the mask token) directly into convert_tokens_to_ids, not as a single text string.

import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

REPO_ID = "kojima-lab/molcrawl-rna-bert-small"
tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
model = AutoModelForMaskedLM.from_pretrained(REPO_ID)
model.eval()

# Gene-expression masked prediction (ENSEMBL gene IDs in rank order)
genes = [
    "ENSG00000000003",
    "ENSG00000000005",
    tokenizer.mask_token,   # mask one position to predict
    "ENSG00000001167",
    "ENSG00000002586",
]
ids = tokenizer.convert_tokens_to_ids(genes)
input_ids = torch.tensor([ids])
mask_index = (input_ids == tokenizer.mask_token_id).nonzero(as_tuple=True)[1]

with torch.no_grad():
    outputs = model(input_ids=input_ids)

predicted_id = outputs.logits[0, mask_index].argmax(dim=-1)
predicted_gene = tokenizer.convert_ids_to_tokens(predicted_id.tolist())[0]
print(f"Predicted gene at mask: {predicted_gene}")
# => Predicted gene at mask: ENSG00000164972