MolCrawl is a pipeline framework to construct biological multimodal foundational models designed to handle chemical and life science data in a unified manner.

---

Detailed Documentation

For comprehensive documentation including environment setup details, dataset guides, training configuration, evaluation procedures, experiment tracking, and more, see the Documentation Index.

---

Pre-trained Models

Pre-trained model checkpoints for all modalities (Genome Sequence, Protein Sequence, RNA, Compounds, and Molecule-related Natural Language) are publicly available on Hugging Face:

https://huggingface.co/collections/kojima-lab/molcrawl

You can use these checkpoints directly for inference or as starting points for fine-tuning, without running the full data preparation and training pipeline described below.

---

Quick Start (Minimum Steps)

Follow these steps after cloning the repository.

---

0. Clone Repository

git clone https://github.com/mmai-framework-lab/MolCrawl.git
cd MolCrawl

---

1. Create Data Directory

Create the directory that will store all processed data and caches.

mkdir -p learning_source

---

2. Set Environment Variable (Required)

export LEARNING_SOURCE_DIR="$(pwd)/learning_source"

To make it permanent (bash example):

echo 'export LEARNING_SOURCE_DIR="$(pwd)/learning_source"' >> ~/.bashrc

Make sure this directory has at least 100GB of free space. Hugging Face cache will be stored inside ${LEARNING_SOURCE_DIR}/.cache/huggingface/.

---

3. Create Environment (Anaconda not allowed)

conda config --remove channels defaults
conda config --add channels conda-forge
conda config --set channel_priority strict

conda env create --name "molcrawl" --file=environment.yaml
conda activate molcrawl
pip install --no-build-isolation -e .

Offline / restricted-network nodes (optional)

The molecule_nat_lang preparation loads the GPT-2 tokenizer via transformers.from_pretrained("gpt2"), which needs to reach Hugging Face at runtime. On nodes that cannot reach huggingface.co this fails with a connection error. Pre-download the tokenizer once into the directory the workflows auto-detect (assets/tokenizers/gpt2):

huggingface-cli download openai-community/gpt2 \
  config.json vocab.json merges.txt tokenizer.json tokenizer_config.json \
  --local-dir assets/tokenizers/gpt2

The 01-/02-molecule_nat_lang workflow scripts export GPT2_TOKENIZER_DIR to this path automatically when it exists; otherwise set it manually.

---

4. Preprocess Raw Data (Download + Tokenize)

python -m molcrawl.preparation.preparation_script_genome_sequence assets/configs/genome_sequence.yaml
python -m molcrawl.preparation.preparation_script_protein_sequence assets/configs/protein_sequence.yaml
python -m molcrawl.preparation.preparation_script_rna assets/configs/rna.yaml
python -m molcrawl.preparation.preparation_script_compounds assets/configs/compounds.yaml
python -m molcrawl.preparation.preparation_script_molecule_nat_lang assets/configs/molecule_nat_lang.yaml

---

5. Build GPT-2/BERT Training Dataset

python molcrawl/data/genome_sequence/dataset/prepare_gpt2.py assets/configs/genome_sequence.yaml
python molcrawl/data/protein_sequence/dataset/prepare_gpt2.py assets/configs/protein_sequence.yaml
python molcrawl/data/rna/dataset/prepare_gpt2.py assets/configs/rna.yaml
python molcrawl/data/compounds/dataset/prepare_gpt2.py assets/configs/compounds.yaml
python molcrawl/data/molecule_nat_lang/dataset/prepare_gpt2.py assets/configs/molecule_nat_lang.yaml

---

6. Training

GPT-2 (Decoder)

Small:

python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/genome_sequence/gpt2_small.py
python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/protein_sequence/gpt2_small.py
python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/rna/gpt2_small.py
python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/compounds/gpt2_small.py
python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/molecule_nat_lang/gpt2_small.py

Medium:

python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/genome_sequence/gpt2_medium.py
python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/protein_sequence/gpt2_medium.py
python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/rna/gpt2_medium.py
python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/compounds/gpt2_medium.py
python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/molecule_nat_lang/gpt2_medium.py

Large:

python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/genome_sequence/gpt2_large.py
python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/protein_sequence/gpt2_large.py
python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/rna/gpt2_large.py
python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/compounds/gpt2_large.py
python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/molecule_nat_lang/gpt2_large.py

Extra-Large (XL):

python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/genome_sequence/gpt2_xl.py
python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/protein_sequence/gpt2_xl.py
python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/rna/gpt2_xl.py
python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/compounds/gpt2_xl.py
python molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/molecule_nat_lang/gpt2_xl.py

---

BERT (Encoder)

python molcrawl/models/bert/main.py molcrawl/tasks/pretrain/configs/genome_sequence/bert_small.py
python molcrawl/models/bert/main.py molcrawl/tasks/pretrain/configs/protein_sequence/bert_small.py
python molcrawl/models/bert/main.py molcrawl/tasks/pretrain/configs/rna/bert_small.py
python molcrawl/models/bert/main.py molcrawl/tasks/pretrain/configs/compounds/bert_small.py
python molcrawl/models/bert/main.py molcrawl/tasks/pretrain/configs/molecule_nat_lang/bert_small.py

