--- license: apache-2.0 library_name: transformers tags: - modernbert - ModChemBERT - cheminformatics - chemical-language-model pipeline_tag: fill-mask --- # ModChemBERT: ModernBERT as a Chemical Language Model ModChemBERT-IR-BASE is a ModernBERT-based chemical language model (CLM) pretrained on SMILES strings using masked language modeling (MLM). This model serves as a base model for training embedding, retrieval, and reranking models for molecular information retrieval tasks. ## Usage Install the `transformers` library starting from v4.56.1: ```bash pip install -U "transformers>=4.56.1,<5.0.0" ``` ### Load Model ```python from transformers import AutoModelForMaskedLM, AutoTokenizer model_id = "Derify/ModChemBERT-IR-BASE" tokenizer = AutoTokenizer.from_pretrained(model_id) model = AutoModelForMaskedLM.from_pretrained( model_id, trust_remote_code=True, dtype="bfloat16", device_map="auto", ) ``` ### Fill-Mask Pipeline ```python from transformers import pipeline fill = pipeline("fill-mask", model=model, tokenizer=tokenizer) print(fill("c1ccccc1[MASK]")) ``` ## Architecture - Backbone: ModernBERT [1] - Hidden size: 1024 - Intermediate size: 1536 - Encoder Layers: 22 - Attention heads: 16 - Max sequence length: 512 tokens - Tokenizer: BPE tokenizer using [MolFormer's vocab](https://github.com/emapco/ModChemBERT/blob/main/modchembert/tokenizers/molformer/vocab.json) (2362 tokens) ## Dataset - Pretraining: [PubChem 110M dataset (canonical SMILES strings)](https://ibm.ent.box.com/v/MoLFormer-data) ## Pooling (Classifier / Regressor Head) Kallergis et al. [2] demonstrated that the CLM embedding method prior to the prediction head was the strongest contributor to downstream performance among evaluated hyperparameters. Behrendt et al. [3] noted that the last few layers contain task-specific information and that pooling methods leveraging information from multiple layers can enhance model performance. Their results further demonstrated that the `max_seq_mha` pooling method was particularly effective in low-data regimes. This base model includes configurable pooling strategies for downstream fine-tuning. When fine-tuned for embedding, retrieval, or reranking tasks (e.g., with Sentence Transformers), various pooling methods can be explored: - `cls`: Last layer [CLS] - `mean`: Mean over last hidden layer - `max_cls`: Max over last k layers of [CLS] - `cls_mha`: MHA with [CLS] as query - `max_seq_mha`: MHA with max pooled sequence as KV and max pooled [CLS] as query - `mean_seq_mha`: MHA with mean pooled sequence as KV and mean pooled [CLS] as query - `sum_mean`: Sum over all layers then mean tokens - `sum_sum`: Sum over all layers then sum tokens - `mean_mean`: Mean over all layers then mean tokens - `mean_sum`: Mean over all layers then sum tokens - `max_seq_mean`: Max over last k layers then mean tokens Note: ModChemBERT's `cls_mha`, `max_seq_mha`, and `mean_seq_mha` differ from MaxPoolBERT [3]. MaxPoolBERT uses PyTorch `nn.MultiheadAttention`, whereas ModChemBERT's `ModChemBertPoolingAttention` adapts ModernBERT's `ModernBertAttention`. On ChemBERTa-3 benchmarks this variant produced stronger validation metrics and avoided the training instabilities (sporadic zero / NaN losses and gradient norms) seen with `nn.MultiheadAttention`. Training instability with ModernBERT has been reported in the past ([discussion 1](https://huggingface.co/answerdotai/ModernBERT-base/discussions/59) and [discussion 2](https://huggingface.co/answerdotai/ModernBERT-base/discussions/63)). ## Intended Use * Primary: Base model for training embedding, retrieval, and reranking models for chemical information retrieval tasks using frameworks such as Sentence Transformers. * Appropriate for: Fine-tuning for semantic search of chemical compounds, molecular similarity tasks, chemical information retrieval systems, and as a foundation for building chemical embedding models. * Not intended for: Direct molecular property prediction without fine-tuning, generating novel molecules, or production use without domain-specific validation. ## Limitations - This is a base model pretrained only on masked language modeling; it requires fine-tuning for specific information retrieval tasks. - Performance on out-of-domain chemical spaces may vary: very long SMILES (>512 tokens), inorganic/organometallic compounds, polymers, or charged/enumerated tautomers may not be well represented in the training corpus. - The model reflects the chemical space distribution of PubChem and may not generalize equally well to all chemical domains. ## Ethical Considerations & Responsible Use - This base model is intended for research and development purposes in chemical information retrieval. - When fine-tuned for downstream applications, users should validate performance on their specific domain and use case. - Do not deploy in clinical, regulatory, or safety-critical settings without rigorous domain-specific validation and appropriate oversight. ## Hardware Training was performed on two NVIDIA RTX 3090 GPUs using `accelerate` for distributed (DDP) training. ## Citation If you use ModChemBERT-IR-BASE in your research, please cite the checkpoint and the following: ``` @software{cortes-2025-modchembert, author = {Emmanuel Cortes}, title = {ModChemBERT: ModernBERT as a Chemical Language Model}, year = {2025}, publisher = {GitHub}, howpublished = {GitHub repository}, url = {https://github.com/emapco/ModChemBERT} } ``` ## References 1. Warner, Benjamin, et al. "Smarter, better, faster, longer: A modern bidirectional encoder for fast, memory efficient, and long context finetuning and inference." arXiv preprint arXiv:2412.13663 (2024). 2. Kallergis, G., Asgari, E., Empting, M. et al. Domain adaptable language modeling of chemical compounds identifies potent pathoblockers for Pseudomonas aeruginosa. Commun Chem 8, 114 (2025). https://doi.org/10.1038/s42004-025-01484-4 3. Behrendt, Maike, Stefan Sylvius Wagner, and Stefan Harmeling. "MaxPoolBERT: Enhancing BERT Classification via Layer-and Token-Wise Aggregation." arXiv preprint arXiv:2505.15696 (2025).