Make sure the GPT-2 dataset preparation step (prepare_gpt2.py) has been completed before BERT training, since BERT uses the same prepared dataset format.

---

Multi-GPU (Single Node, for GPT-2 and BERT)

Both GPT-2 and BERT support Distributed Data Parallel (DDP) via torchrun (use GPU 0 and 2 on a 3-GPU machine).

# GPT2
CUDA_VISIBLE_DEVICES=0,2 \
torchrun --standalone --nproc_per_node=2 \
molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/genome_sequence/gpt2_small.py

CUDA_VISIBLE_DEVICES=0,2 \
torchrun --standalone --nproc_per_node=2 \
molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/protein_sequence/gpt2_small.py

CUDA_VISIBLE_DEVICES=0,2 \
torchrun --standalone --nproc_per_node=2 \
molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/rna/gpt2_small.py

CUDA_VISIBLE_DEVICES=0,2 \
torchrun --standalone --nproc_per_node=2 \
molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/compounds/gpt2_small.py

CUDA_VISIBLE_DEVICES=0,2 \
torchrun --standalone --nproc_per_node=2 \
molcrawl/models/gpt2/train.py molcrawl/tasks/pretrain/configs/molecule_nat_lang/gpt2_small.py

# BERT
CUDA_VISIBLE_DEVICES=0,2 \
torchrun --standalone --nproc_per_node=2 \
molcrawl/models/bert/main.py molcrawl/tasks/pretrain/configs/genome_sequence/bert_small.py

CUDA_VISIBLE_DEVICES=0,2 \
torchrun --standalone --nproc_per_node=2 \
molcrawl/models/bert/main.py molcrawl/tasks/pretrain/configs/protein_sequence/bert_small.py

CUDA_VISIBLE_DEVICES=0,2 \
torchrun --standalone --nproc_per_node=2 \
molcrawl/models/bert/main.py molcrawl/tasks/pretrain/configs/rna/bert_small.py

CUDA_VISIBLE_DEVICES=0,2 \
torchrun --standalone --nproc_per_node=2 \
molcrawl/models/bert/main.py molcrawl/tasks/pretrain/configs/compounds/bert_small.py

CUDA_VISIBLE_DEVICES=0,2 \
torchrun --standalone --nproc_per_node=2 \
molcrawl/models/bert/main.py molcrawl/tasks/pretrain/configs/molecule_nat_lang/bert_small.py

Effective batch size:

effective_batch_size = batch_size × gradient_accumulation_steps × num_gpus

Adjust batch_size or gradient_accumulation_steps accordingly to maintain reproducibility.

---

Hardware and Practical Notes

• Large and XL models require ≥32GB GPU memory
• Multiple GPUs are recommended for faster training
• Reduce batch size if memory errors occur

---

Note: Steps 4–6 can also be run collectively using the workflow scripts in the workflows/ directory. These scripts batch-process multiple modalities (data preparation, dataset building, and training) in a single command. See workflows/README.md for details.

---

7. Testing

GPT-2

Runs automatic evaluation on the test split, including language modeling metrics such as perplexity and token-level accuracy. It can also generate sample outputs and optionally convert the checkpoint to Hugging Face format.

python molcrawl/models/gpt2/test_checkpoint.py \
  --checkpoint_path <path_to_checkpoint> \
  --domain <dataset>

Optional arguments:

• --max_test_samples (default: use full test set) Limit the number of test samples for faster evaluation.

• --convert_to_hf (default: False) Convert the trained checkpoint to Hugging Face format after evaluation.

---

BERT

Runs evaluation on the test split using masked language modeling (MLM) metrics. It also supports prediction inspection and embedding extraction depending on configuration.

python molcrawl/models/bert/test_checkpoint.py \
  --checkpoint_path <path_to_checkpoint> \
  --domain <dataset>

Optional arguments:

• --max_test_samples (default: use full test set) Limit the number of test samples for faster evaluation.

• --extract_embeddings (default: False) Export hidden representations (embeddings) from the model.

• --show_predictions (default: False) Print sample masked-token predictions for qualitative inspection.

---

Dataset Statistics

Genome Sequence

• Number of sequences: 248,678
• Vocabulary size: 4096
• Number of tokens: 3,025,575,847

Protein Sequence (UniRef50)

• Size of the dataset: 18G
• Samples: 66,000,000
• Tokens: 19,182,955,286

RNA-seq (scRNAseq expression data)

• Size of the dataset: 221G
• Samples: 35,822,843
• Tokens: 90,711,564,293

Compounds

• Size of the dataset: 954M
• Samples: 13,288,710
• SMILES tokens: 526,014,485
• Scaffold tokens: 350,007,581

Molecule-related Natural Language

• Size of the dataset: 8.9G
• Training samples: 3,288,855
• Training tokens: 460,447,039
• Validation samples: 20,498
• Validation tokens: 2,568,952
• Test samples: 33,061
• Test tokens: 4,514,586
• Total samples: 3,342,414
• Total tokens: 467,530,577

---